fix(orm): clients derived from a transaction client use their own plugins and auth

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<DB, T>(txBuilder: TransactionBuilder<DB>, callback: (tx: Transaction<DB>) => Promise<T>): Promise<T> {
+        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<T>(
+        consumer: (connection: DatabaseConnection) => Promise<T>,
+        options?: AbortableOperationOptions,
+    ): Promise<T> {
+        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<SchemaDef> | undefined;
+    private transactionConnection: TransactionConnectionProvider | undefined;
     inputValidator: InputValidator<SchemaDef>;
     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);
             });
         }

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<typeof schema>[] = [];
+            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', () => {