[Internal] Improve Lock Error handling

Sources/PowerSync/Kotlin/KotlinPowerSyncDatabaseImpl.swift

@@ -275,9 +275,7 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         return try await wrapPowerSyncException {
             try safeCast(
                 await kotlinDatabase.writeLock(
-                    callback: LockCallback(
-                        callback: callback
-                    )
+                    callback: wrapLockContext(callback: callback)
                 ),
                 to: R.self
             )
@@ -290,9 +288,7 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         return try await wrapPowerSyncException {
             try safeCast(
                 await kotlinDatabase.writeTransaction(
-                    callback: TransactionCallback(
-                        callback: callback
-                    )
+                    callback: wrapTransactionContext(callback: callback)
                 ),
                 to: R.self
             )
@@ -307,9 +303,7 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         return try await wrapPowerSyncException {
             try safeCast(
                 await kotlinDatabase.readLock(
-                    callback: LockCallback(
-                        callback: callback
-                    )
+                    callback: wrapLockContext(callback: callback)
                 ),
                 to: R.self
             )
@@ -322,9 +316,7 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
         return try await wrapPowerSyncException {
             try safeCast(
                 await kotlinDatabase.readTransaction(
-                    callback: TransactionCallback(
-                        callback: callback
-                    )
+                    callback: wrapTransactionContext(callback: callback)
                 ),
                 to: R.self
             )
@@ -372,11 +364,11 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
             }
         )
 
-        let rootPages = rows.compactMap { r in
-            if (r.opcode == "OpenRead" || r.opcode == "OpenWrite") &&
-                r.p3 == 0 && r.p2 != 0
+        let rootPages = rows.compactMap { row in
+            if (row.opcode == "OpenRead" || row.opcode == "OpenWrite") &&
+                row.p3 == 0 && row.p2 != 0
             {
-                return r.p2
+                return row.p2
             }
             return nil
         }
@@ -389,7 +381,7 @@ final class KotlinPowerSyncDatabaseImpl: PowerSyncDatabaseProtocol {
                     message: "Failed to convert pages data to UTF-8 string"
                 )
             }
-            
+
             let tableRows = try await getAll(
                 sql: "SELECT tbl_name FROM sqlite_master WHERE rootpage IN (SELECT json_each.value FROM json_each(?))",
                 parameters: [
@@ -414,3 +406,50 @@ private struct ExplainQueryResult {
     let p2: Int64
     let p3: Int64
 }
+
+extension Error {
+    func toPowerSyncError() -> PowerSyncKotlin.PowerSyncException {
+        return PowerSyncKotlin.PowerSyncException(
+            message: localizedDescription,
+            cause: PowerSyncKotlin.KotlinThrowable(message: localizedDescription)
+        )
+    }
+}
+
+func wrapLockContext(
+    callback: @escaping (any ConnectionContext) throws -> Any
+) throws -> PowerSyncKotlin.ThrowableLockCallback {
+    PowerSyncKotlin.wrapContextHandler { kotlinContext in
+        do {
+            return try PowerSyncKotlin.PowerSyncResult.Success(
+                value: callback(
+                    KotlinConnectionContext(
+                        ctx: kotlinContext
+                    )
+                ))
+        } catch {
+            return PowerSyncKotlin.PowerSyncResult.Failure(
+                exception: error.toPowerSyncError()
+            )
+        }
+    }
+}
+
+func wrapTransactionContext(
+    callback: @escaping (any Transaction) throws -> Any
+) throws -> PowerSyncKotlin.ThrowableTransactionCallback {
+    PowerSyncKotlin.wrapTransactionContextHandler { kotlinContext in
+        do {
+            return try PowerSyncKotlin.PowerSyncResult.Success(
+                value: callback(
+                    KotlinTransactionContext(
+                        ctx: kotlinContext
+                    )
+                ))
+        } catch {
+            return PowerSyncKotlin.PowerSyncResult.Failure(
+                exception: error.toPowerSyncError()
+            )
+        }
+    }
+}