feat(dialect-sqlite-worker): support stream

Author: subframe7536

packages/dialect-sqlite-worker/src/driver.ts

@@ -2,7 +2,7 @@ import { Worker } from 'node:worker_threads'
 import { join } from 'node:path'
 import { EventEmitter } from 'node:events'
 import type { DatabaseConnection, Driver, QueryResult } from 'kysely'
-import { CompiledQuery } from 'kysely'
+import { CompiledQuery, SelectQueryNode } from 'kysely'
 import type { MainMsg, SqliteWorkerDialectConfig, WorkerMsg } from './type'
 
 export class SqliteWorkerDriver implements Driver {
@@ -100,8 +100,48 @@ export class SqliteWorkerConnection implements DatabaseConnection {
     private emit?: EventEmitter,
   ) { }
 
-  async *streamQuery<R>(): AsyncIterableIterator<QueryResult<R>> {
-    throw new Error('sqlite-worker driver doesn\'t support streaming')
+  async *streamQuery<R>(compiledQuery: CompiledQuery): AsyncIterableIterator<QueryResult<R>> {
+    const { parameters, sql, query } = compiledQuery
+    if (!SelectQueryNode.is(query)) {
+      throw new Error('WaSqlite dialect only supported SELECT queries')
+    }
+    this.worker.postMessage(['2', sql, parameters] satisfies MainMsg)
+    let resolver: ((value: IteratorResult<{ rows: QueryResult<R>[] }>) => void) | null = null
+    let rejecter: ((reason: any) => void) | null = null
+
+    this.emit!.on('2', (data, err) => {
+      if (err && rejecter) {
+        rejecter(err)
+      }
+      if (resolver) {
+        resolver({ value: { rows: data! }, done: false })
+        resolver = null
+      }
+    })
+
+    this.emit!.on('3', (_, err) => {
+      if (err && rejecter) {
+        rejecter(err)
+      }
+      if (resolver) {
+        resolver({ value: undefined, done: true })
+      }
+    })
+
+    return {
+      [Symbol.asyncIterator]() {
+        return this
+      },
+      async next() {
+        return new Promise<IteratorResult<any>>((resolve, reject) => {
+          resolver = resolve
+          rejecter = reject
+        })
+      },
+      async return() {
+        return { value: undefined, done: true }
+      },
+    }
   }
 
   async executeQuery<R>(compiledQuery: CompiledQuery<unknown>): Promise<QueryResult<R>> {

packages/dialect-sqlite-worker/src/type.ts

@@ -27,8 +27,9 @@ type RunMsg = [
 ]
 
 type CloseMsg = [type: '1']
+type StreamMsg = [type: '2', sql: string, parameters?: readonly unknown[]]
 
-export type MainMsg = RunMsg | CloseMsg
+export type MainMsg = RunMsg | CloseMsg | StreamMsg
 
 export type WorkerMsg =
   | [
@@ -41,3 +42,13 @@ export type WorkerMsg =
     data: null,
     err: unknown,
   ]
+  | [
+    type: '2',
+    data: QueryResult<any>[] | null,
+    err: unknown,
+  ]
+  | [
+    type: '3',
+    data: null,
+    err: unknown,
+  ]

packages/dialect-sqlite-worker/src/worker.ts

@@ -9,30 +9,42 @@ if (!parentPort) {
 const { src, option } = workerData
 const db = new Database(src, option)
 
-parentPort.on('message', (msg: MainMsg) => {
+parentPort.on('message', ([msg, data1, data2]: MainMsg) => {
   const ret: WorkerMsg = [
-    msg[0],
+    msg,
     null,
     null,
   ]
 
   try {
-    if (msg[0] === '1') {
-      db.close()
-      parentPort!.postMessage(ret)
-      return
-    }
-    const stmt = db.prepare(msg[1])
-    if (stmt.reader) {
-      ret[1] = {
-        rows: stmt.all(msg[2]),
+    switch (msg) {
+      case '0': {
+        const stmt = db.prepare(data1)
+        if (stmt.reader) {
+          ret[1] = {
+            rows: stmt.all(data2),
+          }
+        } else {
+          const { changes, lastInsertRowid } = stmt.run(data2)
+          ret[1] = {
+            rows: [],
+            numAffectedRows: BigInt(changes),
+            insertId: BigInt(lastInsertRowid),
+          }
+        }
+        break
       }
-    } else {
-      const { changes, lastInsertRowid } = stmt.run(msg[2])
-      ret[1] = {
-        rows: [],
-        numAffectedRows: BigInt(changes),
-        insertId: BigInt(lastInsertRowid),
+      case '1':
+        db.close()
+        break
+      case '2': {
+        const stmt = db.prepare(data1)
+        const iter = stmt.iterate(data2)
+        for (const row of iter) {
+          parentPort!.postMessage([msg, [row as any], null] satisfies WorkerMsg)
+        }
+        ret[0] = '3'
+        break
       }
     }
   } catch (error) {

packages/dialect-wasqlite-worker/src/driver.ts

@@ -37,7 +37,7 @@ export class WaSqliteWorkerDriver implements Driver {
       0,
       this.config.fileName,
       // if use OPFS, wasm should use sync version
-      parseWorkerOrURL(this.config.url ?? defaultWasmURL, !useOPFS),
+      parseWorkerOrURL(this.config.url ?? defaultWasmURL, !useOPFS) as string,
       useOPFS,
     ] satisfies MainMsg)
     await new Promise<void>((resolve, reject) => {
@@ -123,12 +123,12 @@ class WaSqliteWorkerConnection implements DatabaseConnection {
     this.mitt = mitt
   }
 
-  streamQuery<R>(compiledQuery: CompiledQuery, chunkSize = 1): AsyncIterableIterator<QueryResult<R>> {
+  streamQuery<R>(compiledQuery: CompiledQuery): AsyncIterableIterator<QueryResult<R>> {
     const { parameters, sql, query } = compiledQuery
     if (!SelectQueryNode.is(query)) {
       throw new Error('WaSqlite dialect only supported SELECT queries')
     }
-    this.worker.postMessage([3, chunkSize, sql, parameters] satisfies MainMsg)
+    this.worker.postMessage([3, sql, parameters] satisfies MainMsg)
     let resolver: ((value: IteratorResult<{ rows: QueryResult<R>[] }>) => void) | null = null
     let rejecter: ((reason: any) => void) | null = null
 

packages/dialect-wasqlite-worker/src/type.ts

@@ -55,7 +55,6 @@ type RunMsg = [
 ]
 type StreamMsg = [
   type: 3,
-  chunkSize: number,
   sql: string,
   parameters?: readonly unknown[],
 ]

packages/dialect-wasqlite-worker/src/worker.ts

@@ -30,13 +30,8 @@ async function exec(isSelect: boolean, sql: string, parameters?: readonly unknow
 async function stream(onData: (data: any) => void, sql: string, parameters?: readonly unknown[]): Promise<void> {
   await db.stream(onData, sql, parameters as any[])
 }
-onmessage = async ({ data }: MessageEvent<MainMsg>) => {
-  const [msg, data1, data2, data3] = data
-  const ret: WorkerMsg = [
-    msg,
-    null,
-    null,
-  ]
+onmessage = async ({ data: [msg, data1, data2, data3] }: MessageEvent<MainMsg>) => {
+  const ret: WorkerMsg = [msg, null, null]
   try {
     switch (msg) {
       case 0:
@@ -48,19 +43,10 @@ onmessage = async ({ data }: MessageEvent<MainMsg>) => {
       case 2:
         await db.close()
         break
-      case 3: {
-        let result: any[] = []
-        await stream((val) => {
-          if (result.length < data1) {
-            result.push(val)
-          } else {
-            postMessage([3, result, null] satisfies WorkerMsg)
-            result = []
-          }
-        }, data2, data3)
+      case 3:
+        await stream(val => postMessage([3, [val], null] satisfies WorkerMsg), data1, data2)
         ret[0] = 4
         break
-      }
     }
   } catch (error) {
     ret[2] = error

test/dialect-sqlite-worker.test.ts

@@ -33,6 +33,12 @@ describe('sqlite worker dialect test', () => {
     expect(age).toStrictEqual(18)
     expect(name).toStrictEqual(`test ${dialect.toString()}`)
     expect(int8).toStrictEqual(Uint8Array.from([1, 2, 3]))
+    const rows = db.selectFrom('test').selectAll().stream()
+    for await (const row of rows) {
+      expect(row.age).toStrictEqual(18)
+      expect(row.name).toStrictEqual(`test ${dialect.toString()}`)
+      expect(row.int8).toStrictEqual(Uint8Array.from([1, 2, 3]))
+    }
   }
   afterAll(async () => {
     await db.destroy()