Browse Source

fix: verify sqlite-vec readiness after extension load. Closes #169

Claude 3 months ago
parent
commit
73136e4f59
2 changed files with 56 additions and 6 deletions
  1. 21 0
      src/store.test.ts
  2. 35 6
      src/store.ts

+ 21 - 0
src/store.test.ts

@@ -8,6 +8,7 @@
 
 
 import { describe, test, expect, beforeAll, afterAll, beforeEach, afterEach, mock, spyOn } from "bun:test";
 import { describe, test, expect, beforeAll, afterAll, beforeEach, afterEach, mock, spyOn } from "bun:test";
 import { Database } from "bun:sqlite";
 import { Database } from "bun:sqlite";
+import * as sqliteVec from "sqlite-vec";
 import { unlink, mkdtemp, rmdir, writeFile } from "node:fs/promises";
 import { unlink, mkdtemp, rmdir, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { join } from "node:path";
@@ -15,6 +16,7 @@ import YAML from "yaml";
 import { disposeDefaultLlamaCpp } from "./llm.js";
 import { disposeDefaultLlamaCpp } from "./llm.js";
 import {
 import {
   createStore,
   createStore,
+  verifySqliteVecLoaded,
   getDefaultDbPath,
   getDefaultDbPath,
   homedir,
   homedir,
   resolve,
   resolve,
@@ -452,6 +454,25 @@ describe("Store Creation", () => {
     await cleanupTestDb(store);
     await cleanupTestDb(store);
   });
   });
 
 
+  test("verifySqliteVecLoaded throws when sqlite-vec is not loaded", () => {
+    const db = new Database(":memory:");
+    try {
+      expect(() => verifySqliteVecLoaded(db)).toThrow("sqlite-vec extension is unavailable");
+    } finally {
+      db.close();
+    }
+  });
+
+  test("verifySqliteVecLoaded succeeds when sqlite-vec is loaded", () => {
+    const db = new Database(":memory:");
+    try {
+      sqliteVec.load(db);
+      expect(() => verifySqliteVecLoaded(db)).not.toThrow();
+    } finally {
+      db.close();
+    }
+  });
+
   test("store.close closes the database connection", async () => {
   test("store.close closes the database connection", async () => {
     const store = await createTestStore();
     const store = await createTestStore();
     store.close();
     store.close();

+ 35 - 6
src/store.ts

@@ -458,17 +458,46 @@ function setSQLiteFromBrewPrefixEnv(): void {
 
 
 setSQLiteFromBrewPrefixEnv();
 setSQLiteFromBrewPrefixEnv();
 
 
+function createSqliteVecUnavailableError(reason: string): Error {
+  return new Error(
+    "sqlite-vec extension is unavailable. " +
+    `${reason}. ` +
+    "Install Homebrew SQLite so the sqlite-vec extension can be loaded, " +
+    "and set BREW_PREFIX if Homebrew is installed in a non-standard location."
+  );
+}
+
+function getErrorMessage(err: unknown): string {
+  return err instanceof Error ? err.message : String(err);
+}
+
+export function verifySqliteVecLoaded(db: Database): void {
+  try {
+    const row = db.prepare(`SELECT vec_version() AS version`).get() as { version?: string } | null;
+    if (!row?.version || typeof row.version !== "string") {
+      throw new Error("vec_version() returned no version");
+    }
+  } catch (err) {
+    const message = getErrorMessage(err);
+    throw createSqliteVecUnavailableError(`sqlite-vec probe failed (${message})`);
+  }
+}
+
 function initializeDatabase(db: Database): void {
 function initializeDatabase(db: Database): void {
   try {
   try {
     sqliteVec.load(db);
     sqliteVec.load(db);
+    verifySqliteVecLoaded(db);
   } catch (err) {
   } catch (err) {
-    if (err instanceof Error && err.message.includes("does not support dynamic extension loading")) {
-      throw new Error(
-        "SQLite build does not support dynamic extension loading. " +
-        "Install Homebrew SQLite so the sqlite-vec extension can be loaded, " +
-        "and set BREW_PREFIX if Homebrew is installed in a non-standard location."
-      );
+    const message = getErrorMessage(err);
+
+    if (message.includes("does not support dynamic extension loading")) {
+      throw createSqliteVecUnavailableError("SQLite build does not support dynamic extension loading");
     }
     }
+
+    if (message.includes("sqlite-vec extension is unavailable")) {
+      throw err;
+    }
+
     throw err;
     throw err;
   }
   }
   db.exec("PRAGMA journal_mode = WAL");
   db.exec("PRAGMA journal_mode = WAL");