Use callbacks in Room driver

Author: simolus3

integrations/room/src/commonMain/kotlin/com/powersync/integrations/room/RoomConnectionPool.kt

@@ -2,48 +2,37 @@ package com.powersync.integrations.room
 
 import androidx.room.RoomDatabase
 import androidx.room.Transactor
+import androidx.room.useReaderConnection
+import androidx.room.useWriterConnection
 import androidx.sqlite.SQLiteStatement
 import com.powersync.db.driver.SQLiteConnectionLease
 import com.powersync.db.driver.SQLiteConnectionPool
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharedFlow
-import kotlinx.coroutines.launch
 
 public class RoomConnectionPool(
     private val db: RoomDatabase,
-    private val scope: CoroutineScope,
 ): SQLiteConnectionPool {
     private val _updates = MutableSharedFlow<Set<String>>()
 
     override suspend fun <R> withAllConnections(action: suspend (SQLiteConnectionLease, List<SQLiteConnectionLease>) -> R) {
         // We can't obtain a list of all connections on Room. That's fine though, we expect this to
         // be used with raw tables, and withAllConnections is only used to apply a PowerSync schema.
-        action(write(), emptyList())
-    }
-
-    override suspend fun read(): SQLiteConnectionLease {
-        return obtainLease(true)
+        write {
+            action(it, emptyList())
+        }
     }
 
-    override suspend fun write(): SQLiteConnectionLease {
-        return obtainLease(false)
+    override suspend fun <T> read(callback: suspend (SQLiteConnectionLease) -> T): T {
+        return db.useReaderConnection {
+            callback(RoomTransactionLease(it))
+        }
     }
 
-    private suspend fun obtainLease(readonly: Boolean): SQLiteConnectionLease {
-        val obtainedLease = CompletableDeferred<SQLiteConnectionLease>()
-        scope.launch {
-            db.useConnection(readonly) { transactor ->
-                val completer = CompletableDeferred<Unit>()
-                obtainedLease.complete(RoomTransactionLease(transactor, completer))
-                completer.await()
-            }
+    override suspend fun <T> write(callback: suspend (SQLiteConnectionLease) -> T): T {
+        return db.useWriterConnection {
+            callback(RoomTransactionLease(it))
         }
-
-        val lease = obtainedLease.await()
-        return lease
     }
 
     override val updates: SharedFlow<Set<String>>
@@ -56,7 +45,6 @@ public class RoomConnectionPool(
 
 private class RoomTransactionLease(
     private val transactor: Transactor,
-    private val completer: CompletableDeferred<Unit>
 ): SQLiteConnectionLease {
     override suspend fun isInTransaction(): Boolean {
         return transactor.inTransaction()
@@ -68,8 +56,4 @@ private class RoomTransactionLease(
     ): R {
         return transactor.usePrepared(sql, block)
     }
-
-    override suspend fun close() {
-        completer.complete(Unit)
-    }
 }

integrations/room/src/commonTest/kotlin/com/powersync/integrations/room/PowerSyncRoomTest.kt

@@ -37,7 +37,7 @@ class PowerSyncRoomTest {
         val logger = Logger(loggerConfigInit())
 
         val powersync = PowerSyncDatabase.opened(
-            pool = RoomConnectionPool(database, this),
+            pool = RoomConnectionPool(database),
             scope = this,
             schema = TestDatabase.schema,
             group = PowerSyncDatabase.databaseGroup(logger, "test"),