diff --git a/packages/orm/src/client/client-impl.ts b/packages/orm/src/client/client-impl.ts index 74b6305a5..0ba70b525 100644 --- a/packages/orm/src/client/client-impl.ts +++ b/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 = { diff --git a/packages/orm/src/client/executor/connection-mutex.ts b/packages/orm/src/client/executor/connection-mutex.ts new file mode 100644 index 000000000..924a6294b --- /dev/null +++ b/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; + #resolve?: () => void; + + async obtainLock(): Promise { + 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?.(); + } +} diff --git a/packages/orm/src/client/executor/zenstack-driver.ts b/packages/orm/src/client/executor/zenstack-driver.ts index 747acdeda..90e9fef3c 100644 --- a/packages/orm/src/client/executor/zenstack-driver.ts +++ b/packages/orm/src/client/executor/zenstack-driver.ts @@ -1,4 +1,13 @@ -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 @@ -6,6 +15,7 @@ import type { CompiledQuery, DatabaseConnection, Driver, Log, QueryResult, Trans export class ZenStackDriver implements Driver { readonly #driver: Driver; readonly #log: Log; + readonly #connectionMutex?: ConnectionMutex; #initPromise?: Promise; #initDone: boolean; @@ -13,10 +23,14 @@ export class ZenStackDriver implements Driver { #connections = new WeakSet(); #txConnections = new WeakMap Promise>>(); - 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 { @@ -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 { - await this.#driver.releaseConnection(connection); + try { + await this.#driver.releaseConnection(connection); + } finally { + this.#connectionMutex?.releaseLock(); + } } async beginTransaction(connection: DatabaseConnection, settings: TransactionSettings): Promise { diff --git a/tests/regression/test/issue-2788/regression.test.ts b/tests/regression/test/issue-2788/regression.test.ts new file mode 100644 index 000000000..22697fa75 --- /dev/null +++ b/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(); + }); +}); diff --git a/tests/regression/test/issue-2788/schema.ts b/tests/regression/test/issue-2788/schema.ts new file mode 100644 index 000000000..30a40da4c --- /dev/null +++ b/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(); diff --git a/tests/regression/test/issue-2788/schema.zmodel b/tests/regression/test/issue-2788/schema.zmodel new file mode 100644 index 000000000..5877e9b02 --- /dev/null +++ b/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 +} \ No newline at end of file