feat(dialect-wasqlite-worker): function to create custom OnMessageCallback in worker

Author: subframe7536

packages/dialect-wasqlite-worker/README.md

@@ -10,6 +10,38 @@ No need to set response header like [official wasm](../dialect-wasm/README.md#of
 pnpm add kysely kysely-wasqlite-worker
 ```
 
+## Usage
+
+```ts
+import {
+  generateDialectOptions,
+  isIdbSupported,
+  isModuleWorkerSupport,
+  isOpfsSupported,
+  useDefaultWasmURL,
+  useDefaultWorker,
+  WaSqliteWorkerDialect
+} from 'kysely-wasqlite-worker'
+
+const dialect = new WaSqliteWorkerDialect({
+  fileName: 'test',
+})
+```
+
+### Custom Worker
+
+in `worker.ts`
+
+```ts
+import { createOnMessageCallback, customFunction } from 'kysely-wasqlite-worker'
+
+onmessage = createOnMessageCallback(
+  async (sqliteDB: SQLiteDB) => {
+    customFunction(sqliteDB.sqlite, sqliteDB.db, 'customFunction', (a, b) => a + b)
+  }
+)
+```
+
 ## Config
 
 ```ts
@@ -59,24 +91,6 @@ export interface WaSqliteWorkerDialectConfig {
 }
 ```
 
-## Usage
-
-```ts
-import {
-  generateDialectOptions,
-  isIdbSupported,
-  isModuleWorkerSupport,
-  isOpfsSupported,
-  useDefaultWasmURL,
-  useDefaultWorker,
-  WaSqliteWorkerDialect
-} from 'kysely-wasqlite-worker'
-
-const dialect = new WaSqliteWorkerDialect({
-  fileName: 'test',
-})
-```
-
 see more in [playground](../../playground/src/modules/wasqliteWorker.ts)
 
 if throw error when using `Vite` to build, add worker config

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

@@ -1,6 +1,6 @@
 import type { DatabaseConnection, Driver, QueryResult } from 'kysely'
 import type { Emitter } from 'zen-mitt'
-import type { EventWithError, MainMsg, WaSqliteWorkerDialectConfig, WorkerMsg } from './type'
+import type { EventWithError, MainToWorkerMsg, WaSqliteWorkerDialectConfig, WorkerToMainMsg } from './type'
 import { isModuleWorkerSupport, isOpfsSupported } from '@subframe7536/sqlite-wasm'
 import { CompiledQuery, SelectQueryNode } from 'kysely'
 import { mitt } from 'zen-mitt'
@@ -29,7 +29,7 @@ export class WaSqliteWorkerDriver implements Driver {
 
     this.worker = parseWorkerOrURL(this.config.worker || defaultWorker, useOPFS || isModuleWorkerSupport())
 
-    this.worker.onmessage = ({ data: [type, ...msg] }: MessageEvent<WorkerMsg>) => {
+    this.worker.onmessage = ({ data: [type, ...msg] }: MessageEvent<WorkerToMainMsg>) => {
       this.mitt?.emit(type, ...msg)
     }
     this.worker.postMessage([
@@ -38,7 +38,7 @@ export class WaSqliteWorkerDriver implements Driver {
       // if use OPFS, wasm should use sync version
       parseWorkerOrURL(this.config.url ?? defaultWasmURL, !useOPFS) as string,
       useOPFS,
-    ] satisfies MainMsg)
+    ] satisfies MainToWorkerMsg)
     await new Promise<void>((resolve, reject) => {
       this.mitt?.once(0, (_, err) => err ? reject(err) : resolve())
     })
@@ -74,7 +74,7 @@ export class WaSqliteWorkerDriver implements Driver {
     if (!this.worker) {
       return
     }
-    this.worker.postMessage([2] satisfies MainMsg)
+    this.worker.postMessage([2] satisfies MainToWorkerMsg)
     return new Promise<void>((resolve, reject) => {
       this.mitt?.once(2, (_, err) => {
         if (err) {
@@ -127,7 +127,7 @@ class WaSqliteWorkerConnection implements DatabaseConnection {
     if (!SelectQueryNode.is(query)) {
       throw new Error('WaSqlite dialect only supported SELECT queries')
     }
-    this.worker.postMessage([3, sql, parameters] satisfies MainMsg)
+    this.worker.postMessage([3, sql, parameters] satisfies MainToWorkerMsg)
     let resolver: ((value: IteratorResult<{ rows: QueryResult<R>[] }>) => void) | null = null
     let rejecter: ((reason: any) => void) | null = null
 
@@ -169,7 +169,7 @@ class WaSqliteWorkerConnection implements DatabaseConnection {
   async executeQuery<R>(compiledQuery: CompiledQuery<unknown>): Promise<QueryResult<R>> {
     const { parameters, sql, query } = compiledQuery
     const isSelect = SelectQueryNode.is(query)
-    this.worker.postMessage([1, isSelect, sql, parameters] satisfies MainMsg)
+    this.worker.postMessage([1, isSelect, sql, parameters] satisfies MainToWorkerMsg)
     return new Promise((resolve, reject) => {
       if (!this.mitt) {
         reject(new Error('kysely instance has been destroyed'))

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

@@ -3,7 +3,15 @@ import type { WaSqliteWorkerDialectConfig } from './type'
 import { SqliteAdapter, SqliteIntrospector, SqliteQueryCompiler } from 'kysely'
 import { WaSqliteWorkerDriver } from './driver'
 
-export { isIdbSupported, isModuleWorkerSupport, isOpfsSupported } from '@subframe7536/sqlite-wasm'
+export { createOnMessageCallback } from './worker/utils'
+
+export {
+  customFunction,
+  isIdbSupported,
+  isModuleWorkerSupport,
+  isOpfsSupported,
+  type SQLiteDB,
+} from '@subframe7536/sqlite-wasm'
 
 export class WaSqliteWorkerDialect implements Dialect {
   /**

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

@@ -67,9 +67,9 @@ type InitMsg = [
 ]
 
 type CloseMsg = [2]
-export type MainMsg = InitMsg | RunMsg | CloseMsg | StreamMsg
+export type MainToWorkerMsg = InitMsg | RunMsg | CloseMsg | StreamMsg
 
-export type WorkerMsg = {
+export type WorkerToMainMsg = {
   [K in keyof Events]: [
     type: K,
     data: Events[K],

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

@@ -0,0 +1,3 @@
+import { createOnMessageCallback } from './utils'
+
+onmessage = createOnMessageCallback()

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

@@ -0,0 +1,80 @@
+import type { SQLiteDB } from '@subframe7536/sqlite-wasm'
+import type { QueryResult } from 'kysely'
+import type { MainToWorkerMsg, WorkerToMainMsg } from '../type'
+import { initSQLite } from '@subframe7536/sqlite-wasm'
+
+let db: SQLiteDB
+
+async function init(
+  fileName: string,
+  url: string,
+  useOPFS: boolean,
+  afterInit?: (sqliteDB: SQLiteDB) => Promise<void>,
+): Promise<void> {
+  db = await initSQLite(
+    (
+      useOPFS
+        ? (await import('@subframe7536/sqlite-wasm/opfs')).useOpfsStorage
+        : (await import('@subframe7536/sqlite-wasm/idb')).useIdbStorage
+    )(
+      fileName,
+      { url },
+    ),
+  )
+  await afterInit?.(db)
+}
+async function exec(isSelect: boolean, sql: string, parameters?: readonly unknown[]): Promise<QueryResult<any>> {
+  const rows = await db.run(sql, parameters as any[])
+  return isSelect || rows.length
+    ? { rows }
+    : {
+        rows,
+        insertId: BigInt(db.lastInsertRowId()),
+        numAffectedRows: BigInt(db.changes()),
+      }
+}
+async function stream(onData: (data: any) => void, sql: string, parameters?: readonly unknown[]): Promise<void> {
+  await db.stream(onData, sql, parameters as any[])
+}
+
+/**
+ * Handle worker message, support custom callback on initialization
+ * @example
+ * // worker.ts
+ * import { createOnMessageCallback, customFunction } from 'kysely-wasqlite-worker'
+ *
+ * onmessage = createOnMessageCallback(
+ *   async (sqliteDB: SQLiteDB) => {
+ *     customFunction(sqliteDB.sqlite, sqliteDB.db, 'customFunction', (a, b) => a + b)
+ *   }
+ * )
+ */
+export function createOnMessageCallback(
+  afterInit?: (sqliteDB: SQLiteDB) => Promise<void>,
+): (event: MessageEvent<MainToWorkerMsg>) => Promise<void> {
+  return async ({
+    data: [msg, data1, data2, data3],
+  }: MessageEvent<MainToWorkerMsg>) => {
+    const ret: WorkerToMainMsg = [msg, null, null]
+    try {
+      switch (msg) {
+        case 0:
+          await init(data1, data2, data3, afterInit)
+          break
+        case 1:
+          ret[1] = await exec(data1, data2, data3)
+          break
+        case 2:
+          await db.close()
+          break
+        case 3:
+          await stream(val => postMessage([3, [val], null] satisfies WorkerToMainMsg), data1, data2)
+          ret[0] = 4
+          break
+      }
+    } catch (error) {
+      ret[2] = error
+    }
+    postMessage(ret)
+  }
+}

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

@@ -1,52 +1,42 @@
 import { copyFile } from 'node:fs/promises'
 import { defineConfig } from 'tsup'
 
-export default defineConfig([
-  {
-    entry: [
-      'src/index.ts',
-    ],
-    clean: true,
-    format: ['cjs', 'esm'],
-    shims: true,
-    dts: true,
-    external: ['kysely'],
-    treeshake: true,
+export default defineConfig({
+  entry: {
+    index: 'src/index.ts',
+    worker: 'src/worker/index.ts',
   },
-  {
-    entry: [
-      'src/worker.ts',
-    ],
-    clean: false,
-    format: ['cjs', 'esm'],
-    dts: true,
-    treeshake: true,
-    plugins: [
-      {
-        name: 'classic worker',
-        renderChunk(code) {
-          if (this.format === 'cjs') {
-            return { code: code.replaceAll('import.meta.url', 'self.location.href') }
-          }
-        },
+  clean: true,
+  format: ['cjs', 'esm'],
+  shims: true,
+  dts: true,
+  external: ['kysely'],
+  treeshake: true,
+  plugins: [
+    {
+      name: 'classic worker',
+      renderChunk(code) {
+        if (this.format === 'cjs') {
+          return { code: code.replaceAll('import.meta.url', 'self.location.href') }
+        }
       },
-      {
-        name: 'copy',
-        buildEnd(this) {
-          if (this.format === 'esm') {
-            Promise.all([
-              copyFile(
-                './node_modules/@subframe7536/sqlite-wasm/dist/wa-sqlite.wasm',
-                './dist/wa-sqlite.wasm',
-              ),
-              copyFile(
-                './node_modules/@subframe7536/sqlite-wasm/dist/wa-sqlite-async.wasm',
-                './dist/wa-sqlite-async.wasm',
-              ),
-            ])
-          }
-        },
+    },
+    {
+      name: 'copy',
+      buildEnd(this) {
+        if (this.format === 'esm') {
+          Promise.all([
+            copyFile(
+              './node_modules/@subframe7536/sqlite-wasm/dist/wa-sqlite.wasm',
+              './dist/wa-sqlite.wasm',
+            ),
+            copyFile(
+              './node_modules/@subframe7536/sqlite-wasm/dist/wa-sqlite-async.wasm',
+              './dist/wa-sqlite-async.wasm',
+            ),
+          ])
+        }
       },
-    ],
-  },
-])
+    },
+  ],
+})