diff --git a/packages/orm/src/client/client-impl.ts b/packages/orm/src/client/client-impl.ts index f6d9f9d76..ceef0e777 100644 --- a/packages/orm/src/client/client-impl.ts +++ b/packages/orm/src/client/client-impl.ts @@ -1,6 +1,12 @@ import { invariant, lowerCaseFirst } from '@zenstackhq/common-helpers'; import type { ProcedureDef, SchemaDef } from '@zenstackhq/schema'; -import type { QueryExecutor } from 'kysely'; +import type { + AbortableOperationOptions, + ConnectionProvider, + DatabaseConnection, + QueryExecutor, + TransactionBuilder, +} from 'kysely'; import { CompiledQuery, DefaultConnectionProvider, @@ -61,6 +67,49 @@ function getExecutorNameMapper(executor: QueryExecutor | undefined) { return executor instanceof ZenStackQueryExecutor ? executor.getNameMapper() : undefined; } +/** + * Provides the connection of a transaction started by `$transaction` to the clients derived from its + * transaction client. Once the transaction has ended it fails like Kysely's transaction executor, so a + * derived client kept past the transaction can't run on a connection that went back to the pool. + */ +class TransactionConnectionProvider implements ConnectionProvider { + private transactionExecutor: QueryExecutor | undefined; + private endedAs: 'committed' | 'rolled back' | undefined; + + async run(txBuilder: TransactionBuilder, callback: (tx: Transaction) => Promise): Promise { + try { + return await txBuilder.execute(async (tx) => { + this.transactionExecutor = tx.getExecutor(); + // ended as soon as the callback settles: Kysely commits or rolls back and then + // releases the connection before `execute` settles + try { + const result = await callback(tx); + this.endedAs = 'committed'; + return result; + } catch (err) { + this.endedAs = 'rolled back'; + throw err; + } + }); + } catch (err) { + // Kysely rolls back when the commit fails + this.endedAs = 'rolled back'; + throw err; + } + } + + async provideConnection( + consumer: (connection: DatabaseConnection) => Promise, + options?: AbortableOperationOptions, + ): Promise { + if (this.endedAs) { + throw new Error(`Transaction is already ${this.endedAs}`); + } + invariant(this.transactionExecutor, 'transaction has not started'); + return this.transactionExecutor.provideConnection(consumer, options); + } +} + /** * ZenStack ORM client. */ @@ -79,6 +128,7 @@ export class ClientImpl { public readonly $schema: SchemaDef; readonly kyselyProps: KyselyProps; private auth: AuthType | undefined; + private transactionConnection: TransactionConnectionProvider | undefined; inputValidator: InputValidator; readonly slowQueries: QueryInfo[] = []; @@ -113,7 +163,11 @@ export class ClientImpl { baseClient.kyselyProps.driver as ZenStackDriver, baseClient.kyselyProps.dialect.createQueryCompiler(), baseClient.kyselyProps.dialect.createAdapter(), - new DefaultConnectionProvider(baseClient.kyselyProps.driver), + // a client derived from a derived client reaches the transaction's connection + // through its base's executor + baseClient.isTransaction + ? (baseClient.transactionConnection ?? baseClient.$qb.getExecutor()) + : new DefaultConnectionProvider(baseClient.kyselyProps.driver), [], false, // A name mapper is derived purely from `$schema` and `$options`, so it can be @@ -157,9 +211,10 @@ export class ClientImpl { } if (baseClient?.isTransaction && !executor) { - // if we're creating a derived client from a transaction client and not replacing - // the executor, reuse the current kysely instance to retain the transaction context - this.kysely = baseClient.$qb; + // A client derived from a transaction client runs on the transaction's connection, but + // with its own executor: the executor resolves Kysely-level and entity mutation hooks + // and `$auth` from the client it is bound to. + this.kysely = new Transaction(this.kyselyProps); } else { this.kysely = new Kysely(this.kyselyProps); } @@ -389,9 +444,11 @@ export class ClientImpl { if (options?.isolationLevel) { txBuilder = txBuilder.setIsolationLevel(options.isolationLevel); } - return txBuilder.execute((tx) => { + const transactionConnection = new TransactionConnectionProvider(); + return transactionConnection.run(txBuilder, (tx) => { const txClient = new ClientImpl(this.schema, this.$options, this); txClient.kysely = tx; + txClient.transactionConnection = transactionConnection; return callback(txClient.$contract); }); } diff --git a/tests/e2e/orm/client-api/transaction.test.ts b/tests/e2e/orm/client-api/transaction.test.ts index 0611aedb3..cd7fd082a 100644 --- a/tests/e2e/orm/client-api/transaction.test.ts +++ b/tests/e2e/orm/client-api/transaction.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; -import type { ClientContract } from '@zenstackhq/orm'; +import { definePlugin, type ClientContract } from '@zenstackhq/orm'; import { schema } from '../schemas/basic'; +import { PolicyPlugin } from '@zenstackhq/plugin-policy'; import { createTestClient } from '@zenstackhq/testtools'; describe('Client raw query tests', () => { @@ -177,6 +178,74 @@ describe('Client raw query tests', () => { await expect(client.user.findMany()).toResolveWithLength(0); }); + + it('$unuse removes a kysely-level plugin from the derived transaction client', async () => { + const db = client.$use(new PolicyPlugin()); + await expect( + db.$transaction(async (tx) => { + const unchecked = tx.$unuse('policy'); + await expect(unchecked.user.create({ data: { email: 'u1@test.com' } })).toResolveTruthy(); + await expect(unchecked.user.count()).resolves.toBe(1); + await expect(tx.user.create({ data: { email: 'u2@test.com' } })).toBeRejectedByPolicy(); + throw new Error('rollback'); + }), + ).rejects.toThrow('rollback'); + + await expect(client.user.findMany()).toResolveWithLength(0); + }); + + it('$setAuth evaluates policies with the new auth inside the transaction', async () => { + const db = client.$use(new PolicyPlugin()); + await db.$setAuth({ id: 'u1' }).$transaction(async (tx) => { + const asU2 = tx.$setAuth({ id: 'u2' }); + await expect(asU2.user.create({ data: { id: 'u2', email: 'u2@test.com' } })).toResolveTruthy(); + await expect(asU2.user.create({ data: { id: 'u1', email: 'u1@test.com' } })).toBeRejectedByPolicy(); + await expect(tx.user.create({ data: { id: 'u1', email: 'u1@test.com' } })).toResolveTruthy(); + }); + + await expect(client.user.findMany()).toResolveWithLength(2); + }); + + it('$use installs kysely-level and entity mutation hooks on the derived transaction client', async () => { + const calls: string[] = []; + const recorder = definePlugin({ + id: 'recorder', + onKyselyQuery: ({ query, proceed }) => { + if (query.kind === 'InsertQueryNode') { + calls.push(query.kind); + } + return proceed(query); + }, + onEntityMutation: { + $all: { + beforeEntityMutation: ({ action }) => { + calls.push(action); + }, + }, + }, + }); + await client.$transaction(async (tx) => { + await (tx as any).$use(recorder).user.create({ data: { email: 'u1@test.com' } }); + await tx.user.create({ data: { email: 'u2@test.com' } }); + }); + + expect(calls).toEqual(['InsertQueryNode', 'create']); + }); + + it('a derived transaction client fails once the transaction has ended', async () => { + const committed = await client.$transaction(async (tx) => tx.$setAuth({ id: 'u1' })); + await expect(committed.user.count()).rejects.toThrow('Transaction is already committed'); + await expect(committed.$unuseAll().user.count()).rejects.toThrow('Transaction is already committed'); + + const rolledBack: ClientContract[] = []; + await expect( + client.$transaction(async (tx) => { + rolledBack.push(tx.$unuseAll().$setAuth({ id: 'u1' })); + throw new Error('rollback'); + }), + ).rejects.toThrow('rollback'); + await expect(rolledBack[0]!.user.count()).rejects.toThrow('Transaction is already rolled back'); + }); }); describe('sequential transaction', () => {