diff --git a/src/index.ts b/src/index.ts index b467d871c..c623940b0 100644 --- a/src/index.ts +++ b/src/index.ts @@ -232,6 +232,7 @@ export { } from './util/type-utils.js' export * from './util/infer-result.js' export { logOnce } from './util/log-once.js' +export { QueryId } from './util/query-id.js' export { SelectExpression, diff --git a/src/query-compiler/compiled-query.ts b/src/query-compiler/compiled-query.ts index b7a7dfc5d..f649a83e4 100644 --- a/src/query-compiler/compiled-query.ts +++ b/src/query-compiler/compiled-query.ts @@ -1,9 +1,11 @@ import { RawNode } from '../operation-node/raw-node.js' import { freeze } from '../util/object-utils.js' +import { createQueryId, QueryId } from '../util/query-id.js' import { RootOperationNode } from './query-compiler.js' export interface CompiledQuery { readonly query: RootOperationNode + readonly queryId: QueryId readonly sql: string readonly parameters: ReadonlyArray } @@ -14,6 +16,7 @@ export const CompiledQuery = freeze({ sql, query: RawNode.createWithSql(sql), parameters: freeze(parameters), + queryId: createQueryId(), }) }, }) diff --git a/src/query-compiler/default-query-compiler.ts b/src/query-compiler/default-query-compiler.ts index 5a04313d6..39a4200df 100644 --- a/src/query-compiler/default-query-compiler.ts +++ b/src/query-compiler/default-query-compiler.ts @@ -112,6 +112,7 @@ import { FetchNode } from '../operation-node/fetch-node.js' import { TopNode } from '../operation-node/top-node.js' import { OutputNode } from '../operation-node/output-node.js' import { RefreshMaterializedViewNode } from '../operation-node/refresh-materialized-view-node.js' +import { QueryId } from '../util/query-id.js' export class DefaultQueryCompiler extends OperationNodeVisitor @@ -124,7 +125,7 @@ export class DefaultQueryCompiler return this.#parameters.length } - compileQuery(node: RootOperationNode): CompiledQuery { + compileQuery(node: RootOperationNode, queryId: QueryId): CompiledQuery { this.#sql = '' this.#parameters = [] this.nodeStack.splice(0, this.nodeStack.length) @@ -133,6 +134,7 @@ export class DefaultQueryCompiler return freeze({ query: node, + queryId, sql: this.getSql(), parameters: [...this.#parameters], }) diff --git a/src/query-compiler/query-compiler.ts b/src/query-compiler/query-compiler.ts index 8a14b7a81..bd9d99a7f 100644 --- a/src/query-compiler/query-compiler.ts +++ b/src/query-compiler/query-compiler.ts @@ -13,6 +13,7 @@ import { MergeQueryNode } from '../operation-node/merge-query-node.js' import { QueryNode } from '../operation-node/query-node.js' import { RawNode } from '../operation-node/raw-node.js' import { RefreshMaterializedViewNode } from '../operation-node/refresh-materialized-view-node.js' +import { QueryId } from '../util/query-id.js' import { CompiledQuery } from './compiled-query.js' export type RootOperationNode = @@ -36,5 +37,5 @@ export type RootOperationNode = * a `QueryCompiler` compiles a query expressed as a tree of `OperationNodes` into SQL. */ export interface QueryCompiler { - compileQuery(node: RootOperationNode): CompiledQuery + compileQuery(node: RootOperationNode, queryId: QueryId): CompiledQuery } diff --git a/src/query-executor/default-query-executor.ts b/src/query-executor/default-query-executor.ts index 65b65c5e3..9f106d8dc 100644 --- a/src/query-executor/default-query-executor.ts +++ b/src/query-executor/default-query-executor.ts @@ -8,6 +8,7 @@ import { import { KyselyPlugin } from '../plugin/kysely-plugin.js' import { QueryExecutorBase } from './query-executor-base.js' import { DialectAdapter } from '../dialect/dialect-adapter.js' +import { QueryId } from '../util/query-id.js' export class DefaultQueryExecutor extends QueryExecutorBase { #compiler: QueryCompiler @@ -31,8 +32,8 @@ export class DefaultQueryExecutor extends QueryExecutorBase { return this.#adapter } - compileQuery(node: RootOperationNode): CompiledQuery { - return this.#compiler.compileQuery(node) + compileQuery(node: RootOperationNode, queryId: QueryId): CompiledQuery { + return this.#compiler.compileQuery(node, queryId) } provideConnection( diff --git a/test/node/src/query-id.test.ts b/test/node/src/query-id.test.ts new file mode 100644 index 000000000..e70b9be45 --- /dev/null +++ b/test/node/src/query-id.test.ts @@ -0,0 +1,75 @@ +import { expect } from 'chai' +import { + DatabaseConnection, + DummyDriver, + Kysely, + PostgresAdapter, + PostgresIntrospector, + PostgresQueryCompiler, + RootOperationNode, + QueryId, +} from '../../..' +import { Database } from './test-setup' + +describe('queryId', () => { + const visits = new Map() + + const db = new Kysely({ + dialect: { + createAdapter: () => new PostgresAdapter(), + createDriver: () => + new (class extends DummyDriver { + async acquireConnection(): Promise { + // @ts-ignore + return { + executeQuery: async ({ queryId }) => { + checkIn(queryId, 'connection.executeQuery') + + return { + rows: [], + } + }, + } + } + })(), + createIntrospector: (db) => new PostgresIntrospector(db), + createQueryCompiler: () => + new (class SomeCompiler extends PostgresQueryCompiler { + compileQuery(node: RootOperationNode, queryId: QueryId) { + checkIn(queryId, 'compiler.compileQuery') + + return super.compileQuery(node, queryId) + } + })(), + }, + plugins: [ + { + transformQuery: (args) => { + checkIn(args.queryId, 'plugin.transformQuery') + + return args.node + }, + transformResult: async (args) => { + checkIn(args.queryId, 'plugin.transformResult') + + return args.result + }, + }, + ], + }) + + it('should pass query id around, allowing async communication between compilers, plugins and connections', async () => { + await db.selectFrom('person').where('id', '=', 1).execute() + + expect(Array.from(visits.values())[0]).to.deep.equal([ + 'plugin.transformQuery', + 'compiler.compileQuery', + 'connection.executeQuery', + 'plugin.transformResult', + ]) + }) + + function checkIn(queryId: QueryId, place: string): void { + visits.set(queryId, [...(visits.get(queryId) || []), place]) + } +})