test: findMany use read lock

Author: stevensJourney

packages/drizzle-driver/README.md

@@ -155,4 +155,3 @@ For more information on how to use Drizzle queries in PowerSync, see [here](http
 ## Known limitations
 
 - The integration does not currently support nested transactions (also known as `savepoints`).
-- `findMany` queries execute with a write connection if not wrapped in an explicit `read only` transaction.

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

@@ -11,15 +11,17 @@ import {
   createTableRelationsHelpers,
   extractTablesRelationalConfig,
   ExtractTablesWithRelations,
+  TableRelationalConfig,
   type RelationalSchemaConfig,
   type TablesRelationalConfig
 } from 'drizzle-orm/relations';
-import { SQLiteTransaction } from 'drizzle-orm/sqlite-core';
+import { SQLiteSession, SQLiteTable, SQLiteTransaction } from 'drizzle-orm/sqlite-core';
 import { BaseSQLiteDatabase } from 'drizzle-orm/sqlite-core/db';
 import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core/dialect';
+import { RelationalQueryBuilder } from 'drizzle-orm/sqlite-core/query-builders/query';
 import type { DrizzleConfig } from 'drizzle-orm/utils';
 import { toCompilableQuery } from './../utils/compilableQuery.js';
-import { PowerSyncSQLiteTransactionConfig } from './PowerSyncSQLiteBaseSession.js';
+import { PowerSyncSQLiteBaseSession, PowerSyncSQLiteTransactionConfig } from './PowerSyncSQLiteBaseSession.js';
 import { PowerSyncSQLiteSession } from './PowerSyncSQLiteSession.js';
 
 export type DrizzleQuery<T> = { toSQL(): Query; execute(): Promise<T | T[]> };
@@ -54,6 +56,37 @@ export class PowerSyncSQLiteDatabase<
 
     super('async', dialect, session as any, schema as any);
     this.db = db;
+
+    /**
+     * A hack in order to use read locks for `db.query.users.findMany()` etc queries.
+     * We don't currently get queryMetadata for these queries, so we can't use the regular session.
+     * This session always uses read locks.
+     */
+    const querySession = new PowerSyncSQLiteBaseSession(
+      {
+        useReadContext: (callback) => db.readLock(callback),
+        useWriteContext: (callback) => db.readLock(callback)
+      },
+      dialect,
+      schema,
+      {
+        logger
+      }
+    );
+    if (this._.schema) {
+      for (const [tableName, columns] of Object.entries(this._.schema)) {
+        this.query[tableName as keyof typeof this.query] = new RelationalQueryBuilder(
+          'async',
+          schema!.fullSchema,
+          this._.schema,
+          this._.tableNamesMap,
+          schema!.fullSchema[tableName] as SQLiteTable,
+          columns as TableRelationalConfig,
+          dialect,
+          querySession as SQLiteSession<any, any, any, any> as any
+        ) as any;
+      }
+    }
   }
 
   transaction<T>(

packages/drizzle-driver/tests/sqlite/lock-usage.test.ts

@@ -81,6 +81,13 @@ describe('Lock Usage Tests', () => {
       expect(readLockSpy).toHaveBeenCalled();
       expect(writeLockSpy).not.toHaveBeenCalled();
     });
+
+    lockUsageTest('should use readLock for query.findMany()', async ({ db, readLockSpy, writeLockSpy }) => {
+      await db.query.users.findMany();
+
+      expect(readLockSpy).toHaveBeenCalled();
+      expect(writeLockSpy).not.toHaveBeenCalled();
+    });
   });
 
   describe('INSERT queries', () => {