Create a sqlite database wrapper as a route to supporting other sqlite drivers

Author: samwillis

packages/d2ts/src/sqlite/database.ts

@@ -0,0 +1,77 @@
+/**
+ * Represents a prepared statement in SQLite
+ */
+export interface SQLiteStatement<Params = unknown, Result = unknown> {
+  /**
+   * Run the prepared statement with parameters
+   * Accepts either an array of parameters or an object of named parameters
+   */
+  run(...params: Params extends any[] ? Params : [Params?]): void
+
+  /**
+   * Get a single row from the prepared statement
+   * Accepts either an array of parameters or an object of named parameters
+   */
+  get(...params: Params extends any[] ? Params : [Params?]): Result | undefined
+
+  /**
+   * Get all rows from the prepared statement
+   * Accepts either an array of parameters or an object of named parameters
+   */
+  all(...params: Params extends any[] ? Params : [Params?]): Result[]
+}
+
+/**
+ * Interface for SQLite database operations
+ */
+export interface SQLiteDb {
+  /**
+   * Execute raw SQL
+   */
+  exec(sql: string): void
+
+  /**
+   * Prepare a statement
+   */
+  prepare<Params = unknown, Result = unknown>(
+    sql: string,
+  ): SQLiteStatement<Params, Result>
+
+  /**
+   * Close the database connection
+   */
+  close(): void
+}
+
+/**
+ * Wrapper for better-sqlite3 to implement SQLiteDb interface
+ */
+export class BetterSQLite3Wrapper implements SQLiteDb {
+  #db: import('better-sqlite3').Database
+
+  constructor(db: import('better-sqlite3').Database) {
+    this.#db = db
+  }
+
+  exec(sql: string): void {
+    this.#db.exec(sql)
+  }
+
+  prepare<Params = unknown, Result = unknown>(
+    sql: string,
+  ): SQLiteStatement<Params, Result> {
+    const stmt = this.#db.prepare(sql)
+    return {
+      run: (...params: Params extends any[] ? Params : [Params?]) => 
+        stmt.run(...params),
+      get: (...params: Params extends any[] ? Params : [Params?]) => 
+        stmt.get(...params) as Result | undefined,
+      all: (...params: Params extends any[] ? Params : [Params?]) => 
+        stmt.all(...params) as Result[],
+    }
+  }
+
+  close(): void {
+    this.#db.close()
+  }
+}

packages/d2ts/src/sqlite/index.ts

@@ -1,2 +1,3 @@
 export * from './version-index.js'
 export * from './operators/index.js'
+export * from './database.js'

packages/d2ts/src/sqlite/operators/consolidate.ts

@@ -12,7 +12,7 @@ import {
   UnaryOperator,
 } from '../../graph.js'
 import { Version, Antichain } from '../../order.js'
-import Database from 'better-sqlite3'
+import { SQLiteDb, SQLiteStatement } from '../database.js'
 
 interface CollectionRow {
   version: string
@@ -29,19 +29,19 @@ interface CollectionParams {
  */
 export class ConsolidateOperatorSQLite<T> extends UnaryOperator<T> {
   #preparedStatements: {
-    insert: Database.Statement<CollectionParams>
-    update: Database.Statement<CollectionParams>
-    get: Database.Statement<string, CollectionRow>
-    delete: Database.Statement<string>
-    getAllVersions: Database.Statement<[], CollectionRow>
+    insert: SQLiteStatement<CollectionParams>
+    update: SQLiteStatement<CollectionParams>
+    get: SQLiteStatement<[string], CollectionRow>
+    delete: SQLiteStatement<[string]>
+    getAllVersions: SQLiteStatement<[], CollectionRow>
   }
 
   constructor(
     id: number,
     inputA: DifferenceStreamReader<T>,
     output: DifferenceStreamWriter<T>,
     initialFrontier: Antichain,
-    db: Database.Database,
+    db: SQLiteDb,
   ) {
     super(id, inputA, output, initialFrontier)
 
@@ -143,7 +143,7 @@ export class ConsolidateOperatorSQLite<T> extends UnaryOperator<T> {
  * Persists state to SQLite
  * @param db - The SQLite database
  */
-export function consolidate<T>(db: Database.Database): PipedOperator<T, T> {
+export function consolidate<T>(db: SQLiteDb): PipedOperator<T, T> {
   return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
     const output = new StreamBuilder<T>(
       stream.graph,

packages/d2ts/src/sqlite/operators/count.ts

@@ -2,8 +2,7 @@ import { StreamBuilder } from '../../d2.js'
 import { IStreamBuilder, KeyValue } from '../../types.js'
 import { DifferenceStreamReader, DifferenceStreamWriter } from '../../graph.js'
 import { Antichain } from '../../order.js'
-
-import Database from 'better-sqlite3'
+import { SQLiteDb } from '../database.js'
 import { ReduceOperatorSQLite } from './reduce.js'
 
 export class CountOperatorSQLite<K, V> extends ReduceOperatorSQLite<
@@ -16,7 +15,7 @@ export class CountOperatorSQLite<K, V> extends ReduceOperatorSQLite<
     inputA: DifferenceStreamReader<[K, V]>,
     output: DifferenceStreamWriter<[K, number]>,
     initialFrontier: Antichain,
-    db: Database.Database,
+    db: SQLiteDb,
   ) {
     const countInner = (vals: [V, number][]): [number, number][] => {
       let count = 0
@@ -39,7 +38,7 @@ export function count<
   K extends T extends KeyValue<infer K, infer _V> ? K : never,
   V extends T extends KeyValue<K, infer V> ? V : never,
   T,
->(db: Database.Database) {
+>(db: SQLiteDb) {
   return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<K, number>> => {
     const output = new StreamBuilder<KeyValue<K, number>>(
       stream.graph,

packages/d2ts/src/sqlite/operators/distinct.ts

@@ -2,7 +2,7 @@ import { StreamBuilder } from '../../d2.js'
 import { IStreamBuilder, KeyValue } from '../../types.js'
 import { DifferenceStreamReader, DifferenceStreamWriter } from '../../graph.js'
 import { Antichain } from '../../order.js'
-import Database from 'better-sqlite3'
+import { SQLiteDb } from '../database.js'
 import { ReduceOperatorSQLite } from './reduce.js'
 
 export class DistinctOperatorSQLite<K, V> extends ReduceOperatorSQLite<
@@ -15,7 +15,7 @@ export class DistinctOperatorSQLite<K, V> extends ReduceOperatorSQLite<
     inputA: DifferenceStreamReader<[K, V]>,
     output: DifferenceStreamWriter<[K, V]>,
     initialFrontier: Antichain,
-    db: Database.Database,
+    db: SQLiteDb,
   ) {
     const distinctInner = (vals: [V, number][]): [V, number][] => {
       const consolidated = new Map<string, number>()
@@ -43,7 +43,7 @@ export function distinct<
   K extends T extends KeyValue<infer K, infer _V> ? K : never,
   V extends T extends KeyValue<K, infer V> ? V : never,
   T,
->(db: Database.Database) {
+>(db: SQLiteDb) {
   return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<K, V>> => {
     const output = new StreamBuilder<KeyValue<K, V>>(
       stream.graph,

packages/d2ts/src/sqlite/operators/join.ts

@@ -12,23 +12,23 @@ import {
   BinaryOperator,
 } from '../../graph.js'
 import { Version, Antichain } from '../../order.js'
-import Database from 'better-sqlite3'
+import { SQLiteDb } from '../database.js'
 import { SQLIndex } from '../version-index.js'
 
 export class JoinOperatorSQLite<K, V1, V2> extends BinaryOperator<
   [K, unknown]
 > {
   #indexA: SQLIndex<K, V1>
   #indexB: SQLIndex<K, V2>
-  #db: Database.Database
+  #db: SQLiteDb
 
   constructor(
     id: number,
     inputA: DifferenceStreamReader<[K, V1]>,
     inputB: DifferenceStreamReader<[K, V2]>,
     output: DifferenceStreamWriter<[K, [V1, V2]]>,
     initialFrontier: Antichain,
-    db: Database.Database,
+    db: SQLiteDb,
   ) {
     super(id, inputA, inputB, output, initialFrontier)
     this.#db = db
@@ -117,7 +117,7 @@ export class JoinOperatorSQLite<K, V1, V2> extends BinaryOperator<
         this.output.sendFrontier(this.outputFrontier)
         this.#indexA.compact(this.outputFrontier)
         this.#indexB.compact(this.outputFrontier)
-      }
+    }
     } finally {
       // Clean up temporary indexes
       deltaA.destroy()
@@ -137,7 +137,7 @@ export function join<
   V1 extends T extends KeyValue<infer _KT, infer VT> ? VT : never,
   V2,
   T,
->(other: IStreamBuilder<KeyValue<K, V2>>, db: Database.Database) {
+>(other: IStreamBuilder<KeyValue<K, V2>>, db: SQLiteDb) {
   return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<K, [V1, V2]>> => {
     if (stream.graph !== other.graph) {
       throw new Error('Cannot join streams from different graphs')

packages/d2ts/src/sqlite/operators/reduce.ts

@@ -12,7 +12,7 @@ import {
   UnaryOperator,
 } from '../../graph.js'
 import { Version, Antichain } from '../../order.js'
-import Database from 'better-sqlite3'
+import { SQLiteDb, SQLiteStatement } from '../database.js'
 import { SQLIndex } from '../version-index.js'
 
 interface KeysTodoRow {
@@ -29,11 +29,11 @@ export class ReduceOperatorSQLite<K, V1, V2> extends UnaryOperator<
   #index: SQLIndex<K, V1>
   #indexOut: SQLIndex<K, V2>
   #preparedStatements: {
-    insertKeyTodo: Database.Statement<[string, string]>
-    getKeysTodo: Database.Statement<[], KeysTodoRow>
-    deleteKeysTodo: Database.Statement<[string]>
-    createKeysTodoTable: Database.Statement
-    dropKeysTodoTable: Database.Statement
+    insertKeyTodo: SQLiteStatement<[string, string]>
+    getKeysTodo: SQLiteStatement<[], KeysTodoRow>
+    deleteKeysTodo: SQLiteStatement<[string]>
+    createKeysTodoTable: SQLiteStatement
+    dropKeysTodoTable: SQLiteStatement
   }
   #f: (values: [V1, number][]) => [V2, number][]
 
@@ -43,7 +43,7 @@ export class ReduceOperatorSQLite<K, V1, V2> extends UnaryOperator<
     output: DifferenceStreamWriter<[K, V2]>,
     f: (values: [V1, number][]) => [V2, number][],
     initialFrontier: Antichain,
-    db: Database.Database,
+    db: SQLiteDb,
   ) {
     super(id, inputA, output, initialFrontier)
     this.#f = f
@@ -211,7 +211,7 @@ export function reduce<
   V1 extends T extends KeyValue<K, infer V> ? V : never,
   R,
   T,
->(f: (values: [V1, number][]) => [R, number][], db: Database.Database) {
+>(f: (values: [V1, number][]) => [R, number][], db: SQLiteDb) {
   return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<K, R>> => {
     const output = new StreamBuilder<KeyValue<K, R>>(
       stream.graph,

packages/d2ts/src/sqlite/version-index.ts

@@ -1,6 +1,6 @@
 import { Version, Antichain, v } from '../order.js'
 import { MultiSet } from '../multiset.js'
-import Database from 'better-sqlite3'
+import { SQLiteDb, SQLiteStatement } from './database.js'
 
 interface IndexRow {
   key: string
@@ -33,37 +33,37 @@ interface JoinResult {
 }
 
 export class SQLIndex<K, V> {
-  #db: Database.Database
+  #db: SQLiteDb
   #tableName: string
   #versionTableName: string
   #isTemp: boolean
   #preparedStatements: {
-    insert: Database.Statement<InsertParams>
-    get: Database.Statement<GetParams, IndexRow>
-    getVersions: Database.Statement<[string], { version: string }>
-    getAllForKey: Database.Statement<[string, GetAllForKeyParams], IndexRow>
-    delete: Database.Statement<[string]>
-    deleteAll: Database.Statement
-    deleteAllVersions: Database.Statement
-    getForCompaction: Database.Statement<[string], IndexRow>
-    consolidateVersions: Database.Statement<[string, string]>
-    insertVersion: Database.Statement<[string], { id: number }>
-    updateVersionMapping: Database.Statement<[string, string]>
-    deleteZeroMultiplicity: Database.Statement
-    getVersionId: Database.Statement<[string], { id: number }>
-    setCompactionFrontier: Database.Statement<[string]>
-    getCompactionFrontier: Database.Statement<[], { value: string }>
-    deleteMeta: Database.Statement
-    getAllKeys: Database.Statement<[], { key: string }>
-    getVersionsForKey: Database.Statement<
+    insert: SQLiteStatement<InsertParams>
+    get: SQLiteStatement<GetParams, IndexRow>
+    getVersions: SQLiteStatement<[string], { version: string }>
+    getAllForKey: SQLiteStatement<[string, GetAllForKeyParams], IndexRow>
+    delete: SQLiteStatement<[string]>
+    deleteAll: SQLiteStatement
+    deleteAllVersions: SQLiteStatement
+    getForCompaction: SQLiteStatement<[string], IndexRow>
+    consolidateVersions: SQLiteStatement<[string, string]>
+    insertVersion: SQLiteStatement<[string], { id: number }>
+    updateVersionMapping: SQLiteStatement<[string, string]>
+    deleteZeroMultiplicity: SQLiteStatement
+    getVersionId: SQLiteStatement<[string], { id: number }>
+    setCompactionFrontier: SQLiteStatement<[string]>
+    getCompactionFrontier: SQLiteStatement<[], { value: string }>
+    deleteMeta: SQLiteStatement
+    getAllKeys: SQLiteStatement<[], { key: string }>
+    getVersionsForKey: SQLiteStatement<
       [string],
       { version: string; version_id: number }
     >
-    moveDataToNewVersion: Database.Statement<[number, string, number]>
-    deleteOldVersionData: Database.Statement<[string, number]>
+    moveDataToNewVersion: SQLiteStatement<[number, string, number]>
+    deleteOldVersionData: SQLiteStatement<[string, number]>
   }
 
-  constructor(db: Database.Database, name: string, isTemp = false) {
+  constructor(db: SQLiteDb, name: string, isTemp = false) {
     this.#db = db
     this.#tableName = `index_${name}`
     this.#versionTableName = `${this.#tableName}_versions`

packages/d2ts/tests/operators-sqlite/consolidate.test.ts

@@ -6,17 +6,19 @@ import { DataMessage, MessageType } from '../../src/types.js'
 import { consolidate } from '../../src/sqlite/operators/consolidate.js'
 import { output } from '../../src/operators/index.js'
 import Database from 'better-sqlite3'
+import { BetterSQLite3Wrapper } from '../../src/sqlite/database.js'
 import fs from 'fs'
 import path from 'path'
 
 const DB_FILENAME = 'test-consolidate.db'
 
 describe('SQLite Operators', () => {
   describe('Consolidate operation', () => {
-    let db: Database.Database
+    let db: BetterSQLite3Wrapper
 
     beforeEach(() => {
-      db = new Database(':memory:')
+      const sqlite = new Database(':memory:')
+      db = new BetterSQLite3Wrapper(sqlite)
     })
 
     afterEach(() => {
@@ -126,13 +128,14 @@ describe('SQLite Operators', () => {
 
   describe('Consolidate operation with persistence', () => {
     const dbPath = path.join(import.meta.dirname, DB_FILENAME)
-    let db: Database.Database
+    let db: BetterSQLite3Wrapper
 
     beforeEach(() => {
       if (fs.existsSync(dbPath)) {
         fs.unlinkSync(dbPath)
       }
-      db = new Database(dbPath)
+      const sqlite = new Database(dbPath)
+      db = new BetterSQLite3Wrapper(sqlite)
     })
 
     afterEach(() => {
@@ -175,7 +178,7 @@ describe('SQLite Operators', () => {
 
       // Create new graph instance with same database
       messages = []
-      db = new Database(dbPath)
+      db = new BetterSQLite3Wrapper(new Database(dbPath))
       graph = new D2({ initialFrontier: v([1, 0]) })
       const newInput = graph.newInput<number>()