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

Author: subframe7536

packages/dialect-sqlite-worker/README.md

@@ -10,6 +10,30 @@ the type is also availiable for [`better-sqlite3-multiple-ciphers`](https://gith
 pnpm add kysely kysely-sqlite-worker better-sqlite3
 ```
 
+## Usage
+
+```ts
+import { SqliteWorkerDialect } from 'kysely-sqlite-worker'
+
+const dialect = new SqliteWorkerDialect({
+  source: ':memory:',
+})
+```
+
+### Custom Worker
+
+in `worker.ts`
+
+```ts
+import { createOnMessageCallback } from 'kysely-sqlite-worker'
+
+createOnMessageCallback(
+  async (db) => {
+    db.loadExtension(/* ... */)
+  }
+)
+```
+
 ## Config
 
 ```ts

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

@@ -3,6 +3,8 @@ import type { SqliteWorkerDialectConfig } from './type'
 import { SqliteAdapter, SqliteIntrospector, SqliteQueryCompiler } from 'kysely'
 import { SqliteWorkerDriver } from './driver'
 
+export { createOnMessageCallback } from './worker/util'
+
 export class SqliteWorkerDialect implements Dialect {
   /**
    * dialect for better-sqlite, execute sql in `node:worker_threads`

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

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

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

@@ -0,0 +1,68 @@
+import type { MainMsg, WorkerMsg } from '../type'
+import { parentPort, workerData } from 'node:worker_threads'
+import Database from 'better-sqlite3'
+
+/**
+ * Handle worker message, support custom callback on initialization
+ * @example
+ * // worker.ts
+ * import { createOnMessageCallback } from 'kysely-sqlite-worker'
+ *
+ * createOnMessageCallback(
+ *   async (db) => {
+ *     db.loadExtension(...)
+ *   }
+ * )
+ */
+export function createOnMessageCallback(onInit?: (db: typeof Database) => void): void {
+  if (!parentPort) {
+    throw new Error('Must be run in a worker thread')
+  }
+
+  const { src, option } = workerData
+  const db = new Database(src, option)
+  onInit?.(db as any)
+  parentPort.on('message', ([msg, data1, data2]: MainMsg) => {
+    const ret: WorkerMsg = [
+      msg,
+      null,
+      null,
+    ]
+
+    try {
+      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
+        }
+        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) {
+      ret[2] = error
+    }
+    parentPort!.postMessage(ret)
+  })
+}

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

@@ -2,10 +2,10 @@ import { defineConfig } from 'tsup'
 
 export default defineConfig([
   {
-    entry: [
-      'src/index.ts',
-      'src/worker.ts',
-    ],
+    entry: {
+      index: 'src/index.ts',
+      worker: 'src/worker.ts',
+    },
     clean: true,
     format: ['cjs', 'esm'],
     shims: true,