Add tests to check Drizzle query concurrency.

Author: rkistner

packages/node/tests/DrizzleNode.test.ts

@@ -1,7 +1,7 @@
 import { sqliteTable, text } from 'drizzle-orm/sqlite-core';
 import { eq, relations } from 'drizzle-orm';
 
-import { databaseTest } from './utils';
+import { customDatabaseTest, databaseTest } from './utils';
 import { wrapPowerSyncWithDrizzle } from '@powersync/drizzle-driver';
 import { PowerSyncDatabase } from '../lib';
 import { expect } from 'vitest';
@@ -53,6 +53,15 @@ databaseTest('should retrieve a list with todos', async ({ database }) => {
   expect(result).toEqual([{ id: '1', name: 'list 1', todos: [{ id: '33', content: 'Post content', list_id: '1' }] }]);
 });
 
+databaseTest('insert returning', async ({ database }) => {
+  const db = await setupDrizzle(database);
+
+  // This is a special case since it's an insert query that returns values
+  const result = await db.insert(drizzleLists).values({ id: '2', name: 'list 2' }).returning();
+
+  expect(result).toEqual([{ id: '2', name: 'list 2' }]);
+});
+
 databaseTest('should retrieve a todo with its list', async ({ database }) => {
   const db = await setupDrizzle(database);
 
@@ -97,3 +106,126 @@ databaseTest('should return a list and todos using fullJoin', async ({ database
   expect(result[0].lists).toEqual({ id: '1', name: 'list 1' });
   expect(result[0].todos).toEqual({ id: '33', content: 'Post content', list_id: '1' });
 });
+
+customDatabaseTest({ database: { readWorkerCount: 2 } as any })(
+  'should execute transactions concurrently',
+  {},
+  async ({ database }) => {
+    // This test opens one write transaction and two read transactions, testing that they can all execute
+    // concurrently.
+    const db = await setupDrizzle(database);
+
+    const openedWrite = deferred();
+    const openedRead1 = deferred();
+    const openedRead2 = deferred();
+    const completedWrite = deferred();
+
+    const t1 = db.transaction(async (tx) => {
+      await tx.insert(drizzleLists).values({ id: '2', name: 'list 2' });
+
+      openedWrite.resolve();
+
+      await openedRead1.promise;
+      await openedRead2.promise;
+
+      completedWrite.resolve();
+    });
+
+    await openedWrite.promise;
+
+    const t2 = db.transaction(
+      async (tx) => {
+        const result = await tx.query.lists.findMany();
+        expect(result).toEqual([{ id: '1', name: 'list 1' }]);
+        openedRead1.resolve();
+        await completedWrite.promise;
+      },
+      { accessMode: 'read only' }
+    );
+
+    const t3 = db.transaction(
+      async (tx) => {
+        await openedRead1.promise;
+        const result = await tx.query.lists.findMany();
+        expect(result).toEqual([{ id: '1', name: 'list 1' }]);
+
+        openedRead2.resolve();
+        await completedWrite.promise;
+      },
+      { accessMode: 'read only' }
+    );
+
+    await Promise.all([t1, t2, t3]);
+
+    const result = await db.query.lists.findMany();
+    expect(result).toEqual([
+      { id: '1', name: 'list 1' },
+      { id: '2', name: 'list 2' }
+    ]);
+  }
+);
+
+customDatabaseTest({ database: { readWorkerCount: 2 } as any })(
+  'should execute select queries concurrently',
+  async ({ database }) => {
+    // This test opens one write transaction and two read transactions, testing that they can all execute
+    // concurrently.
+    const db = await setupDrizzle(database);
+
+    const openedWrite = deferred();
+    const completedRead = deferred();
+    const completedWrite = deferred();
+
+    const t1 = db.transaction(async (tx) => {
+      await tx.insert(drizzleLists).values({ id: '2', name: 'list 2' });
+
+      openedWrite.resolve();
+
+      await completedRead.promise;
+
+      completedWrite.resolve();
+    });
+
+    await openedWrite.promise;
+
+    const result1 = await db.select().from(drizzleLists);
+    expect(result1).toEqual([{ id: '1', name: 'list 1' }]);
+
+    const result2 = await db
+      .select()
+      .from(drizzleLists)
+      .fullJoin(drizzleTodos, eq(drizzleLists.id, drizzleTodos.list_id));
+
+    expect(result2[0].lists).toEqual({ id: '1', name: 'list 1' });
+    expect(result2[0].todos).toEqual({ id: '33', content: 'Post content', list_id: '1' });
+
+    // Note: This case is not supported yet (drizzle 0.44.7), since it doesn't set
+    // queryMetadata for these queries
+    // const result3 = await db.query.lists.findMany();
+    // expect(result3).toEqual([{ id: '1', name: 'list 1' }]);
+
+    completedRead.resolve();
+
+    await t1;
+
+    const resultAfter = await db.query.lists.findMany();
+    expect(resultAfter).toEqual([
+      { id: '1', name: 'list 1' },
+      { id: '2', name: 'list 2' }
+    ]);
+  }
+);
+
+function deferred<T = void>() {
+  let resolve: (value: T) => void;
+  let reject: (err) => void;
+  const promise = new Promise<T>((a, b) => {
+    resolve = a;
+    reject = b;
+  });
+  return {
+    promise,
+    resolve,
+    reject
+  };
+}

packages/node/tests/utils.ts

@@ -14,6 +14,7 @@ import {
   PowerSyncBackendConnector,
   PowerSyncCredentials,
   PowerSyncDatabase,
+  PowerSyncDatabaseOptions,
   Schema,
   StreamingSyncCheckpoint,
   StreamingSyncLine,
@@ -65,27 +66,32 @@ export async function createDatabase(
 
   const database = new PowerSyncDatabase({
     schema: AppSchema,
+    ...options,
+    logger: defaultLogger,
     database: {
       dbFilename: 'test.db',
       dbLocation: tmpdir,
       // Using a single read worker (instead of multiple, the default) seems to improve the reliability of tests in GH
       // actions. So far, we've not been able to reproduce these failures locally.
-      readWorkerCount: 1
-    },
-    logger: defaultLogger,
-    ...options
+      readWorkerCount: 1,
+      ...options.database
+    }
   });
   await database.init();
   return database;
 }
 
-export const databaseTest = tempDirectoryTest.extend<{ database: PowerSyncDatabase }>({
-  database: async ({ tmpdir }, use) => {
-    const db = await createDatabase(tmpdir);
-    await use(db);
-    await db.close();
-  }
-});
+export const customDatabaseTest = (options?: Partial<NodePowerSyncDatabaseOptions>) => {
+  return tempDirectoryTest.extend<{ database: PowerSyncDatabase }>({
+    database: async ({ tmpdir }, use) => {
+      const db = await createDatabase(tmpdir, options);
+      await use(db);
+      await db.close();
+    }
+  });
+};
+
+export const databaseTest = customDatabaseTest();
 
 // TODO: Unify this with the test setup for the web SDK.
 export const mockSyncServiceTest = tempDirectoryTest.extend<{