Tests, fixes, cleanup.

Author: rkistner

.gitignore

@@ -0,0 +1,3 @@
+node_modules/
+test-db/
+*.db

package.json

@@ -1,4 +1,7 @@
 {
+  "scripts": {
+    "test": "vitest"
+  },
   "dependencies": {
     "better-sqlite3": "^9.5.0",
     "event-iterator": "^2.0.0"

src/driver-api.ts

@@ -2,6 +2,7 @@ import { SqliteArguments, SqliteValue } from "./common.js";
 
 export interface SqliteDriverConnection {
   prepare(query: string): SqliteDriverStatement;
+  close(): Promise<void>;
 }
 
 export interface SqliteDriverConnectionPool {
@@ -14,6 +15,8 @@ export interface SqliteDriverConnectionPool {
   reserveConnection(
     options?: ReserveConnectionOptions
   ): Promise<ReservedConnection>;
+
+  close(): Promise<void>;
 }
 
 export interface ReservedConnection {
@@ -27,16 +30,27 @@ export interface ReserveConnectionOptions {
 }
 
 export interface SqliteDriverStatement {
-  execute(args?: SqliteArguments): Promise<void>;
+  run(args?: SqliteArguments): Promise<void>;
+  runWithResults(args?: SqliteArguments): Promise<RunResults>;
 
-  stream(
+  selectStreamed(
     args?: SqliteArguments,
     options?: ExecuteOptions
   ): AsyncIterable<ResultSet>;
 
+  selectAll(
+    args?: SqliteArguments,
+    options?: ExecuteOptions
+  ): Promise<ResultSet>;
+
   dispose(): void;
 }
 
+export interface RunResults {
+  changes: number;
+  lastInsertRowId: bigint;
+}
+
 export interface ResultSet {
   columns: string[];
   rows: SqliteValue[][];

src/driver-util.ts

@@ -17,6 +17,10 @@ export class SingleConnectionPool implements SqliteDriverConnectionPool {
 
   constructor(private connection: SqliteDriverConnection) {}
 
+  async close() {
+    await this.connection.close();
+  }
+
   reserveConnection(
     options?: ReserveConnectionOptions
   ): Promise<ReservedConnection> {
@@ -150,6 +154,13 @@ class MultiConnectionPool implements SqliteDriverConnectionPool {
       },
     });
   }
+
+  async close() {
+    // TODO: Wait for statements to finish
+    for (let con of this._allConnections) {
+      await con.close();
+    }
+  }
 }
 
 export class ReadWriteConnectionPool implements SqliteDriverConnectionPool {
@@ -177,4 +188,9 @@ export class ReadWriteConnectionPool implements SqliteDriverConnectionPool {
       return this.writePool!.reserveConnection(options);
     }
   }
+
+  async close() {
+    await this.readPool.close();
+    await this.writePool?.close();
+  }
 }

src/drivers/better-sqlite3-async-driver.ts

@@ -1,6 +1,8 @@
 import type * as bsqlite from "better-sqlite3";
 import * as worker_threads from "worker_threads";
 import {
+  ResultSet,
+  RunResults,
   SqliteDriverConnection,
   SqliteDriverConnectionPool,
   SqliteDriverStatement,
@@ -37,19 +39,37 @@ export class BetterSqliteAsyncConnection implements SqliteDriverConnection {
     this.worker = worker;
   }
 
+  async close() {
+    this.worker.postMessage(["close"]);
+    await new Promise<void>((resolve, reject) => {
+      this.worker.once("message", (value) => {
+        resolve();
+      });
+    });
+    await this.worker.terminate();
+  }
+
   prepare(query: string): SqliteDriverStatement {
     const worker = this.worker;
     return {
       dispose() {},
-      execute: async (args) => {
-        worker.postMessage(["execute", { query, args }]);
+      run: async (args) => {
+        worker.postMessage(["run", { query, args }]);
         return new Promise<void>((resolve, reject) => {
           worker.once("message", (value) => {
             resolve();
           });
         });
       },
-      async *stream(args, options) {
+      runWithResults: async (args) => {
+        worker.postMessage(["run", { query, args }]);
+        return new Promise<RunResults>((resolve, reject) => {
+          worker.once("message", (value) => {
+            resolve(value);
+          });
+        });
+      },
+      async *selectStreamed(args, options) {
         worker.postMessage(["stream", { query, args, options }]);
         const iter = new EventIterator(({ push }) => {
           worker.addListener("message", push);
@@ -67,6 +87,18 @@ export class BetterSqliteAsyncConnection implements SqliteDriverConnection {
           }
         }
       },
+
+      async selectAll(args, options): Promise<ResultSet> {
+        let results: ResultSet | undefined = undefined;
+        for await (let rs of this.selectStreamed(args, options)) {
+          if (results == null) {
+            results = rs;
+          } else {
+            results!.rows.push(...rs.rows);
+          }
+        }
+        return results!;
+      },
     };
   }
 }

src/drivers/better-sqlite3-driver.ts

@@ -2,6 +2,7 @@ import { SqliteArguments, SqliteValue } from "../common.js";
 import {
   ExecuteOptions,
   ResultSet,
+  RunResults,
   SqliteDriverConnection,
   SqliteDriverConnectionPool,
   SqliteDriverStatement,
@@ -29,20 +30,44 @@ export class BetterSqliteConnection implements SqliteDriverConnection {
   con: bsqlite.Database;
 
   constructor(path: string, options?: bsqlite.Options) {
-    console.log("options", options);
     this.con = new Database(path, options);
   }
 
   prepare(query: string): SqliteDriverStatement {
     const stmt = this.con.prepare(query);
     return new BetterSqliteStatement(stmt);
   }
+
+  async close() {
+    this.con.close();
+  }
 }
 
 export class BetterSqliteStatement implements SqliteDriverStatement {
   constructor(private statement: bsqlite.Statement) {}
 
-  async *stream(
+  async selectAll(
+    args?: SqliteArguments | undefined,
+    options?: ExecuteOptions | undefined
+  ): Promise<ResultSet> {
+    const bindArgs = args == undefined ? [] : [args];
+    if (!this.statement.reader) {
+      this.statement.run(...bindArgs);
+      return { columns: [], rows: [] };
+    }
+    this.statement.raw();
+    if (options?.bigint) {
+      this.statement.safeIntegers();
+    }
+    const columns = this.statement.columns().map((c) => c.name);
+    const rows = this.statement.all(...bindArgs) as SqliteValue[][];
+    return {
+      columns,
+      rows,
+    };
+  }
+
+  async *selectStreamed(
     args?: SqliteArguments,
     options?: ExecuteOptions
   ): AsyncGenerator<ResultSet, any, undefined> {
@@ -58,7 +83,7 @@ export class BetterSqliteStatement implements SqliteDriverStatement {
     const columns = this.statement.columns().map((c) => c.name);
     let buffer: SqliteValue[][] = [];
     let didYield = false;
-    for (let row of this.statement.all(...bindArgs)) {
+    for (let row of this.statement.iterate(...bindArgs)) {
       buffer.push(row as SqliteValue[]);
       if (buffer.length > (options?.chunkSize ?? 10)) {
         yield {
@@ -77,8 +102,18 @@ export class BetterSqliteStatement implements SqliteDriverStatement {
     }
   }
 
-  async execute(args?: SqliteArguments): Promise<void> {
-    this.statement.run();
+  async run(args?: SqliteArguments): Promise<void> {
+    const bindArgs = args == undefined ? [] : [args];
+    this.statement.run(...bindArgs);
+  }
+
+  async runWithResults(args?: SqliteArguments): Promise<RunResults> {
+    const bindArgs = args == undefined ? [] : [args];
+    const r = this.statement.run(...bindArgs);
+    return {
+      changes: r.changes,
+      lastInsertRowId: BigInt(r.lastInsertRowid),
+    };
   }
 
   dispose(): void {

src/drivers/better-sqlite3-worker.js

@@ -8,14 +8,20 @@ worker_threads.parentPort.addListener("message", (value) => {
   const [message, args] = value;
   if (message == "open") {
     db = new Database(args.path);
-  } else if (message == "execute") {
-    if (args.args == null) {
-      const rs = db.prepare(args.query).all();
-      worker_threads.parentPort.postMessage(rs);
-    } else {
-      const rs = db.prepare(args.query).all(args.args);
-      worker_threads.parentPort.postMessage(rs);
-    }
+  } else if (message == "close") {
+    db?.close();
+    worker_threads.parentPort.postMessage(["closed"]);
+  } else if (message == "run") {
+    const bindArgs = args.args == undefined ? [] : [args.args];
+    const rs = db.prepare(args.query).run(...bindArgs);
+    worker_threads.parentPort.postMessage({
+      changes: rs.changes,
+      lastInsertRowId: BigInt(rs.lastInsertRowid),
+    });
+  } else if (message == "selectAll") {
+    const bindArgs = args.args == undefined ? [] : [args.args];
+    const rs = db.prepare(args.query).all(...bindArgs);
+    worker_threads.parentPort.postMessage(rs);
   } else if (message == "stream") {
     const bindArgs = args.args == undefined ? [] : [args.args];
     const statement = db.prepare(args.query);

src/impl.ts

@@ -85,15 +85,15 @@ export class ConnectionImpl implements SqliteConnection {
     callback: (tx: SqliteTransaction) => Promise<T>,
     options: TransactionOptions
   ): Promise<T> {
-    await this.driver.prepare("BEGIN").execute();
+    await this.driver.prepare("BEGIN").run();
     try {
       const tx = new TransactionImpl(this);
       const result = await callback(tx);
 
-      await this.driver.prepare("COMMIT").execute();
+      await this.driver.prepare("COMMIT").run();
       return result;
     } catch (e) {
-      await this.driver.prepare("ROLLBACK").execute();
+      await this.driver.prepare("ROLLBACK").run();
       throw e;
     }
   }
@@ -135,7 +135,7 @@ export class ConnectionImpl implements SqliteConnection {
 
     const q = this.driver.prepare(query as string);
 
-    for await (let rs of q.stream(args)) {
+    for await (let rs of q.selectStreamed(args)) {
       if (result == null) {
         result = new ResultSetImpl(rs.columns, [...rs.rows]);
       } else {
@@ -152,7 +152,7 @@ export class ConnectionImpl implements SqliteConnection {
   ): AsyncGenerator<ResultSet<T>, void, unknown> {
     const q = this.driver.prepare(query as string);
 
-    for await (let rs of q.stream(args, options)) {
+    for await (let rs of q.selectStreamed(args, options)) {
       yield new ResultSetImpl(rs.columns, rs.rows);
     }
   }

test/better-sqlite3-async.test.ts

@@ -0,0 +1,4 @@
+import { betterSqliteAsyncPool } from "../src/drivers/better-sqlite3-async-driver.js";
+import { describeDriverTests } from "./tests/drivers.js";
+
+describeDriverTests("better-sqlite3-async-pool", betterSqliteAsyncPool);

test/better-sqlite3.test.ts

@@ -0,0 +1,4 @@
+import { betterSqlitePool } from "../src/drivers/better-sqlite3-driver.js";
+import { describeDriverTests } from "./tests/drivers.js";
+
+describeDriverTests("better-sqlite3", betterSqlitePool);