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..d89016a3a 100644 --- a/packages/orm/src/client/executor/zenstack-driver.ts +++ b/packages/orm/src/client/executor/zenstack-driver.ts @@ -9,97 +9,126 @@ 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 { + // 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.#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 +141,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); + }); +});