fix: drizzle locks

Author: stevensJourney

.changeset/nice-mangos-fail.md

@@ -0,0 +1,5 @@
+---
+'@powersync/drizzle-driver': patch
+---
+
+Fixed issue where read queries would not use a readLock.

packages/drizzle-driver/src/sqlite/PowerSyncSQLiteBaseSession.ts

@@ -1,4 +1,4 @@
-import { LockContext, QueryResult } from '@powersync/common';
+import { QueryResult } from '@powersync/common';
 import { entityKind } from 'drizzle-orm/entity';
 import type { Logger } from 'drizzle-orm/logger';
 import { NoopLogger } from 'drizzle-orm/logger';
@@ -7,13 +7,13 @@ import { type Query } from 'drizzle-orm/sql/sql';
 import type { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core/dialect';
 import type { SelectedFieldsOrdered } from 'drizzle-orm/sqlite-core/query-builders/select.types';
 import {
-  type PreparedQueryConfig as PreparedQueryConfigBase,
-  type SQLiteExecuteMethod,
   SQLiteSession,
   SQLiteTransaction,
+  type PreparedQueryConfig as PreparedQueryConfigBase,
+  type SQLiteExecuteMethod,
   type SQLiteTransactionConfig
 } from 'drizzle-orm/sqlite-core/session';
-import { PowerSyncSQLitePreparedQuery } from './PowerSyncSQLitePreparedQuery.js';
+import { ContextProvider, PowerSyncSQLitePreparedQuery } from './PowerSyncSQLitePreparedQuery.js';
 
 export interface PowerSyncSQLiteSessionOptions {
   logger?: Logger;
@@ -39,7 +39,7 @@ export class PowerSyncSQLiteBaseSession<
   protected logger: Logger;
 
   constructor(
-    protected db: LockContext,
+    protected contextProvider: ContextProvider,
     protected dialect: SQLiteAsyncDialect,
     protected schema: RelationalSchemaConfig<TSchema> | undefined,
     protected options: PowerSyncSQLiteSessionOptions = {}
@@ -56,7 +56,7 @@ export class PowerSyncSQLiteBaseSession<
     customResultMapper?: (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => unknown
   ): PowerSyncSQLitePreparedQuery<T> {
     return new PowerSyncSQLitePreparedQuery(
-      this.db,
+      this.contextProvider,
       query,
       this.logger,
       fields,

packages/drizzle-driver/src/sqlite/PowerSyncSQLitePreparedQuery.ts

@@ -1,18 +1,35 @@
 import { LockContext, QueryResult } from '@powersync/common';
-import { Column, DriverValueDecoder, getTableName, SQL } from 'drizzle-orm';
+import { Column, DriverValueDecoder, SQL, getTableName } from 'drizzle-orm';
 import { entityKind, is } from 'drizzle-orm/entity';
 import type { Logger } from 'drizzle-orm/logger';
 import { fillPlaceholders, type Query } from 'drizzle-orm/sql/sql';
 import { SQLiteColumn } from 'drizzle-orm/sqlite-core';
 import type { SelectedFieldsOrdered } from 'drizzle-orm/sqlite-core/query-builders/select.types';
 import {
+  SQLitePreparedQuery,
   type PreparedQueryConfig as PreparedQueryConfigBase,
-  type SQLiteExecuteMethod,
-  SQLitePreparedQuery
+  type SQLiteExecuteMethod
 } from 'drizzle-orm/sqlite-core/session';
 
 type PreparedQueryConfig = Omit<PreparedQueryConfigBase, 'statement' | 'run'>;
 
+/**
+ * Callback which uses a LockContext for database operations.
+ */
+export type LockCallback<T> = (ctx: LockContext) => Promise<T>;
+
+/**
+ * Provider for specific database contexts.
+ * Handlers are provided a context to the provided callback.
+ * This does not necessarily need to acquire a database lock for each call.
+ * Calls might use the same lock context for multiple operations.
+ * The read/write context may relate to a single read OR write context.
+ */
+export type ContextProvider = {
+  useReadContext: <T>(fn: LockCallback<T>) => Promise<T>;
+  useWriteContext: <T>(fn: LockCallback<T>) => Promise<T>;
+};
+
 export class PowerSyncSQLitePreparedQuery<
   T extends PreparedQueryConfig = PreparedQueryConfig
 > extends SQLitePreparedQuery<{
@@ -26,7 +43,7 @@ export class PowerSyncSQLitePreparedQuery<
   static readonly [entityKind]: string = 'PowerSyncSQLitePreparedQuery';
 
   constructor(
-    private db: LockContext,
+    private contextProvider: ContextProvider,
     query: Query,
     private logger: Logger,
     private fields: SelectedFieldsOrdered | undefined,
@@ -40,17 +57,23 @@ export class PowerSyncSQLitePreparedQuery<
   async run(placeholderValues?: Record<string, unknown>): Promise<QueryResult> {
     const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
     this.logger.logQuery(this.query.sql, params);
-    const rs = await this.db.execute(this.query.sql, params);
-    return rs;
+    /**
+     * Run operations are teated as potential mutations, so they use the write context.
+     */
+    return this.contextProvider.useWriteContext(async (ctx) => {
+      const rs = await ctx.execute(this.query.sql, params);
+      return rs;
+    });
   }
 
   async all(placeholderValues?: Record<string, unknown>): Promise<T['all']> {
     const { fields, query, logger, customResultMapper } = this;
     if (!fields && !customResultMapper) {
       const params = fillPlaceholders(query.params, placeholderValues ?? {});
       logger.logQuery(query.sql, params);
-      const rs = await this.db.execute(this.query.sql, params);
-      return rs.rows?._array ?? [];
+      return await this.contextProvider.useReadContext(async (ctx) => {
+        return await ctx.getAll(this.query.sql, params);
+      });
     }
 
     const rows = (await this.values(placeholderValues)) as unknown[][];
@@ -69,7 +92,9 @@ export class PowerSyncSQLitePreparedQuery<
     const { fields, customResultMapper } = this;
     const joinsNotNullableMap = (this as any).joinsNotNullableMap;
     if (!fields && !customResultMapper) {
-      return this.db.get(this.query.sql, params);
+      return this.contextProvider.useReadContext(async (ctx) => {
+        return await ctx.get(this.query.sql, params);
+      });
     }
 
     const rows = (await this.values(placeholderValues)) as unknown[][];
@@ -90,7 +115,9 @@ export class PowerSyncSQLitePreparedQuery<
     const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
     this.logger.logQuery(this.query.sql, params);
 
-    return await this.db.executeRaw(this.query.sql, params);
+    return await this.contextProvider.useReadContext(async (ctx) => {
+      return await ctx.executeRaw(this.query.sql, params);
+    });
   }
 
   isResponseInArrayMode(): boolean {

packages/drizzle-driver/src/sqlite/PowerSyncSQLiteSession.ts

@@ -1,4 +1,4 @@
-import { AbstractPowerSyncDatabase, DBAdapter } from '@powersync/common';
+import { AbstractPowerSyncDatabase, LockContext } from '@powersync/common';
 import { entityKind } from 'drizzle-orm/entity';
 import type { RelationalSchemaConfig, TablesRelationalConfig } from 'drizzle-orm/relations';
 import type { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core/dialect';
@@ -21,7 +21,16 @@ export class PowerSyncSQLiteSession<
     schema: RelationalSchemaConfig<TSchema> | undefined,
     options: PowerSyncSQLiteSessionOptions = {}
   ) {
-    super(db, dialect, schema, options);
+    super(
+      // Top level operations use the respective locks.
+      {
+        useReadContext: (callback) => db.readLock(callback),
+        useWriteContext: (callback) => db.writeLock(callback)
+      },
+      dialect,
+      schema,
+      options
+    );
     this.client = db;
   }
 
@@ -39,14 +48,23 @@ export class PowerSyncSQLiteSession<
   }
 
   protected async internalTransaction<T>(
-    connection: DBAdapter,
+    connection: LockContext,
     fn: (tx: PowerSyncSQLiteTransaction<TFullSchema, TSchema>) => T,
     config: PowerSyncSQLiteTransactionConfig = {}
   ): Promise<T> {
     const tx = new PowerSyncSQLiteTransaction<TFullSchema, TSchema>(
       'async',
       (this as any).dialect,
-      new PowerSyncSQLiteBaseSession(connection, this.dialect, this.schema, this.options),
+      new PowerSyncSQLiteBaseSession(
+        {
+          // We already have a fixed context here. We need to use it for both "read" and "write" operations.
+          useReadContext: (callback) => callback(connection),
+          useWriteContext: (callback) => callback(connection)
+        },
+        this.dialect,
+        this.schema,
+        this.options
+      ),
       this.schema
     );