Actually, use callbacks

Author: simolus3

core/src/commonIntegrationTest/kotlin/com/powersync/DatabaseTest.kt

@@ -22,6 +22,8 @@ import io.kotest.matchers.string.shouldContain
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.takeWhile
 import kotlinx.coroutines.launch
@@ -503,37 +505,45 @@ class DatabaseTest {
 
     @Test
     @OptIn(ExperimentalPowerSyncAPI::class)
-    fun testLeaseReadOnly() =
+    fun testUseRawReadOnly() =
         databaseTest {
             database.execute(
                 "INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)",
                 listOf("a", "[email protected]"),
             )
 
-
-            val raw = database.leaseConnection(readOnly = true)
-            raw.usePrepared("SELECT * FROM users") { stmt ->
-                stmt.step() shouldBe true
-                stmt.getText(1) shouldBe "a"
-                stmt.getText(2) shouldBe "[email protected]"
+            database.useConnection(true) {
+                it.usePrepared("SELECT * FROM users") { stmt ->
+                    stmt.step() shouldBe true
+                    stmt.getText(1) shouldBe "a"
+                    stmt.getText(2) shouldBe "a@example.org"
+                }
             }
-            raw.close()
         }
 
     @Test
     @OptIn(ExperimentalPowerSyncAPI::class)
-    fun testLeaseWrite() =
+    fun testUseRawWrite() =
         databaseTest {
-            val raw = database.leaseConnection(readOnly = false)
-            raw.usePrepared("INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)") { stmt ->
-                stmt.bindText(1, "name")
-                stmt.bindText(2, "email")
-                stmt.step() shouldBe false
-
-                stmt.reset()
-                stmt.step() shouldBe false
+            val didWrite = CompletableDeferred<Unit>()
+
+            val job = scope.launch {
+                database.useConnection(readOnly = false) { raw ->
+                    raw.usePrepared("INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)") { stmt ->
+                        stmt.bindText(1, "name")
+                        stmt.bindText(2, "email")
+                        stmt.step() shouldBe false
+
+                        stmt.reset()
+                        stmt.step() shouldBe false
+                    }
+
+                    didWrite.complete(Unit)
+                    awaitCancellation()
+                }
             }
 
+            didWrite.await()
             database.getAll("SELECT * FROM users") { it.getString("name") } shouldHaveSize 2
 
             // Verify that the statement indeed holds a lock on the database.
@@ -548,7 +558,7 @@ class DatabaseTest {
 
             delay(100.milliseconds)
             hadOtherWrite.isCompleted shouldBe false
-            raw.close()
+            job.cancelAndJoin()
             hadOtherWrite.await()
         }
 }

core/src/commonMain/kotlin/com/powersync/PowerSyncDatabase.kt

@@ -233,5 +233,9 @@ public interface PowerSyncDatabase : Queries {
                 group,
             )
         }
+
+        public fun databaseGroup(logger: Logger, identifier: String): Pair<ActiveDatabaseResource, Any> {
+            return ActiveDatabaseGroup.referenceDatabase(logger, identifier)
+        }
     }
 }

core/src/commonMain/kotlin/com/powersync/db/PowerSyncDatabaseImpl.kt

@@ -323,11 +323,15 @@ internal class PowerSyncDatabaseImpl(
     }
 
     @ExperimentalPowerSyncAPI
-override suspend fun leaseConnection(readOnly: Boolean): SQLiteConnectionLease {
+    override suspend fun <T> useConnection(
+        readOnly: Boolean,
+        block: suspend (SQLiteConnectionLease) -> T
+    ): T {
         waitReady()
-        return internalDb.leaseConnection(readOnly)
+        return internalDb.useConnection(readOnly, block)
     }
 
+
     override suspend fun <RowType : Any> get(
         sql: String,
         parameters: List<Any?>?,

core/src/commonMain/kotlin/com/powersync/db/Queries.kt

@@ -192,15 +192,11 @@ public interface Queries {
      *
      * This is useful when you need full control over the raw statements to use.
      *
-     * The connection needs to be released by calling [SQLiteConnectionLease.close] as soon as
-     * you're done with it, because the connection will occupy a read resource or the write lock
-     * while active.
-     *
      * Misusing this API, for instance by not cleaning up transactions started on the underlying
      * connection with a `BEGIN` statement or forgetting to close it, can disrupt the rest of the
      * PowerSync SDK. For this reason, this method should only be used if absolutely necessary.
      */
     @ExperimentalPowerSyncAPI()
     @HiddenFromObjC()
-    public suspend fun leaseConnection(readOnly: Boolean = false): SQLiteConnectionLease
+    public suspend fun <T> useConnection(readOnly: Boolean = false, block: suspend (SQLiteConnectionLease) -> T): T
 }

core/src/commonMain/kotlin/com/powersync/db/driver/InternalConnectionPool.kt

@@ -11,6 +11,7 @@ import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 
 @OptIn(ExperimentalPowerSyncAPI::class)
 internal class InternalConnectionPool(
@@ -65,40 +66,37 @@ internal class InternalConnectionPool(
         return connection
     }
 
-    override suspend fun read(): SQLiteConnectionLease {
-        return readPool.obtainConnection()
+    override suspend fun <T> read(callback: suspend (SQLiteConnectionLease) -> T): T {
+        return readPool.read(callback)
     }
 
-    override suspend fun write(): SQLiteConnectionLease {
-        writeLockMutex.lock()
-        return RawConnectionLease(writeConnection) {
-            // When we've leased a write connection, we may have to update table update flows
-            // after users ran their custom statements.
-            writeConnection.prepare("SELECT powersync_update_hooks('get')").use {
-                check(it.step())
-                val updatedTables = JsonUtil.json.decodeFromString<Set<String>>(it.getText(0))
-                if (updatedTables.isNotEmpty()) {
-                    scope.launch {
-                        tableUpdatesFlow.emit(updatedTables)
+    override suspend fun <T> write(callback: suspend (SQLiteConnectionLease) -> T): T {
+        return writeLockMutex.withLock {
+            try {
+                callback(RawConnectionLease(writeConnection))
+            } finally {
+                // When we've leased a write connection, we may have to update table update flows
+                // after users ran their custom statements.
+                writeConnection.prepare("SELECT powersync_update_hooks('get')").use {
+                    check(it.step())
+                    val updatedTables = JsonUtil.json.decodeFromString<Set<String>>(it.getText(0))
+                    if (updatedTables.isNotEmpty()) {
+                        scope.launch {
+                            tableUpdatesFlow.emit(updatedTables)
+                        }
                     }
                 }
             }
-
-            writeLockMutex.unlock()
         }
     }
 
     override suspend fun <R> withAllConnections(action: suspend (SQLiteConnectionLease, List<SQLiteConnectionLease>) -> R) {
         // First get a lock on all read connections
         readPool.withAllConnections { rawReadConnections ->
-            val readers = rawReadConnections.map { RawConnectionLease(it) {} }
+            val readers = rawReadConnections.map { RawConnectionLease(it) }
             // Then get access to the write connection
-            val writer = write()
-
-            try {
+            write { writer ->
                 action(writer, readers)
-            } finally {
-                writer.close()
             }
         }
     }

core/src/commonMain/kotlin/com/powersync/db/driver/RawConnectionLease.kt

@@ -11,7 +11,6 @@ import com.powersync.ExperimentalPowerSyncAPI
 @OptIn(ExperimentalPowerSyncAPI::class)
 internal class RawConnectionLease(
     private val connection: SQLiteConnection,
-    private val returnConnection: () -> Unit,
 ) : SQLiteConnectionLease {
     private var isCompleted = false
 
@@ -39,12 +38,4 @@ internal class RawConnectionLease(
         checkNotCompleted()
         return connection.prepare(sql).use(block)
     }
-
-    override suspend fun close() {
-        // Note: This is a lease, don't close the underlying connection.
-        if (!isCompleted) {
-            isCompleted = true
-            returnConnection()
-        }
-    }
 }

core/src/commonMain/kotlin/com/powersync/db/driver/ReadPool.kt

@@ -43,18 +43,21 @@ internal class ReadPool(
             }
         }
 
-    suspend fun obtainConnection(): RawConnectionLease {
-        val (connection, done) =
-            try {
-                available.receive()
-            } catch (e: PoolClosedException) {
-                throw PowerSyncException(
-                    message = "Cannot process connection pool request",
-                    cause = e,
-                )
-            }
+    suspend fun <T> read(block: suspend (SQLiteConnectionLease) -> T): T {
+        val (connection, done) = try {
+            available.receive()
+        } catch (e: PoolClosedException) {
+            throw PowerSyncException(
+                message = "Cannot process connection pool request",
+                cause = e,
+            )
+        }
 
-        return RawConnectionLease(connection) { done.complete(Unit) }
+        try {
+            return block(RawConnectionLease(connection))
+        } finally {
+            done.complete(Unit)
+        }
     }
 
     suspend fun <R> withAllConnections(action: suspend (connections: List<SQLiteConnection>) -> R): R {

core/src/commonMain/kotlin/com/powersync/db/driver/SQLiteConnectionPool.kt

@@ -7,15 +7,14 @@ import kotlinx.coroutines.runBlocking
 
 @ExperimentalPowerSyncAPI()
 public interface SQLiteConnectionPool {
-    public suspend fun read(): SQLiteConnectionLease
+    public suspend fun <T> read(callback: suspend (SQLiteConnectionLease) -> T): T
+    public suspend fun <T> write(callback: suspend (SQLiteConnectionLease) -> T): T
 
     public suspend fun <R> withAllConnections(action: suspend (
         writer: SQLiteConnectionLease,
         readers: List<SQLiteConnectionLease>
     ) -> R)
 
-    public suspend fun write(): SQLiteConnectionLease
-
     public val updates: SharedFlow<Set<String>>
 
     public suspend fun close()
@@ -50,9 +49,4 @@ public interface SQLiteConnectionLease {
             it.step()
         }
     }
-
-    /**
-     * Returns the leased connection to the pool.
-     */
-    public suspend fun close()
 }

core/src/commonMain/kotlin/com/powersync/db/internal/InternalDatabaseImpl.kt

@@ -152,12 +152,14 @@ internal class InternalDatabaseImpl(
                 }
         }
 
-
-    override suspend fun leaseConnection(readOnly: Boolean): SQLiteConnectionLease {
+    override suspend fun <T> useConnection(
+        readOnly: Boolean,
+        block: suspend (SQLiteConnectionLease) -> T
+    ): T {
         return if (readOnly) {
-            pool.read()
+            pool.read(block)
         } else {
-            pool.write()
+            pool.write(block)
         }
     }
 
@@ -168,14 +170,10 @@ internal class InternalDatabaseImpl(
     private suspend fun <R> internalReadLock(callback: suspend (SQLiteConnectionLease) -> R): R =
         withContext(dbContext) {
             runWrapped {
-                val connection = leaseConnection(readOnly = true)
-                try {
+                useConnection(true) { connection ->
                     catchSwiftExceptions {
                         callback(connection)
                     }
-                } finally {
-                    // Closing the lease will release the connection back into the pool.
-                    connection.close()
                 }
             }
         }
@@ -197,16 +195,12 @@ internal class InternalDatabaseImpl(
     @OptIn(ExperimentalPowerSyncAPI::class)
     private suspend fun <R> internalWriteLock(callback: suspend (SQLiteConnectionLease) -> R): R =
         withContext(dbContext) {
-            val lease = pool.write()
-            try {
+            pool.write { writer ->
                 runWrapped {
                     catchSwiftExceptions {
-                        callback(lease)
+                        callback(writer)
                     }
                 }
-            } finally {
-                // Returning the lease will unlock the writeLockMutex
-                lease.close()
             }
         }