From 5901ff8786b011f669f91c90285ed023e89ba33b Mon Sep 17 00:00:00 2001 From: kuldeepsinh19 Date: Mon, 2 Nov 2026 22:42:43 +0530 Subject: [PATCH 1/2] fix(orm): respect the log option of clients derived with $setOptions Derived clients reused the base client's ZenStackDriver, which patches pooled connections for logging only if the base client's log enables it and always logs with the base client's Log, so `$setOptions` couldn't enable, change or disable logging. ZenStackDriver now keeps its pool-level state (underlying driver, lifecycle, sqlite mutex, transaction bookkeeping) in a shared object. `withLog()` creates a driver over the same state with its own Log. The acquiring driver's log is recorded per connection, and patched connections look it up on each query. A derived client whose `log` differs from its base uses `withLog()` for its executor and `$qbRaw`. Fixes #2355 --- packages/orm/src/client/client-impl.ts | 25 +- .../src/client/executor/zenstack-driver.ts | 279 ++++++++++-------- tests/regression/test/issue-2355.test.ts | 138 +++++++++ 3 files changed, 316 insertions(+), 126 deletions(-) create mode 100644 tests/regression/test/issue-2355.test.ts diff --git a/packages/orm/src/client/client-impl.ts b/packages/orm/src/client/client-impl.ts index f6d9f9d76..c1d71c6a9 100644 --- a/packages/orm/src/client/client-impl.ts +++ b/packages/orm/src/client/client-impl.ts @@ -104,16 +104,25 @@ export class ClientImpl { // here we use kysely's props constructor so we can pass a custom query executor if (baseClient) { + // A derived client shares the base client's connections. If its `log` option differs, it gets a + // driver over the same connections that logs with its own log. Deliberately an identity check: + // `$use`, `$setAuth` and friends spread the options and keep the same `log`, so they keep the + // base driver. + const logChanged = options.log !== baseClient.$options.log; + const driver = logChanged + ? (baseClient.kyselyProps.driver as ZenStackDriver).withLog(new Log(options.log ?? [])) + : baseClient.kyselyProps.driver; this.kyselyProps = { ...baseClient.kyselyProps, + driver, executor: executor ?? new ZenStackQueryExecutor( this, - baseClient.kyselyProps.driver as ZenStackDriver, + driver as ZenStackDriver, baseClient.kyselyProps.dialect.createQueryCompiler(), baseClient.kyselyProps.dialect.createAdapter(), - new DefaultConnectionProvider(baseClient.kyselyProps.driver), + new DefaultConnectionProvider(driver), [], false, // A name mapper is derived purely from `$schema` and `$options`, so it can be @@ -126,7 +135,17 @@ export class ClientImpl { : undefined, ), }; - this.kyselyRaw = baseClient.kyselyRaw; + this.kyselyRaw = logChanged + ? new Kysely({ + ...this.kyselyProps, + executor: new DefaultQueryExecutor( + baseClient.kyselyProps.dialect.createQueryCompiler(), + baseClient.kyselyProps.dialect.createAdapter(), + new DefaultConnectionProvider(driver), + [], + ), + }) + : baseClient.kyselyRaw; this.auth = baseClient.auth; this.slowQueries = baseClient.slowQueries; } else { diff --git a/packages/orm/src/client/executor/zenstack-driver.ts b/packages/orm/src/client/executor/zenstack-driver.ts index 90e9fef3c..adea4313d 100644 --- a/packages/orm/src/client/executor/zenstack-driver.ts +++ b/packages/orm/src/client/executor/zenstack-driver.ts @@ -9,97 +9,122 @@ import type { } from 'kysely'; import { ConnectionMutex } from './connection-mutex'; +/** + * State shared by all `ZenStackDriver` instances created over the same underlying driver. + */ +export type ZenStackDriverState = { + driver: Driver; + adapter: DialectAdapter; + connectionMutex?: ConnectionMutex; + initPromise?: Promise; + initDone: boolean; + destroyPromise?: Promise; + // connections whose query methods have been patched for logging + loggedConnections: WeakSet; + // log of the driver instance that most recently acquired each connection + connectionLogs: WeakMap; + txConnections: WeakMap Promise>>; +}; + /** * Copied from kysely's RuntimeDriver */ export class ZenStackDriver implements Driver { - readonly #driver: Driver; + readonly #state: ZenStackDriverState; readonly #log: Log; - readonly #connectionMutex?: ConnectionMutex; - - #initPromise?: Promise; - #initDone: boolean; - #destroyPromise?: Promise; - #connections = new WeakSet(); - #txConnections = new WeakMap Promise>>(); - constructor(driver: Driver, log: Log, adapter: DialectAdapter) { - this.#initDone = false; - this.#driver = driver; + constructor(driver: Driver, log: Log, adapter: DialectAdapter, sharedState?: ZenStackDriverState) { this.#log = log; + this.#state = sharedState ?? { + driver, + adapter, + connectionMutex: adapter.supportsMultipleConnections ? undefined : new ConnectionMutex(), + initDone: false, + loggedConnections: new WeakSet(), + connectionLogs: new WeakMap(), + txConnections: new WeakMap(), + }; + } - if (!adapter.supportsMultipleConnections) { - this.#connectionMutex = new ConnectionMutex(); - } + /** + * Creates a driver that shares this driver's underlying driver, connections, lifecycle and + * transaction bookkeeping, but logs with the given log. + */ + withLog(log: Log): ZenStackDriver { + return new ZenStackDriver(this.#state.driver, log, this.#state.adapter, this.#state); } async init(): Promise { - if (this.#destroyPromise) { + const state = this.#state; + if (state.destroyPromise) { throw new Error('driver has already been destroyed'); } - if (!this.#initPromise) { - this.#initPromise = this.#driver + if (!state.initPromise) { + state.initPromise = state.driver .init() .then(() => { - this.#initDone = true; + state.initDone = true; }) .catch((err) => { - this.#initPromise = undefined; + state.initPromise = undefined; return Promise.reject(err); }); } - await this.#initPromise; + await state.initPromise; } async acquireConnection(): Promise { - if (this.#destroyPromise) { + const state = this.#state; + if (state.destroyPromise) { throw new Error('driver has already been destroyed'); } - if (!this.#initDone) { + if (!state.initDone) { await this.init(); } - await this.#connectionMutex?.obtainLock(); + await state.connectionMutex?.obtainLock(); try { - const connection = await this.#driver.acquireConnection(); - if (!this.#connections.has(connection)) { - if (this.#needsLogging()) { - this.#addLogging(connection); - } - - this.#connections.add(connection); + const connection = await state.driver.acquireConnection(); + // Kysely drivers hand out a connection to one holder at a time (until it's released), so the + // connection's queries are logged with the log of the driver that acquired it. Drivers created + // with `withLog` share connections but not logs. + state.connectionLogs.set(connection, this.#log); + if (needsLogging(this.#log) && !state.loggedConnections.has(connection)) { + addLogging(connection, state); + state.loggedConnections.add(connection); } return connection; } catch (error) { - this.#connectionMutex?.releaseLock(); + state.connectionMutex?.releaseLock(); throw error; } } async releaseConnection(connection: DatabaseConnection): Promise { try { - await this.#driver.releaseConnection(connection); + await this.#state.driver.releaseConnection(connection); } finally { - this.#connectionMutex?.releaseLock(); + this.#state.connectionMutex?.releaseLock(); } } async beginTransaction(connection: DatabaseConnection, settings: TransactionSettings): Promise { - const result = await this.#driver.beginTransaction(connection, settings); - this.#txConnections.set(connection, []); + const result = await this.#state.driver.beginTransaction(connection, settings); + this.#state.txConnections.set(connection, []); return result; } async commitTransaction(connection: DatabaseConnection): Promise { + const txConnections = this.#state.txConnections; try { - const result = await this.#driver.commitTransaction(connection); - const callbacks = this.#txConnections.get(connection); + const result = await this.#state.driver.commitTransaction(connection); + const callbacks = txConnections.get(connection); // delete from the map immediately to avoid accidental re-triggering - this.#txConnections.delete(connection); + txConnections.delete(connection); if (callbacks) { for (const callback of callbacks) { try { @@ -112,125 +137,133 @@ export class ZenStackDriver implements Driver { } return result; } catch (err) { - this.#txConnections.delete(connection); + txConnections.delete(connection); throw err; } } async rollbackTransaction(connection: DatabaseConnection): Promise { try { - return await this.#driver.rollbackTransaction(connection); + return await this.#state.driver.rollbackTransaction(connection); } finally { - this.#txConnections.delete(connection); + this.#state.txConnections.delete(connection); } } async destroy(): Promise { - if (!this.#initPromise) { + const state = this.#state; + if (!state.initPromise) { return; } - await this.#initPromise; + await state.initPromise; - if (!this.#destroyPromise) { - this.#destroyPromise = this.#driver.destroy().catch((err) => { - this.#destroyPromise = undefined; + if (!state.destroyPromise) { + state.destroyPromise = state.driver.destroy().catch((err) => { + state.destroyPromise = undefined; return Promise.reject(err); }); } - await this.#destroyPromise; - } - - #needsLogging(): boolean { - return this.#log.isLevelEnabled('query') || this.#log.isLevelEnabled('error'); - } - - // This method monkey patches the database connection's executeQuery method - // by adding logging code around it. Monkey patching is not pretty, but it's - // the best option in this case. - #addLogging(connection: DatabaseConnection): void { - const executeQuery = connection.executeQuery; - const streamQuery = connection.streamQuery; - // eslint-disable-next-line @typescript-eslint/no-this-alias - const dis = this; - - connection.executeQuery = async (compiledQuery): Promise> => { - let caughtError: unknown; - const startTime = performanceNow(); - - try { - return await executeQuery.call(connection, compiledQuery); - } catch (error) { - caughtError = error; - await dis.#logError(error, compiledQuery, startTime); - throw error; - } finally { - if (!caughtError) { - await dis.#logQuery(compiledQuery, startTime); - } - } - }; - - connection.streamQuery = async function* (compiledQuery, chunkSize): AsyncIterableIterator> { - let caughtError: unknown; - const startTime = performanceNow(); - - try { - for await (const result of streamQuery.call(connection, compiledQuery, chunkSize)) { - yield result; - } - } catch (error) { - caughtError = error; - await dis.#logError(error, compiledQuery, startTime); - throw error; - } finally { - if (!caughtError) { - await dis.#logQuery(compiledQuery, startTime, true); - } - } - }; - } - - async #logError(error: unknown, compiledQuery: CompiledQuery, startTime: number): Promise { - await this.#log.error(() => ({ - level: 'error', - error, - query: compiledQuery, - queryDurationMillis: this.#calculateDurationMillis(startTime), - })); - } - - async #logQuery(compiledQuery: CompiledQuery, startTime: number, isStream = false): Promise { - await this.#log.query(() => ({ - level: 'query', - isStream, - query: compiledQuery, - queryDurationMillis: this.#calculateDurationMillis(startTime), - })); - } - - #calculateDurationMillis(startTime: number): number { - return performanceNow() - startTime; + await state.destroyPromise; } isTransactionConnection(connection: DatabaseConnection): boolean { - return this.#txConnections.has(connection); + return this.#state.txConnections.has(connection); } registerTransactionCommitCallback(connection: DatabaseConnection, callback: () => Promise): void { - if (!this.#txConnections.has(connection)) { + const txConnections = this.#state.txConnections; + if (!txConnections.has(connection)) { return; } - const callbacks = this.#txConnections.get(connection); + const callbacks = txConnections.get(connection); if (callbacks) { callbacks.push(callback); } else { - this.#txConnections.set(connection, [callback]); + txConnections.set(connection, [callback]); } } } +function needsLogging(log: Log | undefined): log is Log { + return !!log && (log.isLevelEnabled('query') || log.isLevelEnabled('error')); +} + +// This function monkey patches the database connection's executeQuery method +// by adding logging code around it. Monkey patching is not pretty, but it's +// the best option in this case. The log is looked up per query, so the connection +// logs with the log of whichever driver currently holds it. +function addLogging(connection: DatabaseConnection, state: ZenStackDriverState): void { + const executeQuery = connection.executeQuery; + const streamQuery = connection.streamQuery; + + connection.executeQuery = async (compiledQuery): Promise> => { + const log = state.connectionLogs.get(connection); + if (!needsLogging(log)) { + return executeQuery.call(connection, compiledQuery); + } + + let caughtError: unknown; + const startTime = performanceNow(); + + try { + return await executeQuery.call(connection, compiledQuery); + } catch (error) { + caughtError = error; + await logError(log, error, compiledQuery, startTime); + throw error; + } finally { + if (!caughtError) { + await logQuery(log, compiledQuery, startTime); + } + } + }; + + connection.streamQuery = async function* (compiledQuery, chunkSize): AsyncIterableIterator> { + const log = state.connectionLogs.get(connection); + if (!needsLogging(log)) { + yield* streamQuery.call(connection, compiledQuery, chunkSize); + return; + } + + let caughtError: unknown; + const startTime = performanceNow(); + + try { + for await (const result of streamQuery.call(connection, compiledQuery, chunkSize)) { + yield result; + } + } catch (error) { + caughtError = error; + await logError(log, error, compiledQuery, startTime); + throw error; + } finally { + if (!caughtError) { + await logQuery(log, compiledQuery, startTime, true); + } + } + }; +} + +async function logError(log: Log, error: unknown, compiledQuery: CompiledQuery, startTime: number): Promise { + await log.error(() => ({ + level: 'error', + error, + query: compiledQuery, + queryDurationMillis: performanceNow() - startTime, + })); +} + +async function logQuery(log: Log, compiledQuery: CompiledQuery, startTime: number, isStream = false): Promise { + await log.query(() => ({ + level: 'query', + isStream, + query: compiledQuery, + queryDurationMillis: performanceNow() - startTime, + })); +} + export function performanceNow() { if (typeof performance !== 'undefined' && typeof performance.now === 'function') { return performance.now(); diff --git a/tests/regression/test/issue-2355.test.ts b/tests/regression/test/issue-2355.test.ts new file mode 100644 index 000000000..0161444c0 --- /dev/null +++ b/tests/regression/test/issue-2355.test.ts @@ -0,0 +1,138 @@ +import { createTestClient } from '@zenstackhq/testtools'; +import { afterEach, describe, expect, it } from 'vitest'; + +type LogEvent = { level: 'query' | 'error'; query: { sql: string } }; + +// https://github.com/zenstackhq/zenstack/issues/2355 +describe('Regression for issue 2355', () => { + const schema = ` +model User { + id Int @id @default(autoincrement()) + email String @unique +} + `; + + let db: any; + + afterEach(async () => { + await db?.$disconnect(); + }); + + function collect() { + const events: LogEvent[] = []; + const log = (event: LogEvent) => { + events.push(event); + }; + return { events, log, queries: () => events.filter((e) => e.level === 'query').map((e) => e.query.sql) }; + } + + it('enables logging with $setOptions', async () => { + db = await createTestClient(schema); + const logger = collect(); + + const client = db.$setOptions({ ...db.$options, log: logger.log }); + await client.user.findMany(); + + expect(logger.queries().some((sql) => /select/i.test(sql))).toBe(true); + }); + + it('logs to the logger of the client that ran the query', async () => { + const base = collect(); + db = await createTestClient(schema, { log: base.log }); + // ignore the queries run while setting up the test database + base.events.length = 0; + const derived = collect(); + const client = db.$setOptions({ ...db.$options, log: derived.log }); + + await client.user.create({ data: { email: 'derived@test.com' } }); + expect(derived.queries().some((sql) => /insert/i.test(sql))).toBe(true); + expect(base.events).toHaveLength(0); + + await db.user.count(); + expect(base.queries().some((sql) => /count/i.test(sql))).toBe(true); + expect(derived.queries().some((sql) => /count/i.test(sql))).toBe(false); + }); + + it('disables logging with $setOptions', async () => { + const base = collect(); + db = await createTestClient(schema, { log: base.log }); + // ignore the queries run while setting up the test database + base.events.length = 0; + + const { log: _log, ...options } = db.$options; + await db.$setOptions(options).user.findMany(); + + expect(base.events).toHaveLength(0); + }); + + it('logs queries of transactions started from the derived client', async () => { + db = await createTestClient(schema); + const logger = collect(); + const client = db.$setOptions({ ...db.$options, log: logger.log }); + + await client.$transaction(async (tx: any) => { + await tx.user.create({ data: { email: 'tx@test.com' } }); + }); + + expect(logger.queries().some((sql) => /insert/i.test(sql))).toBe(true); + }); + + it('logs failed queries of the derived client as errors', async () => { + db = await createTestClient(schema); + const logger = collect(); + const client = db.$setOptions({ ...db.$options, log: logger.log }); + + await client.user.create({ data: { email: 'dup@test.com' } }); + // violates the unique constraint on `email` + await expect(client.user.create({ data: { email: 'dup@test.com' } })).rejects.toThrow(); + + expect(logger.events.some((e) => e.level === 'error')).toBe(true); + }); + + it('logs raw query builder queries of the derived client', async () => { + db = await createTestClient(schema); + const logger = collect(); + const client = db.$setOptions({ ...db.$options, log: logger.log }); + + await client.$qbRaw.selectFrom('User').selectAll().execute(); + + expect(logger.queries().some((sql) => /select/i.test(sql))).toBe(true); + }); + + it('logs concurrent queries to the logger of the client that ran them', async () => { + const base = collect(); + db = await createTestClient(schema, { log: base.log }); + base.events.length = 0; + const derived = collect(); + const client = db.$setOptions({ ...db.$options, log: derived.log }); + + await Promise.all( + Array.from({ length: 10 }, (_, i) => + i % 2 === 0 + ? db.user.create({ data: { email: `base${i}@test.com` } }) + : client.user.create({ data: { email: `derived${i}@test.com` } }), + ), + ); + + const params = (events: LogEvent[]) => events.flatMap((e) => (e.query as any).parameters ?? []).join(' '); + expect(params(base.events)).not.toMatch(/derived/); + expect(params(derived.events)).not.toMatch(/base/); + expect(params(base.events)).toMatch(/base0/); + expect(params(derived.events)).toMatch(/derived1/); + }); + + it('keeps the logger for clients derived without changing it', async () => { + const base = collect(); + db = await createTestClient(schema, { log: base.log }); + base.events.length = 0; + + await db.$setAuth({ id: 1 }).user.findMany(); + expect(base.queries().some((sql) => /select/i.test(sql))).toBe(true); + + base.events.length = 0; + await db.$transaction(async (tx: any) => { + await tx.user.count(); + }); + expect(base.queries().some((sql) => /count/i.test(sql))).toBe(true); + }); +}); From 07264a97a2a02f043cc6310d61d7cfc8c8555821 Mon Sep 17 00:00:00 2001 From: kuldeepsinh19 Date: Tue, 3 Nov 2026 08:40:35 +0530 Subject: [PATCH 2/2] fix(orm): clear a connection's log before releasing it A pool can hand a released connection to a queued acquisition during release (pg-pool resolves the waiting request right away), and queries run on it before that acquisition records its own log, such as a dialect's onReserveConnection hook, were logged with the previous holder's log. Drop the recorded log before releasing the connection. Reported by CodeRabbit in review. --- packages/orm/src/client/executor/zenstack-driver.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/orm/src/client/executor/zenstack-driver.ts b/packages/orm/src/client/executor/zenstack-driver.ts index adea4313d..d89016a3a 100644 --- a/packages/orm/src/client/executor/zenstack-driver.ts +++ b/packages/orm/src/client/executor/zenstack-driver.ts @@ -105,6 +105,10 @@ export class ZenStackDriver implements Driver { } async releaseConnection(connection: DatabaseConnection): Promise { + // Forget the holder's log first: a pool (e.g. pg-pool) can hand the connection to a queued + // acquisition during release, and queries run on it before that acquisition records its own log + // (such as a dialect's `onReserveConnection` hook) must not be logged with the previous holder's log. + this.#state.connectionLogs.delete(connection); try { await this.#state.driver.releaseConnection(connection); } finally {