fix(kysely-generic-sqlite): bind queryId into worker message

subframe7536 · subframe7536 · commit 6b18db47ca0d · 2026-06-29T20:08:36.000+08:00
diff --git a/packages/dialect-generic-sqlite/src/worker/driver.ts b/packages/dialect-generic-sqlite/src/worker/driver.ts
@@ -40,7 +40,7 @@ export class GenericSqliteWorkerDriver<
       this.worker.postMessage([initEvent, exec.data || ({} as any)] satisfies InitMsg<R>)
 
       await new Promise<void>((resolve, reject) => {
-        this.mitt!.once(initEvent, (_, err) => (err ? reject(err) : resolve()))
+        this.mitt!.once(initEvent, (_qid, _data, err) => (err ? reject(err) : resolve()))
       })
 
       this.conn = new GenericSqliteWorkerConnection(this.worker, this.mitt)
@@ -54,7 +54,7 @@ export class GenericSqliteWorkerDriver<
     }
     this.worker.postMessage([closeEvent] satisfies CloseMsg)
     return new Promise<void>((resolve, reject) =>
-      this.mitt?.once(closeEvent, (_, err) => (err ? reject(err) : resolve())),
+      this.mitt?.once(closeEvent, (_qid, _data, err) => (err ? reject(err) : resolve())),
     ).finally(() => {
       this.worker?.terminate()
       this.mitt?.off()
@@ -64,13 +64,71 @@ export class GenericSqliteWorkerDriver<
 }
 
 class GenericSqliteWorkerConnection implements DatabaseConnection {
+  private pendingRuns = new Map<
+    string,
+    { resolve: (data: any) => void; reject: (err: any) => void }
+  >()
+  private pendingStreams = new Map<
+    string,
+    { resolve: (data: any) => void; reject: (err: any) => void }
+  >()
+
   constructor(
     readonly worker: IGenericWorker,
     readonly mitt: IGenericEventEmitter,
-  ) {}
+  ) {
+    // Central dispatcher — register once, route by queryId
+    this.mitt.on(runEvent, (queryId: string | undefined, data: any, err: unknown) => {
+      if (!queryId) {
+        return
+      }
+      const pending = this.pendingRuns.get(queryId)
+      if (!pending) {
+        return
+      }
+      this.pendingRuns.delete(queryId)
+      if (err) {
+        pending.reject(err)
+      } else {
+        pending.resolve(data)
+      }
+    })
+
+    this.mitt.on(dataEvent, (queryId: string | undefined, data: any, err: unknown) => {
+      if (!queryId) {
+        return
+      }
+      const state = this.pendingStreams.get(queryId)
+      if (!state) {
+        return
+      }
+      if (err) {
+        this.pendingStreams.delete(queryId)
+        state.reject(err)
+      } else {
+        state.resolve([{ rows: [data] }, false])
+      }
+    })
+
+    this.mitt.on(endEvent, (queryId: string | undefined, _data: any, err: unknown) => {
+      if (!queryId) {
+        return
+      }
+      const state = this.pendingStreams.get(queryId)
+      if (!state) {
+        return
+      }
+      this.pendingStreams.delete(queryId)
+      if (err) {
+        state.reject(err)
+      } else {
+        state.resolve([undefined, true])
+      }
+    })
+  }
 
   async *streamQuery<R>(
-    { parameters, sql, query }: CompiledQuery,
+    { parameters, sql, query, queryId }: CompiledQuery,
     chunkSize?: number,
     options?: AbortableOperationOptions,
   ): AsyncIterableIterator<QueryResult<R>> {
@@ -80,40 +138,31 @@ class GenericSqliteWorkerConnection implements DatabaseConnection {
 
     this.worker.postMessage([
       dataEvent,
+      queryId.queryId,
       SelectQueryNode.is(query),
       sql,
       parameters,
       chunkSize,
     ] satisfies StreamMsg)
     type ResolveData = [data: QueryResult<R> | undefined, done: boolean]
     let done = false
-    let resolveFn: (value: ResolveData) => void
     let rejectFn: (reason?: any) => void
 
     const onAbort = (): void => rejectFn(new Error('Query aborted'))
     options?.signal?.addEventListener('abort', onAbort, { once: true })
 
-    this.mitt!.on(dataEvent, (data, err): void => {
-      if (err) {
-        rejectFn(err)
-      } else {
-        resolveFn([{ rows: [data] }, false])
-      }
-    })
-
-    this.mitt!.on(endEvent, (_, err): void => {
-      if (err) {
-        rejectFn(err)
-      } else {
-        resolveFn([undefined, true])
-      }
-    })
+    const streamState: { resolve: (data: any) => void; reject: (err: any) => void } = {
+      resolve: undefined!,
+      reject: undefined!,
+    }
+    this.pendingStreams.set(queryId.queryId, streamState)
 
     try {
       while (!done) {
         const [data, isDone] = await new Promise<ResolveData>((res, rej) => {
-          resolveFn = res
           rejectFn = rej
+          streamState.resolve = res
+          streamState.reject = rej
         })
 
         if (isDone) {
@@ -124,20 +173,22 @@ class GenericSqliteWorkerConnection implements DatabaseConnection {
       }
     } finally {
       options?.signal?.removeEventListener('abort', onAbort)
-      this.mitt?.off(dataEvent)
-      this.mitt?.off(endEvent)
+      this.pendingStreams.delete(queryId.queryId)
     }
   }
 
   async executeQuery<R>(compiledQuery: CompiledQuery<unknown>): Promise<QueryResult<R>> {
-    const { parameters, sql, query } = compiledQuery
-    this.worker.postMessage([runEvent, SelectQueryNode.is(query), sql, parameters] satisfies RunMsg)
-    return await new Promise((resolve, reject) => {
-      if (!this.mitt) {
-        reject(new Error('kysely instance has been destroyed'))
-      }
+    const { parameters, sql, query, queryId } = compiledQuery
+    this.worker.postMessage([
+      runEvent,
+      queryId.queryId,
+      SelectQueryNode.is(query),
+      sql,
+      parameters,
+    ] satisfies RunMsg)
 
-      this.mitt!.once(runEvent, (data, err) => (!err && data ? resolve(data) : reject(err)))
+    return new Promise((resolve, reject) => {
+      this.pendingRuns.set(queryId.queryId, { resolve, reject })
     })
   }
 }
diff --git a/packages/dialect-generic-sqlite/src/worker/types.ts b/packages/dialect-generic-sqlite/src/worker/types.ts
@@ -12,13 +12,15 @@ export type InitMsg<T extends Record<string, unknown>> = [type: typeof initEvent
 
 export type RunMsg = [
   type: typeof runEvent,
+  queryId: string,
   isSelect: boolean,
   sql: string,
   parameters?: readonly unknown[],
 ]
 
 export type StreamMsg = [
   type: typeof dataEvent,
+  queryId: string,
   isSelect: boolean,
   sql: string,
   parameters?: readonly unknown[],
@@ -34,7 +36,7 @@ export type MainToWorkerMsg<T extends Record<string, unknown>> =
   | StreamMsg
 
 export type WorkerToMainMsg = {
-  [K in keyof Events]: [type: `${K}`, data: Events[K], err: unknown]
+  [K in keyof Events]: [type: `${K}`, queryId: string | null, data: Events[K], err: unknown]
 }[keyof Events]
 
 type Events = {
diff --git a/packages/dialect-generic-sqlite/src/worker/utils.ts b/packages/dialect-generic-sqlite/src/worker/utils.ts
@@ -15,16 +15,17 @@ export function createGenericOnMessageCallback<T extends Record<string, unknown>
   message?: MessageHandleFn<DB>,
 ): (data: MainToWorkerMsg<T>) => Promise<void> {
   let db: IGenericSqlite<DB>
-  return async ([type, data1, data2, data3, data4]) => {
-    const ret: WorkerToMainMsg = [type, null, null]
+  return async ([type, data1, data2, data3, data4, data5]) => {
+    const ret: WorkerToMainMsg = [type, null, null, null]
     try {
       switch (type) {
         case initEvent:
           db = await init(data1)
           break
 
         case runEvent:
-          ret[1] = await db.query(data1, data2, data3)
+          ret[1] = data1 // queryId
+          ret[2] = await db.query(data2, data3, data4)
           break
 
         case closeEvent:
@@ -35,23 +36,24 @@ export function createGenericOnMessageCallback<T extends Record<string, unknown>
           if (!db.iterator) {
             throw new Error('streamQuery() is not supported.')
           }
-          const it = db.iterator(data1, data2, data3, data4)
+          const it = db.iterator(data2, data3, data4, data5)
           for await (const row of it) {
-            post([type, row as any, null] satisfies WorkerToMainMsg)
+            post([type, data1, row as any, null] satisfies WorkerToMainMsg)
           }
           ret[0] = endEvent
+          ret[1] = data1 // queryId
           break
         }
         default:
           if (message) {
             const data = await message(db, type, data1, data2, data3, data4)
             if (data !== undefined && data !== null) {
-              ret[1] = data
+              ret[2] = data
             }
           }
       }
     } catch (error) {
-      ret[2] = error
+      ret[3] = error
     }
     post(ret)
   }
diff --git a/packages/dialect-generic-sqlite/test/index.test.ts b/packages/dialect-generic-sqlite/test/index.test.ts
@@ -3,6 +3,8 @@ import { describe, expect, it } from 'vitest'
 
 import { testCase } from '../../test-utils'
 import { buildQueryFn, GenericSqliteDialect, parseBigInt } from '../src'
+import { createGenericOnMessageCallback } from '../src/worker'
+import { initEvent, dataEvent, runEvent } from '../src/worker/types'
 
 describe('generic sqlite dialect test', () => {
   it('better-sqlite3 executor', async () => {
@@ -31,4 +33,36 @@ describe('generic sqlite dialect test', () => {
 
     await testCase(dialect, expect, false)
   })
+  it('keeps worker query responses isolated by query id', async () => {
+    const messages: unknown[] = []
+    const onMessage = createGenericOnMessageCallback(
+      async () => ({
+        db: {},
+        close: () => {},
+        query: async (_isSelect, _sql, parameters) => ({
+          rows: [{ id: parameters?.[0] }],
+        }),
+        *iterator(_isSelect, _sql, parameters, chunkSize) {
+          yield {
+            chunkSize,
+            id: parameters?.[0],
+          }
+        },
+      }),
+      (message) => messages.push(message),
+    )
+
+    await onMessage([initEvent, {}])
+    await onMessage([runEvent, 'query-a', true, 'select ?', ['a']])
+    await onMessage([runEvent, 'query-b', true, 'select ?', ['b']])
+    await onMessage([dataEvent, 'stream-a', true, 'select ?', ['a'], 2])
+
+    expect(messages).toStrictEqual([
+      [initEvent, null, null, null],
+      [runEvent, 'query-a', { rows: [{ id: 'a' }] }, null],
+      [runEvent, 'query-b', { rows: [{ id: 'b' }] }, null],
+      [dataEvent, 'stream-a', { chunkSize: 2, id: 'a' }, null],
+      ['4', 'stream-a', null, null],
+    ])
+  })
 })
