From 2be61f6f9ec1add4d75aed7d2328230098dab238 Mon Sep 17 00:00:00 2001 From: evgenovalov Date: Fri, 30 Oct 2026 17:42:29 +0100 Subject: [PATCH 1/2] fix(orm): clients derived from a transaction client use their own plugins and auth Since #2497, a client derived from a transaction client ($use, $unuse, $unuseAll, $setAuth, $setOptions) reused the transaction's Kysely instance. Its executor is bound to the original transaction client, so Kysely-level and entity mutation hooks, executor options and the auth handed to the hooks were the transaction client's: tx.$unuse('policy') still enforced access policies, tx.$use(plugin) never ran the plugin's onKyselyQuery/onEntityMutation hooks, and tx.$setAuth(user) evaluated policies with the old auth. Give the derived client its own executor, which takes its connection from the transaction's executor, inside a Kysely Transaction so it stays in the transaction. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RwcqsRcacQHX18hXKbJN4n --- packages/orm/src/client/client-impl.ts | 11 ++-- tests/e2e/orm/client-api/transaction.test.ts | 55 ++++++++++++++++++++ 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/packages/orm/src/client/client-impl.ts b/packages/orm/src/client/client-impl.ts index f6d9f9d76..27cf517a2 100644 --- a/packages/orm/src/client/client-impl.ts +++ b/packages/orm/src/client/client-impl.ts @@ -113,7 +113,9 @@ export class ClientImpl { baseClient.kyselyProps.driver as ZenStackDriver, baseClient.kyselyProps.dialect.createQueryCompiler(), baseClient.kyselyProps.dialect.createAdapter(), - new DefaultConnectionProvider(baseClient.kyselyProps.driver), + baseClient.isTransaction + ? 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 +159,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); } diff --git a/tests/e2e/orm/client-api/transaction.test.ts b/tests/e2e/orm/client-api/transaction.test.ts index 0611aedb3..59f68f6b2 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 { schema } from '../schemas/basic'; +import { PolicyPlugin } from '@zenstackhq/plugin-policy'; import { createTestClient } from '@zenstackhq/testtools'; describe('Client raw query tests', () => { @@ -177,6 +178,60 @@ 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[] = []; + await client.$transaction(async (tx) => { + await tx + .$use({ + id: 'recorder', + onKyselyQuery: ({ query, proceed }) => { + if (query.kind === 'InsertQueryNode') { + calls.push(query.kind); + } + return proceed(query); + }, + onEntityMutation: { + $all: { + beforeEntityMutation: ({ action }) => { + calls.push(action); + }, + }, + }, + }) + .user.create({ data: { email: 'u1@test.com' } }); + await tx.user.create({ data: { email: 'u2@test.com' } }); + }); + + expect(calls).toEqual(['InsertQueryNode', 'create']); + }); }); describe('sequential transaction', () => { From be8ab11b02efab6c5c8d5eadbab20999a75cf55a Mon Sep 17 00:00:00 2001 From: evgenovalov Date: Fri, 30 Oct 2026 19:30:49 +0100 Subject: [PATCH 2/2] fix(orm): fail derived transaction clients once the transaction has ended A client derived from a transaction client took the transaction's connection through Kysely's transaction executor used as a connection provider. That bypasses Kysely's "already committed / rolled back" check, which only guards executeQuery and stream, so a derived client kept past the transaction ran on a connection that had gone back to the pool instead of throwing. $transaction now gives its transaction client a connection provider that derived clients share. It throws Kysely's errors once the transaction callback has settled, before Kysely commits or rolls back and releases the connection. Also type the $use test like the other transaction tests: $use is not part of TransactionClientContract, so the e2e type check rejected it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RwcqsRcacQHX18hXKbJN4n --- packages/orm/src/client/client-impl.ts | 60 +++++++++++++++++++- tests/e2e/orm/client-api/transaction.test.ts | 52 ++++++++++------- 2 files changed, 90 insertions(+), 22 deletions(-) diff --git a/packages/orm/src/client/client-impl.ts b/packages/orm/src/client/client-impl.ts index 27cf517a2..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,8 +163,10 @@ export class ClientImpl { baseClient.kyselyProps.driver as ZenStackDriver, baseClient.kyselyProps.dialect.createQueryCompiler(), baseClient.kyselyProps.dialect.createAdapter(), + // a client derived from a derived client reaches the transaction's connection + // through its base's executor baseClient.isTransaction - ? baseClient.$qb.getExecutor() + ? (baseClient.transactionConnection ?? baseClient.$qb.getExecutor()) : new DefaultConnectionProvider(baseClient.kyselyProps.driver), [], false, @@ -392,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 59f68f6b2..cd7fd082a 100644 --- a/tests/e2e/orm/client-api/transaction.test.ts +++ b/tests/e2e/orm/client-api/transaction.test.ts @@ -1,5 +1,5 @@ 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'; @@ -208,30 +208,44 @@ describe('Client raw query tests', () => { it('$use installs kysely-level and entity mutation hooks on the derived transaction client', async () => { const calls: string[] = []; - await client.$transaction(async (tx) => { - await tx - .$use({ - id: 'recorder', - onKyselyQuery: ({ query, proceed }) => { - if (query.kind === 'InsertQueryNode') { - calls.push(query.kind); - } - return proceed(query); - }, - onEntityMutation: { - $all: { - beforeEntityMutation: ({ action }) => { - calls.push(action); - }, - }, + 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); }, - }) - .user.create({ data: { email: 'u1@test.com' } }); + }, + }, + }); + 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', () => {