fix(orm): add mutex for single-connection adapters

packages/orm/src/client/client-impl.ts

@@ -122,9 +122,13 @@ export class ClientImpl {
             this.auth = baseClient.auth;
             this.slowQueries = baseClient.slowQueries;
         } else {
-            const driver = new ZenStackDriver(options.dialect.createDriver(), new Log(this.$options.log ?? []));
-            const compiler = options.dialect.createQueryCompiler();
             const adapter = options.dialect.createAdapter();
+            const driver = new ZenStackDriver(
+                options.dialect.createDriver(),
+                new Log(this.$options.log ?? []),
+                adapter,
+            );
+            const compiler = options.dialect.createQueryCompiler();
             const connectionProvider = new DefaultConnectionProvider(driver);
 
             this.kyselyProps = {

packages/orm/src/client/executor/connection-mutex.ts

@@ -0,0 +1,30 @@
+/**
+ * This mutex is used to ensure that only one operation at a time can
+ * acquire a connection from the driver. This is necessary when the
+ * driver only has a single connection, like SQLite and PGlite.
+ *
+ * @see {@link https://github.com/kysely-org/kysely/blob/478ec67b2de2568f5590a015d3e120644e81bd87/src/driver/connection-mutex.ts|Kysely Source}
+ */
+export class ConnectionMutex {
+    #promise?: Promise<void>;
+    #resolve?: () => void;
+
+    async obtainLock(): Promise<void> {
+        while (this.#promise) {
+            await this.#promise;
+        }
+
+        this.#promise = new Promise((resolve) => {
+            this.#resolve = resolve;
+        });
+    }
+
+    releaseLock(): void {
+        const resolve = this.#resolve;
+
+        this.#promise = undefined;
+        this.#resolve = undefined;
+
+        resolve?.();
+    }
+}

packages/orm/src/client/executor/zenstack-driver.ts

@@ -1,22 +1,36 @@
-import type { CompiledQuery, DatabaseConnection, Driver, Log, QueryResult, TransactionSettings } from 'kysely';
+import type {
+    CompiledQuery,
+    DatabaseConnection,
+    DialectAdapter,
+    Driver,
+    Log,
+    QueryResult,
+    TransactionSettings,
+} from 'kysely';
+import { ConnectionMutex } from './connection-mutex';
 
 /**
  * Copied from kysely's RuntimeDriver
  */
 export class ZenStackDriver implements Driver {
     readonly #driver: Driver;
     readonly #log: Log;
+    readonly #connectionMutex?: ConnectionMutex;
 
     #initPromise?: Promise<void>;
     #initDone: boolean;
     #destroyPromise?: Promise<void>;
     #connections = new WeakSet<DatabaseConnection>();
     #txConnections = new WeakMap<DatabaseConnection, Array<() => Promise<unknown>>>();
 
-    constructor(driver: Driver, log: Log) {
+    constructor(driver: Driver, log: Log, adapter: DialectAdapter) {
         this.#initDone = false;
         this.#driver = driver;
         this.#log = log;
+
+        if (!adapter.supportsMultipleConnections) {
+            this.#connectionMutex = new ConnectionMutex();
+        }
     }
 
     async init(): Promise<void> {
@@ -48,21 +62,30 @@ export class ZenStackDriver implements Driver {
             await this.init();
         }
 
-        const connection = await this.#driver.acquireConnection();
+        await this.#connectionMutex?.obtainLock();
 
-        if (!this.#connections.has(connection)) {
-            if (this.#needsLogging()) {
-                this.#addLogging(connection);
-            }
+        try {
+            const connection = await this.#driver.acquireConnection();
+            if (!this.#connections.has(connection)) {
+                if (this.#needsLogging()) {
+                    this.#addLogging(connection);
+                }
 
-            this.#connections.add(connection);
+                this.#connections.add(connection);
+            }
+            return connection;
+        } catch (error) {
+            this.#connectionMutex?.releaseLock();
+            throw error;
         }
-
-        return connection;
     }
 
     async releaseConnection(connection: DatabaseConnection): Promise<void> {
-        await this.#driver.releaseConnection(connection);
+        try {
+            await this.#driver.releaseConnection(connection);
+        } finally {
+            this.#connectionMutex?.releaseLock();
+        }
     }
 
     async beginTransaction(connection: DatabaseConnection, settings: TransactionSettings): Promise<void> {

tests/regression/test/issue-2788/regression.test.ts

@@ -0,0 +1,68 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+import { schema } from './schema';
+
+// https://github.com/zenstackhq/zenstack/issues/2788
+
+describe('Regression for issue #2788', () => {
+    it('does not error during concurrent upserts', async () => {
+        const db = await createTestClient(schema);
+        const user = await db.user.create({
+            data: {
+                email: 'test@zenstack.dev',
+                posts: {
+                    create: [
+                        {
+                            title: 'Post 1',
+                            content: 'This is a test post',
+                        },
+                    ],
+                },
+            },
+            include: { posts: true },
+        });
+
+        let posts: any[] = await db.post.findMany();
+
+        posts[0].title = 'Post 1 Updated';
+
+        posts.push({
+            id: 'cmstai1q2000104js3i7s2d8l',
+            title: 'Post 2',
+            content: 'This is a test post',
+            authorId: user.id,
+        });
+
+        await Promise.all(
+            posts.map(async (p) => {
+                await db.post.upsert({
+                    where: { id: p.id },
+                    update: { ...p },
+                    create: { title: p.title!, content: p.content!, authorId: p.authorId! },
+                });
+            }),
+        );
+
+        posts = await db.post.findMany();
+
+        expect(posts.find((p) => p.title === 'Post 1 Updated')).toMatchObject({
+            title: 'Post 1 Updated',
+            content: 'This is a test post',
+            published: false,
+        });
+
+        expect(posts.find((p) => p.title === 'Post 2')).toMatchObject({
+            title: 'Post 2',
+            content: 'This is a test post',
+            published: false,
+        });
+
+        await expect(
+            db.post.findUnique({
+                where: {
+                    id: 'cmstai1q2000104js3i7s2d8l',
+                },
+            }),
+        ).resolves.toBeNull();
+    });
+});

tests/regression/test/issue-2788/schema.ts

@@ -0,0 +1,103 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaDef, type AttributeApplication, type FieldDefault, ExpressionUtils } from "@zenstackhq/schema";
+export class SchemaType implements SchemaDef {
+    provider = {
+        type: "sqlite"
+    } as const;
+    models = {
+        User: {
+            name: "User",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "String",
+                    id: true,
+                    attributes: [{ name: "@id" }, { name: "@default", args: [{ name: "value", value: ExpressionUtils.call("cuid") }] }] as readonly AttributeApplication[],
+                    default: ExpressionUtils.call("cuid") as FieldDefault
+                },
+                email: {
+                    name: "email",
+                    type: "String",
+                    unique: true,
+                    attributes: [{ name: "@unique" }, { name: "@email" }, { name: "@length", args: [{ name: "min", value: ExpressionUtils.literal(6) }, { name: "max", value: ExpressionUtils.literal(32) }] }] as readonly AttributeApplication[]
+                },
+                posts: {
+                    name: "posts",
+                    type: "Post",
+                    array: true,
+                    relation: { opposite: "author" }
+                }
+            },
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "String" },
+                email: { type: "String" }
+            }
+        },
+        Post: {
+            name: "Post",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "String",
+                    id: true,
+                    attributes: [{ name: "@id" }, { name: "@default", args: [{ name: "value", value: ExpressionUtils.call("cuid") }] }] as readonly AttributeApplication[],
+                    default: ExpressionUtils.call("cuid") as FieldDefault
+                },
+                createdAt: {
+                    name: "createdAt",
+                    type: "DateTime",
+                    attributes: [{ name: "@default", args: [{ name: "value", value: ExpressionUtils.call("now") }] }] as readonly AttributeApplication[],
+                    default: ExpressionUtils.call("now") as FieldDefault
+                },
+                updatedAt: {
+                    name: "updatedAt",
+                    type: "DateTime",
+                    updatedAt: true,
+                    attributes: [{ name: "@updatedAt" }] as readonly AttributeApplication[]
+                },
+                title: {
+                    name: "title",
+                    type: "String",
+                    attributes: [{ name: "@length", args: [{ name: "min", value: ExpressionUtils.literal(1) }, { name: "max", value: ExpressionUtils.literal(256) }] }] as readonly AttributeApplication[]
+                },
+                content: {
+                    name: "content",
+                    type: "String"
+                },
+                published: {
+                    name: "published",
+                    type: "Boolean",
+                    attributes: [{ name: "@default", args: [{ name: "value", value: ExpressionUtils.literal(false) }] }] as readonly AttributeApplication[],
+                    default: false as FieldDefault
+                },
+                author: {
+                    name: "author",
+                    type: "User",
+                    attributes: [{ name: "@relation", args: [{ name: "fields", value: ExpressionUtils.array("String", [ExpressionUtils.field("authorId")]) }, { name: "references", value: ExpressionUtils.array("String", [ExpressionUtils.field("id")]) }, { name: "onDelete", value: ExpressionUtils.literal("Cascade") }] }] as readonly AttributeApplication[],
+                    relation: { opposite: "posts", fields: ["authorId"], references: ["id"], onDelete: "Cascade" }
+                },
+                authorId: {
+                    name: "authorId",
+                    type: "String",
+                    foreignKeyFor: [
+                        "author"
+                    ] as readonly string[]
+                }
+            },
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "String" }
+            }
+        }
+    } as const;
+    authType = "User" as const;
+    plugins = {};
+}
+export const schema = new SchemaType();

tests/regression/test/issue-2788/schema.zmodel

@@ -0,0 +1,26 @@
+// This is a sample model to get you started.
+
+/// A sample data source using local sqlite db.
+datasource db {
+    provider = 'sqlite'
+    url = 'file:./dev.db'
+}
+
+/// User model
+model User {
+    id       String @id @default(cuid())
+    email    String @unique @email @length(6, 32)
+    posts    Post[]
+}
+
+/// Post model
+model Post {
+    id        String   @id @default(cuid())
+    createdAt DateTime @default(now())
+    updatedAt DateTime @updatedAt
+    title     String   @length(1, 256)
+    content   String
+    published Boolean  @default(false)
+    author    User     @relation(fields: [authorId], references: [id], onDelete: Cascade)
+    authorId  String
+}