feat: custom function

Author: subframe7536

playground/src/worker.ts

@@ -1,18 +1,22 @@
 import url from 'wa-sqlite/dist/wa-sqlite.wasm?url'
-import { initSQLite, isOpfsSupported } from '../../src'
+import { def, initSQLite, isOpfsSupported } from '../../src'
 import { useOpfsStorage } from '../../src/vfs/opfs'
 import { runSQLStream } from './runSQL'
 
 onmessage = async () => {
  if (!await isOpfsSupported()) {
  return
  }
- const { run, stream, lastInsertRowId, changes } = await initSQLite(useOpfsStorage(
+ const { run, stream, lastInsertRowId, changes, sqlite, db } = await initSQLite(useOpfsStorage(
  'test',
  { url },
  // 'https://cdn.jsdelivr.net/gh/rhashimoto/wa-sqlite@v0.9.9/dist/wa-sqlite.wasm',
  ))
  await runSQLStream(run, stream, data => postMessage(data))
  console.log(lastInsertRowId(), changes())
+ def(sqlite, db, 'testtest', (a: number, b: number) => a + b)
+ console.log(
+ await run('select testtest(1,2)'),
+ )
  postMessage('done')
 }

src/utils.ts

@@ -1,3 +1,5 @@
+import { SQLITE_DETERMINISTIC, SQLITE_DIRECTONLY, SQLITE_UTF8 } from 'wa-sqlite'
+
 /**
  * check if IndexedDB and Web Locks API supported
  */
@@ -82,4 +84,38 @@ export function isModuleWorkerSupport(): boolean {
  }
 }
 
+export function def<N extends string, T extends SQLiteCompatibleType[]>(
+ sqlite: SQLiteAPI,
+ db: number,
+ fnName: N,
+ fn: N extends '' ? never : (...args: T) => (SQLiteCompatibleType | number[]) | null,
+ option: {
+ deterministic?: boolean
+ directOnly?: boolean
+ varargs?: boolean
+ } = {},
+): void {
+ let flags = SQLITE_UTF8
+ if (option.deterministic) {
+ flags |= SQLITE_DETERMINISTIC
+ }
+ if (option.directOnly) {
+ flags |= SQLITE_DIRECTONLY
+ }
+ sqlite.create_function(
+ db,
+ fnName,
+ (option.varargs || fn.length === 0) ? -1 : fn.length,
+ flags,
+ 0,
+ (ctx, value) => {
+ const args = [] as unknown as T
+ for (let i = 0; i < fn.length; i++) {
+ args.push(sqlite.value(value[i]))
+ }
+ return sqlite.result(ctx, fn(...args))
+ },
+ )
+}
+
 // todo: import/export