diff --git a/packages/openmemory-js/src/core/migrate.ts b/packages/openmemory-js/src/core/migrate.ts index f21b9136..2865a744 100644 --- a/packages/openmemory-js/src/core/migrate.ts +++ b/packages/openmemory-js/src/core/migrate.ts @@ -21,6 +21,7 @@ const LEGACY_SQLITE_VECTOR_TABLE = "vectors"; interface Migration { version: string; desc: string; + completionColumn: string; sqlite: (vectorTable: string) => string[]; postgres: string[]; } @@ -29,6 +30,7 @@ const migrations: Migration[] = [ { version: "1.2.0", desc: "Multi-user tenant support", + completionColumn: "user_id", sqlite: (vectorTable: string) => [ `ALTER TABLE memories ADD COLUMN user_id TEXT`, `CREATE INDEX IF NOT EXISTS idx_memories_user ON memories(user_id)`, @@ -76,6 +78,7 @@ const migrations: Migration[] = [ { version: "1.3.0", desc: "Project-level isolation support", + completionColumn: "project_id", sqlite: (vectorTable: string) => [ `ALTER TABLE memories ADD COLUMN project_id TEXT`, `CREATE INDEX IF NOT EXISTS idx_memories_project ON memories(project_id)`, @@ -193,14 +196,14 @@ async function run_sqlite_migration( ): Promise { log(`Running migration: ${m.version} - ${m.desc}`); - const has_user_id = await check_column_exists_sqlite( + const is_complete = await check_column_exists_sqlite( db, "memories", - "user_id", + m.completionColumn, ); - if (has_user_id) { + if (is_complete) { log( - `Migration ${m.version} already applied (user_id exists), skipping`, + `Migration ${m.version} already applied (${m.completionColumn} exists), skipping`, ); await set_db_version_sqlite(db, m.version); return; @@ -296,11 +299,15 @@ async function run_pg_migration(pool: Pool, m: Migration): Promise { process.env.OM_VECTOR_TABLE || DEFAULT_VECTOR_TABLE, "OM_VECTOR_TABLE", ); - const has_user_id = await check_column_exists_pg(pool, mt, "user_id"); + const is_complete = await check_column_exists_pg( + pool, + mt, + m.completionColumn, + ); - if (has_user_id) { + if (is_complete) { log( - `Migration ${m.version} already applied (user_id exists), skipping`, + `Migration ${m.version} already applied (${m.completionColumn} exists), skipping`, ); await set_db_version_pg(pool, m.version); return; @@ -392,6 +399,49 @@ async function quarantine_orphan_temporal_facts_pg(pool: Pool): Promise { } } +export async function run_sqlite_migrations( + db: sqlite3.Database, +): Promise { + const current = await get_db_version_sqlite(db); + log(`Current database version: ${current || "none"}`); + + for (const m of migrations) { + const is_complete = await check_column_exists_sqlite( + db, + "memories", + m.completionColumn, + ); + if (!is_complete || !current || m.version > current) { + await run_sqlite_migration(db, m); + } + } + + await quarantine_orphan_temporal_facts_sqlite(db); +} + +async function run_pg_migrations(pool: Pool): Promise { + const current = await get_db_version_pg(pool); + log(`Current database version: ${current || "none"}`); + + const mt = assertSafeIdentifier( + process.env.OM_PG_TABLE || "openmemory_memories", + "OM_PG_TABLE", + ); + + for (const m of migrations) { + const is_complete = await check_column_exists_pg( + pool, + mt, + m.completionColumn, + ); + if (!is_complete || !current || m.version > current) { + await run_pg_migration(pool, m); + } + } + + await quarantine_orphan_temporal_facts_pg(pool); +} + export async function run_migrations() { log("Checking for pending migrations..."); @@ -411,32 +461,14 @@ export async function run_migrations() { ssl, }); - const current = await get_db_version_pg(pool); - log(`Current database version: ${current || "none"}`); - - for (const m of migrations) { - if (!current || m.version > current) { - await run_pg_migration(pool, m); - } - } - - await quarantine_orphan_temporal_facts_pg(pool); + await run_pg_migrations(pool); await pool.end(); } else { const db_path = process.env.OM_DB_PATH || "./data/openmemory.sqlite"; const db = new sqlite3.Database(db_path); - const current = await get_db_version_sqlite(db); - log(`Current database version: ${current || "none"}`); - - for (const m of migrations) { - if (!current || m.version > current) { - await run_sqlite_migration(db, m); - } - } - - await quarantine_orphan_temporal_facts_sqlite(db); + await run_sqlite_migrations(db); await new Promise((ok) => db.close(() => ok())); } diff --git a/packages/openmemory-js/tests/migrate.test.ts b/packages/openmemory-js/tests/migrate.test.ts new file mode 100644 index 00000000..42db2011 --- /dev/null +++ b/packages/openmemory-js/tests/migrate.test.ts @@ -0,0 +1,94 @@ +import sqlite3 from "sqlite3"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { run_sqlite_migrations } from "../src/core/migrate"; + +const exec = (db: sqlite3.Database, sql: string): Promise => + new Promise((resolve, reject) => { + db.exec(sql, (err) => (err ? reject(err) : resolve())); + }); + +const all = (db: sqlite3.Database, sql: string): Promise => + new Promise((resolve, reject) => { + db.all(sql, (err, rows) => (err ? reject(err) : resolve(rows as T[]))); + }); + +const columns = async (db: sqlite3.Database, table: string) => { + const rows = await all<{ name: string }>(db, `PRAGMA table_info(${table})`); + return rows.map((row) => row.name); +}; + +const create_corrupted_v130_schema = async (db: sqlite3.Database) => { + await exec( + db, + ` + CREATE TABLE memories (id TEXT PRIMARY KEY, user_id TEXT, content TEXT); + CREATE TABLE vectors (id TEXT, sector TEXT, user_id TEXT); + CREATE TABLE waypoints (src_id TEXT, dst_id TEXT, user_id TEXT); + CREATE TABLE temporal_facts (id TEXT PRIMARY KEY, user_id TEXT); + CREATE TABLE schema_version (version TEXT PRIMARY KEY, applied_at INTEGER); + + INSERT INTO memories VALUES ('memory-1', 'user-1', 'keep me'); + INSERT INTO vectors VALUES ('memory-1', 'semantic', 'user-1'); + INSERT INTO waypoints VALUES ('memory-1', 'memory-1', 'user-1'); + INSERT INTO temporal_facts VALUES ('fact-1', 'user-1'); + INSERT INTO schema_version VALUES ('1.3.0', 1); + `, + ); +}; + +describe("SQLite schema migrations", () => { + let db: sqlite3.Database; + + beforeEach(async () => { + db = new sqlite3.Database(":memory:"); + await create_corrupted_v130_schema(db); + }); + + afterEach(async () => { + await new Promise((resolve, reject) => { + db.close((err) => (err ? reject(err) : resolve())); + }); + }); + + it("repairs a database marked 1.3.0 when project_id columns are missing", async () => { + await run_sqlite_migrations(db); + + for (const table of [ + "memories", + "vectors", + "waypoints", + "temporal_facts", + ]) { + expect(await columns(db, table)).toContain("project_id"); + } + + const memories = await all<{ content: string }>( + db, + "SELECT content FROM memories", + ); + expect(memories).toEqual([{ content: "keep me" }]); + }); + + it("is idempotent after repairing the missing 1.3.0 columns", async () => { + await run_sqlite_migrations(db); + await run_sqlite_migrations(db); + + const versions = await all<{ version: string }>( + db, + "SELECT version FROM schema_version ORDER BY version", + ); + expect(versions).toEqual([{ version: "1.3.0" }]); + + for (const table of [ + "memories", + "vectors", + "waypoints", + "temporal_facts", + ]) { + const projectColumns = (await columns(db, table)).filter( + (column) => column === "project_id", + ); + expect(projectColumns).toHaveLength(1); + } + }); +});