Lease API that works better with Room

Author: simolus3

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

@@ -8,6 +8,7 @@ import co.touchlab.kermit.ExperimentalKermitApi
 import com.powersync.db.ActiveDatabaseGroup
 import com.powersync.db.crud.CrudEntry
 import com.powersync.db.crud.CrudTransaction
+import com.powersync.db.getString
 import com.powersync.db.schema.Schema
 import com.powersync.testutils.UserRow
 import com.powersync.testutils.databaseTest
@@ -18,17 +19,18 @@ import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.collections.shouldHaveSize
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.string.shouldContain
-import io.kotest.matchers.throwable.shouldHaveMessage
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.takeWhile
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withContext
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
+import kotlin.time.Duration.Companion.milliseconds
 
 @OptIn(ExperimentalKermitApi::class)
 class DatabaseTest {
@@ -500,36 +502,52 @@ class DatabaseTest {
         }
 
     @Test
-    fun testRawConnection() =
+    @OptIn(ExperimentalPowerSyncAPI::class)
+    fun testLeaseReadOnly() =
         databaseTest {
             database.execute(
                 "INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)",
                 listOf("a", "[email protected]"),
             )
-            var capturedConnection: SQLiteConnection? = null
 
-            database.readLock {
-                it.rawConnection.prepare("SELECT * FROM users").use { stmt ->
-                    stmt.step() shouldBe true
-                    stmt.getText(1) shouldBe "a"
-                    stmt.getText(2) shouldBe "[email protected]"
-                }
+            val raw = database.leaseConnection(readOnly = true)
+            raw.prepare("SELECT * FROM users").use { stmt ->
+                stmt.step() shouldBe true
+                stmt.getText(1) shouldBe "a"
+                stmt.getText(2) shouldBe "[email protected]"
+            }
+            raw.close()
+        }
 
-                capturedConnection = it.rawConnection
+    @Test
+    @OptIn(ExperimentalPowerSyncAPI::class)
+    fun testLeaseWrite() =
+        databaseTest {
+            val raw = database.leaseConnection(readOnly = false)
+            raw.prepare("INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)").use { stmt ->
+                stmt.bindText(1, "name")
+                stmt.bindText(2, "email")
+                stmt.step() shouldBe false
+
+                stmt.reset()
+                stmt.step() shouldBe false
             }
 
-            // When we exit readLock, the connection should no longer be usable
-            shouldThrow<IllegalStateException> { capturedConnection!!.execSQL("DELETE FROM users") } shouldHaveMessage
-                "Connection lease already closed"
+            database.getAll("SELECT * FROM users") { it.getString("name") } shouldHaveSize 2
 
-            capturedConnection = null
-            database.writeLock {
-                it.rawConnection.execSQL("DELETE FROM users")
-                capturedConnection = it.rawConnection
+            // Verify that the statement indeed holds a lock on the database.
+            val hadOtherWrite = CompletableDeferred<Unit>()
+            scope.launch {
+                database.execute(
+                    "INSERT INTO users (id, name, email) VALUES (uuid(), ?, ?)",
+                    listOf("another", "[email protected]"),
+                )
+                hadOtherWrite.complete(Unit)
             }
 
-            // Same thing for writes
-            shouldThrow<IllegalStateException> { capturedConnection!!.prepare("SELECT * FROM users") } shouldHaveMessage
-                "Connection lease already closed"
+            delay(100.milliseconds)
+            hadOtherWrite.isCompleted shouldBe false
+            raw.close()
+            hadOtherWrite.await()
         }
 }

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

@@ -1,7 +1,9 @@
 package com.powersync.db
 
+import androidx.sqlite.SQLiteConnection
 import co.touchlab.kermit.Logger
 import com.powersync.DatabaseDriverFactory
+import com.powersync.ExperimentalPowerSyncAPI
 import com.powersync.PowerSyncDatabase
 import com.powersync.PowerSyncException
 import com.powersync.bucket.BucketPriority
@@ -329,6 +331,12 @@ internal class PowerSyncDatabaseImpl(
         return powerSyncVersion
     }
 
+    @ExperimentalPowerSyncAPI
+    override suspend fun leaseConnection(readOnly: Boolean): SQLiteConnection {
+        waitReady()
+        return internalDb.leaseConnection(readOnly)
+    }
+
     override suspend fun <RowType : Any> get(
         sql: String,
         parameters: List<Any?>?,

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

@@ -1,10 +1,13 @@
 package com.powersync.db
 
+import androidx.sqlite.SQLiteConnection
+import com.powersync.ExperimentalPowerSyncAPI
 import com.powersync.PowerSyncException
 import com.powersync.db.internal.ConnectionContext
 import com.powersync.db.internal.PowerSyncTransaction
 import kotlinx.coroutines.flow.Flow
 import kotlin.coroutines.cancellation.CancellationException
+import kotlin.native.HiddenFromObjC
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.milliseconds
 
@@ -183,4 +186,21 @@ public interface Queries {
      */
     @Throws(PowerSyncException::class, CancellationException::class)
     public suspend fun <R> readTransaction(callback: ThrowableTransactionCallback<R>): R
+
+    /**
+     * Obtains a connection from the read pool or an exclusive reference on the write connection.
+     *
+     * This is useful when you need full control over the raw statements to use.
+     *
+     * The connection needs to be released by calling [SQLiteConnection.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): SQLiteConnection
 }

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

@@ -5,12 +5,8 @@ import androidx.sqlite.SQLiteStatement
 import com.powersync.PowerSyncException
 import com.powersync.db.SqlCursor
 import com.powersync.db.StatementBasedCursor
-import kotlin.native.HiddenFromObjC
 
 public interface ConnectionContext {
-    @HiddenFromObjC
-    public val rawConnection: SQLiteConnection
-
     @Throws(PowerSyncException::class)
     public fun execute(
         sql: String,
@@ -40,7 +36,7 @@ public interface ConnectionContext {
 }
 
 internal class ConnectionContextImplementation(
-    override val rawConnection: SQLiteConnection,
+    private val rawConnection: SQLiteConnection,
 ) : ConnectionContext {
     override fun execute(
         sql: String,

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

@@ -38,7 +38,7 @@ internal class ConnectionPool(
             }
         }
 
-    suspend fun <R> withConnection(action: suspend (connection: SQLiteConnection) -> R): R {
+    suspend fun obtainConnection(): RawConnectionLease {
         val (connection, done) =
             try {
                 available.receive()
@@ -49,11 +49,7 @@ internal class ConnectionPool(
                 )
             }
 
-        try {
-            return action(connection)
-        } finally {
-            done.complete(Unit)
-        }
+        return RawConnectionLease(connection) { done.complete(Unit) }
     }
 
     suspend fun <R> withAllConnections(action: suspend (connections: List<SQLiteConnection>) -> R): R {

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

@@ -4,6 +4,7 @@ import androidx.sqlite.SQLiteConnection
 import androidx.sqlite.execSQL
 import co.touchlab.kermit.Logger
 import com.powersync.DatabaseDriverFactory
+import com.powersync.ExperimentalPowerSyncAPI
 import com.powersync.PowerSyncException
 import com.powersync.db.SqlCursor
 import com.powersync.db.ThrowableLockCallback
@@ -22,7 +23,6 @@ import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.onSubscription
 import kotlinx.coroutines.flow.transform
 import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
 import kotlin.time.Duration.Companion.milliseconds
 
@@ -206,17 +206,30 @@ internal class InternalDatabaseImpl(
                 }
         }
 
+    @ExperimentalPowerSyncAPI
+    override suspend fun leaseConnection(readOnly: Boolean): SQLiteConnection =
+        if (readOnly) {
+            readPool.obtainConnection()
+        } else {
+            writeLockMutex.lock()
+            RawConnectionLease(writeConnection, writeLockMutex::unlock)
+        }
+
     /**
      * Creates a read lock while providing an internal transactor for transactions
      */
+    @OptIn(ExperimentalPowerSyncAPI::class)
     private suspend fun <R> internalReadLock(callback: (SQLiteConnection) -> R): R =
         withContext(dbContext) {
             runWrapped {
-                readPool.withConnection {
+                val connection = leaseConnection(readOnly = true)
+                try {
                     catchSwiftExceptions {
-                        val lease = RawConnectionLease(it)
-                        callback(lease).also { lease.completed = true }
+                        callback(connection)
                     }
+                } finally {
+                    // Closing the lease will release the connection back into the pool.
+                    connection.close()
                 }
             }
         }
@@ -235,11 +248,11 @@ internal class InternalDatabaseImpl(
             }
         }
 
+    @OptIn(ExperimentalPowerSyncAPI::class)
     private suspend fun <R> internalWriteLock(callback: (SQLiteConnection) -> R): R =
         withContext(dbContext) {
-            writeLockMutex.withLock {
-                val lease = RawConnectionLease(writeConnection)
-
+            val lease = leaseConnection(readOnly = false)
+            try {
                 runWrapped {
                     catchSwiftExceptions {
                         callback(lease)
@@ -248,8 +261,10 @@ internal class InternalDatabaseImpl(
                     // Trigger watched queries
                     // Fire updates inside the write lock
                     updates.fireTableUpdates()
-                    lease.completed = true
                 }
+            } finally {
+                // Returning the lease will unlock the writeLockMutex
+                lease.close()
             }
         }
 

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

@@ -8,7 +8,7 @@ import com.powersync.db.SqlCursor
 public interface PowerSyncTransaction : ConnectionContext
 
 internal class PowerSyncTransactionImpl(
-    override val rawConnection: SQLiteConnection,
+    private val rawConnection: SQLiteConnection,
 ) : PowerSyncTransaction,
     ConnectionContext {
     private val delegate = ConnectionContextImplementation(rawConnection)

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

@@ -8,10 +8,12 @@ import androidx.sqlite.SQLiteStatement
  */
 internal class RawConnectionLease(
     private val connection: SQLiteConnection,
-    var completed: Boolean = false,
+    private val returnConnection: () -> Unit,
 ) : SQLiteConnection {
+    private var isCompleted = false
+
     private fun checkNotCompleted() {
-        check(!completed) { "Connection lease already closed" }
+        check(!isCompleted) { "Connection lease already closed" }
     }
 
     override fun inTransaction(): Boolean {
@@ -26,5 +28,9 @@ internal class RawConnectionLease(
 
     override fun close() {
         // Note: This is a lease, don't close the underlying connection.
+        if (!isCompleted) {
+            isCompleted = true
+            returnConnection()
+        }
     }
 }

drivers/common/build.gradle.kts

@@ -45,7 +45,7 @@ kotlin {
 }
 
 android {
-    namespace = "com.powersync.compose"
+    namespace = "com.powersync.drivers.common"
     compileSdk =
         libs.versions.android.compileSdk
             .get()

drivers/common/src/androidMain/kotlin/com/powersync/internal/driver/AndroidDriver.kt

@@ -4,7 +4,9 @@ import android.content.Context
 import java.util.Properties
 import java.util.concurrent.atomic.AtomicBoolean
 
-public class AndroidDriver(private val context: Context): JdbcDriver() {
+public class AndroidDriver(
+    private val context: Context,
+) : JdbcDriver() {
     override fun addDefaultProperties(properties: Properties) {
         val isFirst = IS_FIRST_CONNECTION.getAndSet(false)
         if (isFirst) {