Await initialization before connecting

Author: simolus3

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Changelog
 
-## 1.14.0-Beta.0
+## 1.14.0
 
 * Remove internal dependency on the PowerSync Kotlin SDK. Going forward, the Swift SDK is implemented in Swift!
   __Important__: While these changes are tested, they are a full rewrite of the internal connection pool logic.

Sources/PowerSync/CurrentVersion.swift

@@ -1,3 +1,3 @@
 // The current version of the PowerSync Swift SDK. This should be updated to the latest version in `CHANGELOG.md` when a new version is released.
-let libraryVersion = "1.14.0-Beta.0"
+let libraryVersion = "1.14.0"
 

Sources/PowerSync/Implementation/PowerSyncDatabaseImpl.swift

@@ -8,7 +8,7 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
     private let dbFilename: String?
     private let httpClient: HttpClient
     private let initializer = DatabaseInitializationAction()
-    fileprivate let queries: ConnectionPoolQueries
+    let pool: any SQLiteConnectionPoolProtocol
     let schema: AsyncMutex<Schema>
 
     init(
@@ -24,7 +24,7 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         self.logger = logger
         self.schema = AsyncMutex(schema)
         self.httpClient = httpClient
-        self.queries = ConnectionPoolQueries(pool: pool)
+        self.pool = pool
         self.group = activeInstanceStore.referenceGroup(identifier: identifier, logger: logger)
     }
 
@@ -44,7 +44,7 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
 
     fileprivate func resolveOfflineSyncStatus() async throws {
         // We can't use get() here because it runs as part of the initialization step.
-        let offlineSyncStatus = try await queries.readLock { connection in
+        let offlineSyncStatus = try await readLockInner { connection in
             try connection.get(sql: "SELECT powersync_offline_sync_status()", parameters: []) { cursor in
                 let raw = try cursor.getString(index: 0)
                 guard let data = raw.data(using: .utf8) else {
@@ -70,16 +70,17 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
     }
 
     fileprivate func applySchema(schema: Schema) async throws {
-        try await queries.withAll { writer, readers in
+        try await pool.withAllConnections { writer, readers in
             let encoded = try StreamingSyncClient.jsonEncoder.encode(schema)
             guard let asString = String(data: encoded, encoding: .utf8) else {
                 throw PowerSyncError.operationFailed(message: "Could not serialize schema")
             }
-            try writer.execute(sql: "SELECT powersync_replace_schema(?)", parameters: [asString])
+
+            let _ = try writer.execute(sql: "SELECT powersync_replace_schema(?)", parameters: [.string(asString)])
 
             for reader in readers {
                 // Update the schema on all read connections
-                try reader.execute(sql: "pragma table_info('sqlite_master')", parameters: [])
+                let _ = try reader.execute(sql: "pragma table_info('sqlite_master')", parameters: [])
             }
         }
     }
@@ -113,7 +114,7 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         try await initialize()
         try await initializer.close {
             await group.syncCoordinator.disconnect()
-            try await queries.pool.close()
+            try await pool.close()
         }
     }
 
@@ -127,6 +128,7 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
     }
 
     func connect(connector: any PowerSyncBackendConnectorProtocol, options: ConnectOptions?) async throws {
+        try await initialize()
         await group.syncCoordinator.connect(db: self, connector: connector, options: options ?? ConnectOptions(), client: httpClient)
     }
 
@@ -143,23 +145,35 @@ final class PowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
 
             do {
                 let flags = flags
-                let _ = try await queries.writeLock { ctx in try ctx.execute(sql: "SELECT powersync_clear(?)", parameters: [flags]) }
+                let _ = try await writeLockInner { ctx in try ctx.execute(sql: "SELECT powersync_clear(?)", parameters: [flags]) }
             }
         }
     }
 
     func writeLock<R: Sendable>(callback: @escaping @Sendable (any ConnectionContext) throws -> R) async throws -> R {
         try await initialize()
-        return try await queries.writeLock(callback: callback)
+        return try await writeLockInner(callback: callback)
+    }
+
+    fileprivate func writeLockInner<R: Sendable>(callback: @escaping @Sendable (any ConnectionContext) throws -> R) async throws -> R {
+        return try await self.pool.write { connection in
+            try callback(ConnectionLeaseContext(lease: connection))
+        }
     }
 
     func readLock<R: Sendable>(callback: @escaping @Sendable (any ConnectionContext) throws -> R) async throws -> R {
         try await initialize()
-        return try await queries.readLock(callback: callback)
+        return try await readLockInner(callback: callback)
+    }
+
+    fileprivate func readLockInner<R: Sendable>(callback: @escaping @Sendable (any ConnectionContext) throws -> R) async throws -> R {
+        return try await pool.read { connection in
+            try callback(ConnectionLeaseContext(lease: connection))
+        }
     }
 
     func watch<RowType: Sendable>(options: WatchOptions<RowType>) throws -> AsyncThrowingStream<[RowType], any Error> {
-        return try queries.watch(options: options)
+        return watchImpl(db: self, options: options)
     }
 
     static let maxOpId = Int64.max
@@ -178,7 +192,7 @@ private actor DatabaseInitializationAction {
             return
         }
 
-        powerSyncVersion = try await db.queries.writeLock { conn in
+        powerSyncVersion = try await db.writeLockInner { conn in
             let sqliteVersion = try conn.get(sql: "SELECT sqlite_version()", parameters: []) { try $0.getString(index: 0) }
             let powerSyncVersion = try conn.get(sql: "SELECT powersync_rs_version()", parameters: []) { try $0.getString(index: 0) }
 

Sources/PowerSync/Implementation/queries/ConnectionPoolQueries.swift

@@ -0,0 +1,108 @@
+import AsyncAlgorithms
+
+func watchImpl<RowType: Sendable>(db: PowerSyncDatabaseImpl, options: WatchOptions<RowType>) -> AsyncThrowingStream<[RowType], any Error> {
+    AsyncThrowingStream { continuation in
+        // Create an outer task to monitor cancellation
+        let task = Task {
+            do {
+                let watchedTables = try await getQuerySourceTables(
+                    db: db,
+                    sql: options.sql,
+                    parameters: options.parameters
+                )
+
+                let updateNotifications = db.pool.tableUpdates.filter { changedTables in
+                    changedTables.contains(where: watchedTables.contains)
+                }.map { _ in () }
+                // Allows emitting the first result even if there aren't changes
+                let withInitial = AsyncAlgorithms.merge([()].async, updateNotifications)
+                let merged = MergeItemSequence(inner: withInitial)
+
+                for try await _ in merged {
+                    // Check if the outer task is cancelled
+                    try Task.checkCancellation()
+
+                    try continuation.yield(await db.getAll(
+                        sql: options.sql,
+                        parameters: options.parameters,
+                        mapper: options.mapper
+                    ))
+                    try await sleepForSeconds(seconds: options.throttle)
+                }
+
+                continuation.finish()
+            } catch {
+                if error is CancellationError {
+                    continuation.finish()
+                } else {
+                    continuation.finish(throwing: error)
+                }
+            }
+        }
+
+        // Propagate cancellation from the outer task to the inner task
+        continuation.onTermination = { @Sendable _ in
+            task.cancel() // This cancels the inner task when the stream is terminated
+        }
+    }
+}
+
+private func getQuerySourceTables(
+    db: PowerSyncDatabaseImpl,
+    sql: String,
+    parameters: [Sendable?]
+) async throws -> Set<String> {
+    let rows = try await db.getAll(
+        sql: "EXPLAIN \(sql)",
+        parameters: parameters,
+        mapper: { cursor in
+            try ExplainQueryResult(
+                addr: cursor.getString(index: 0),
+                opcode: cursor.getString(index: 1),
+                p1: cursor.getInt64(index: 2),
+                p2: cursor.getInt64(index: 3),
+                p3: cursor.getInt64(index: 4)
+            )
+        }
+    )
+
+    let rootPages = rows.compactMap { row in
+        if (row.opcode == "OpenRead" || row.opcode == "OpenWrite") &&
+            row.p3 == 0 && row.p2 != 0
+        {
+            return row.p2
+        }
+        return nil
+    }
+
+    do {
+        let pagesData = try StreamingSyncClient.jsonEncoder.encode(rootPages)
+        guard let pagesString = String(data: pagesData, encoding: .utf8) else {
+            throw PowerSyncError.operationFailed(
+                message: "Failed to convert pages data to UTF-8 string"
+            )
+        }
+
+        let tableRows = try await db.getAll(
+            sql: "SELECT tbl_name FROM sqlite_master WHERE rootpage IN (SELECT json_each.value FROM json_each(?))",
+            parameters: [
+                pagesString,
+            ]
+        ) { try $0.getString(index: 0) }
+
+        return Set(tableRows)
+    } catch {
+        throw PowerSyncError.operationFailed(
+            message: "Could not determine watched query tables",
+            underlyingError: error
+        )
+    }
+}
+
+private struct ExplainQueryResult {
+    let addr: String
+    let opcode: String
+    let p1: Int64
+    let p2: Int64
+    let p3: Int64
+}