add queryId to CompiledQuery.

...

src/dialect/mysql/mysql-driver.ts

@@ -7,6 +7,7 @@ import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
 import { CompiledQuery } from '../../query-compiler/compiled-query.js'
 import { QueryCompiler } from '../../query-compiler/query-compiler.js'
 import { isFunction, isObject, freeze } from '../../util/object-utils.js'
+import { createQueryId } from '../../util/query-id.js'
 import { extendStackTrace } from '../../util/stack-trace-utils.js'
 import {
   MysqlDialectConfig,
@@ -98,7 +99,10 @@ export class MysqlDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+      compileQuery(
+        parseSavepointCommand('savepoint', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -108,7 +112,10 @@ export class MysqlDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+      compileQuery(
+        parseSavepointCommand('rollback to', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -118,7 +125,10 @@ export class MysqlDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('release savepoint', savepointName)),
+      compileQuery(
+        parseSavepointCommand('release savepoint', savepointName),
+        createQueryId(),
+      ),
     )
   }
 

src/dialect/postgres/postgres-driver.ts

@@ -12,6 +12,7 @@ import {
   RootOperationNode,
 } from '../../query-compiler/query-compiler.js'
 import { isFunction, freeze } from '../../util/object-utils.js'
+import { createQueryId } from '../../util/query-id.js'
 import { extendStackTrace } from '../../util/stack-trace-utils.js'
 import {
   PostgresCursorConstructor,
@@ -91,7 +92,10 @@ export class PostgresDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+      compileQuery(
+        parseSavepointCommand('savepoint', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -101,7 +105,10 @@ export class PostgresDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+      compileQuery(
+        parseSavepointCommand('rollback to', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -111,7 +118,10 @@ export class PostgresDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('release', savepointName)),
+      compileQuery(
+        parseSavepointCommand('release', savepointName),
+        createQueryId(),
+      ),
     )
   }
 

src/dialect/sqlite/sqlite-driver.ts

@@ -8,6 +8,7 @@ import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
 import { CompiledQuery } from '../../query-compiler/compiled-query.js'
 import { QueryCompiler } from '../../query-compiler/query-compiler.js'
 import { freeze, isFunction } from '../../util/object-utils.js'
+import { createQueryId } from '../../util/query-id.js'
 import { SqliteDatabase, SqliteDialectConfig } from './sqlite-dialect-config.js'
 
 export class SqliteDriver implements Driver {
@@ -58,7 +59,10 @@ export class SqliteDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+      compileQuery(
+        parseSavepointCommand('savepoint', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -68,7 +72,10 @@ export class SqliteDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+      compileQuery(
+        parseSavepointCommand('rollback to', savepointName),
+        createQueryId(),
+      ),
     )
   }
 
@@ -78,7 +85,10 @@ export class SqliteDriver implements Driver {
     compileQuery: QueryCompiler['compileQuery'],
   ): Promise<void> {
     await connection.executeQuery(
-      compileQuery(parseSavepointCommand('release', savepointName)),
+      compileQuery(
+        parseSavepointCommand('release', savepointName),
+        createQueryId(),
+      ),
     )
   }
 

src/index.ts

@@ -238,6 +238,7 @@ export {
 } from './util/type-utils.js'
 export * from './util/infer-result.js'
 export { logOnce } from './util/log-once.js'
+export { createQueryId, QueryId } from './util/query-id.js'
 
 export {
   SelectExpression,

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<O = unknown> {
   readonly query: RootOperationNode
+  readonly queryId: QueryId
   readonly sql: string
   readonly parameters: ReadonlyArray<unknown>
 }
@@ -14,6 +16,7 @@ export const CompiledQuery = freeze({
       sql,
       query: RawNode.createWithSql(sql),
       parameters: freeze(parameters),
+      queryId: createQueryId(),
     })
   },
 })

src/query-compiler/default-query-compiler.ts

@@ -115,6 +115,7 @@ import { RefreshMaterializedViewNode } from '../operation-node/refresh-materiali
 import { OrActionNode } from '../operation-node/or-action-node.js'
 import { logOnce } from '../util/log-once.js'
 import { CollateNode } from '../operation-node/collate-node.js'
+import { QueryId } from '../util/query-id.js'
 
 export class DefaultQueryCompiler
   extends OperationNodeVisitor
@@ -127,7 +128,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)
@@ -136,6 +137,7 @@ export class DefaultQueryCompiler
 
     return freeze({
       query: node,
+      queryId,
       sql: this.getSql(),
       parameters: [...this.#parameters],
     })

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
 }

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<T>(

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<Database>({
+    dialect: {
+      createAdapter: () => new PostgresAdapter(),
+      createDriver: () =>
+        new (class extends DummyDriver {
+          async acquireConnection(): Promise<DatabaseConnection> {
+            // @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])
+  }
+})