@axiom-lattice/pg-stores 1.0.83 → 1.0.85

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.83",
+  "version": "1.0.85",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -23,8 +23,8 @@
     "@langchain/langgraph-checkpoint-postgres": "1.0.0",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.92",
-    "@axiom-lattice/protocols": "2.1.46"
+    "@axiom-lattice/core": "2.1.94",
+    "@axiom-lattice/protocols": "2.1.48"
   },
   "devDependencies": {
     "@types/node": "^20.11.24",

package/scripts/check-migration-versions.mjs

@@ -1,6 +1,10 @@
 /**
- * Check for duplicate migration version numbers.
- * Exits with code 1 if any version collision is found.
+ * Migration identity checker.
+ * 
+ * Migration identity is now based on **name**, not version.
+ * Duplicate names are ERRORS (same name = same migration = can't register twice).
+ * Duplicate versions are WARNINGS (allowed but should be intentional).
+ * 
  * Run as part of lint pipeline.
  */
 import { readFileSync, readdirSync } from "node:fs";
@@ -10,22 +14,47 @@ import { fileURLToPath } from "node:url";
 const MIGRATIONS_DIR = join(dirname(fileURLToPath(import.meta.url)), "../src/migrations");
 
 const files = readdirSync(MIGRATIONS_DIR).filter(f => f.endsWith(".ts"));
+const nameMap = new Map();
 const versionMap = new Map();
 
 let hasError = false;
 
 for (const file of files) {
   const content = readFileSync(join(MIGRATIONS_DIR, file), "utf-8");
-  const matches = content.matchAll(/^\s*version:\s*(\d+)/gm);
-  for (const m of matches) {
-    const version = parseInt(m[1]);
-    if (versionMap.has(version)) {
+
+  // Extract name and version from migration definitions
+  // Matches patterns like:
+  //   name: "create_threads_table",
+  //   version: 1,
+  const nameMatches = content.matchAll(/^\s*name:\s*"([^"]+)"/gm);
+  const versionMatches = content.matchAll(/^\s*version:\s*(\d+)/gm);
+
+  const names = [...nameMatches].map(m => m[1]);
+  const versions = [...versionMatches].map(m => parseInt(m[1]));
+
+  for (let i = 0; i < names.length; i++) {
+    const name = names[i];
+    const version = versions[i];
+
+    // NAME conflict = ERROR (identity collision)
+    if (nameMap.has(name)) {
       console.error(
-        `\x1b[31mMIGRATION VERSION CONFLICT:\x1b[0m v${version} in both "${versionMap.get(version)}" and "${file}"`
+        `\x1b[31mMIGRATION NAME CONFLICT:\x1b[0m "${name}" (v${version} in "${file}") ` +
+        `collides with "${name}" in "${nameMap.get(name)}"`
       );
       hasError = true;
     } else {
-      versionMap.set(version, file);
+      nameMap.set(name, file);
+    }
+
+    // VERSION duplicate = WARNING (allowed but flag for review)
+    if (versionMap.has(version)) {
+      console.warn(
+        `\x1b[33mDuplicate version:\x1b[0m v${version} used by "${name}" (${file}) ` +
+        `and in "${versionMap.get(version)}" — ensure no ordering dependency`
+      );
+    } else {
+      versionMap.set(version, `${name} (${file})`);
     }
   }
 }
@@ -34,4 +63,8 @@ if (hasError) {
   process.exit(1);
 }
 
-console.log(`\x1b[32mMigration versions OK\x1b[0m (${versionMap.size} unique versions in ${files.length} files)`);
+console.log(`\x1b[32mMigration names OK\x1b[0m (${nameMap.size} unique names in ${files.length} files)`);
+const dupVersions = [...versionMap.keys()].length;
+if (dupVersions !== nameMap.size) {
+  console.log(`\x1b[33m${nameMap.size - dupVersions} version(s) shared across migrations\x1b[0m`);
+}

package/src/createPgStoreConfig.ts

@@ -1,3 +1,5 @@
+import { Pool } from "pg";
+import { MigrationManager } from "./migrations/migration";
 import { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";
 import { PostgreSQLAssistantStore } from "./stores/PostgreSQLAssistantStore";
 import { PostgreSQLDatabaseConfigStore } from "./stores/PostgreSQLDatabaseConfigStore";
@@ -20,14 +22,106 @@ import { MenuStore } from "./stores/MenuStore";
 import { PostgresSharedResourceStore } from "./stores/PostgresSharedResourceStore";
 import { PostgresSaver } from "@langchain/langgraph-checkpoint-postgres";
 
-export function createPgStoreConfig(connectionString: string) {
-  const opts = { poolConfig: connectionString, autoMigrate: true } as const;
+// Migration imports — all migrations registered centrally to guarantee
+// ordering and prevent name collisions across features.
+import { createThreadsTable } from "./migrations/thread_migrations";
+import { addThreadTenantId } from "./migrations/thread_tenant_migration";
+import { changeThreadPrimaryKey } from "./migrations/thread_pk_migration";
+import { createScheduledTasksTable } from "./migrations/schedule_migrations";
+import { addScheduleTenantId } from "./migrations/schedule_tenant_migration";
+import { createAssistantsTable } from "./migrations/assistant_migrations";
+import { addAssistantTenantId } from "./migrations/assistant_tenant_migration";
+import { changeAssistantPrimaryKey } from "./migrations/assistant_pk_migration";
+import { addAssistantOwnerUserId } from "./migrations/assistant_owner_user_migration";
+import { createSkillsTable } from "./migrations/skill_migrations";
+import { addSkillTenantId } from "./migrations/skill_tenant_migration";
+import { changeSkillPrimaryKey } from "./migrations/skill_pk_migration";
+import { createWorkspacesTable } from "./migrations/workspace_migrations";
+import { createProjectsTable } from "./migrations/project_migrations";
+import { addProjectConfigColumn } from "./migrations/add_project_config_column";
+import { createUsersTable } from "./migrations/user_migrations";
+import { addUserStatusColumn } from "./migrations/user_status_migration";
+import { createTenantsTable } from "./migrations/tenant_migrations";
+import { createUserTenantLinkTable, migrateRemoveTenantIdFromUsers } from "./migrations/user_tenant_link_migrations";
+import { createChannelIdentityMappingTables } from "./migrations/channel_identity_mapping_migration";
+import { createChannelInstallationsTable } from "./migrations/channel_installation_migrations";
+import { alterChannelInstallationsTable } from "./migrations/channel_installations_alter_migration";
+import { createChannelBindingsTable } from "./migrations/channel_bindings_migration";
+import { createDatabaseConfigsTable } from "./migrations/database_config_migrations";
+import { createMetricsConfigsTable } from "./migrations/metrics_config_migrations";
+import { createMcpServerConfigsTable } from "./migrations/mcp_server_config_migrations";
+import { createThreadMessageQueueTable } from "./migrations/thread_message_queue_migrations";
+import { addPriorityAndCommandColumns } from "./migrations/add_priority_command_columns";
+import { addCustomRunConfigColumn } from "./migrations/add_custom_run_config_column";
+import { alterMessageQueueIdColumn } from "./migrations/alter_message_queue_id_column";
+import { addWorkspaceProjectToQueue } from "./migrations/add_workspace_project_to_queue";
+import { createWorkflowTrackingTables, addStepThreadId } from "./migrations/workflow_tracking_migrations";
+import { evalMigrations } from "./migrations/eval_migrations";
+import { createA2AApiKeysTable } from "./migrations/a2a_api_key_migration";
+import { createTasksTable } from "./migrations/task_migration";
+import { createMenuItemsTable } from "./migrations/menu_items_migration";
+import { addFileContentType } from "./migrations/menu_items_add_file_type";
+import { createSharedResourcesTable } from "./migrations/shared_resources_migration";
+
+export async function createPgStoreConfig(connectionString: string) {
+  const pool = new Pool({ connectionString });
+
+  // Centralized migration manager — all migrations registered here
+  // in one place. Name conflicts throw immediately. Execution order
+  // follows version numbers globally.
+  const mm = new MigrationManager(pool);
+  mm.register(createThreadsTable);          // v1
+  mm.register(createScheduledTasksTable);   // v2
+  mm.register(createAssistantsTable);       // v3
+  mm.register(createSkillsTable);           // v4
+  mm.register(createWorkspacesTable);       // v5
+  mm.register(createProjectsTable);         // v6
+  mm.register(createUsersTable);            // v7
+  mm.register(createTenantsTable);          // v8
+  mm.register(addUserStatusColumn);         // v9
+  mm.register(createUserTenantLinkTable);   // v10
+  mm.register(migrateRemoveTenantIdFromUsers); // v11
+  mm.register(createMetricsConfigsTable);   // v12
+  mm.register(createMcpServerConfigsTable); // v13
+  mm.register(addAssistantTenantId);        // v14
+  mm.register(addSkillTenantId);            // v16
+  mm.register(changeSkillPrimaryKey);       // v17
+  mm.register(addScheduleTenantId);         // v18
+  mm.register(changeThreadPrimaryKey);      // v19
+  mm.register(createChannelIdentityMappingTables); // v20
+  mm.register(createChannelInstallationsTable);    // v21
+  mm.register(addThreadTenantId);           // v23
+  mm.register(addProjectConfigColumn);      // v24
+  mm.register(createThreadMessageQueueTable); // v100
+  mm.register(addPriorityAndCommandColumns);  // v101
+  mm.register(alterMessageQueueIdColumn);     // v102
+  mm.register(createWorkflowTrackingTables);  // v103
+  for (const m of evalMigrations) {           // v104-v108
+    mm.register(m);
+  }
+  mm.register(alterChannelInstallationsTable); // v109
+  mm.register(createChannelBindingsTable);      // v110
+  mm.register(createDatabaseConfigsTable);      // v111
+  mm.register(changeAssistantPrimaryKey);       // v112
+  mm.register(addStepThreadId);                 // v113
+  mm.register(addCustomRunConfigColumn);        // v114
+  mm.register(createA2AApiKeysTable);           // v130
+  mm.register(addWorkspaceProjectToQueue);      // v131
+  mm.register(addAssistantOwnerUserId);         // v132
+  mm.register(createTasksTable);                 // v133
+  mm.register(createMenuItemsTable);             // v134
+  mm.register(createSharedResourcesTable);       // v135
+  mm.register(addFileContentType);               // v136
+
+  await mm.migrate();
 
   const checkpoint = PostgresSaver.fromConnString(connectionString);
   checkpoint.setup().catch((err) => {
     console.error("[pg-stores] Failed to setup checkpoint table:", err.message || err);
   });
 
+  const opts = { pool };
+
   return {
     workspace: new PostgreSQLWorkspaceStore(opts),
     project: new PostgreSQLProjectStore(opts),
@@ -51,4 +145,4 @@ export function createPgStoreConfig(connectionString: string) {
     sharedResource: new PostgresSharedResourceStore(opts),
     checkpoint,
   };
-}
+}

package/src/migrations/migration.ts

@@ -1,30 +1,29 @@
 /**
  * Migration system for database schema management
  *
- * ## Migration Version Rules
+ * ## Identity vs Ordering
  *
- * Version numbers are globally unique across all migrations applied to the same
- * database. They are NOT auto-assigned — you must pick one manually.
+ * **name = identity** — each migration is uniquely identified by its name.
+ * Two migrations MAY share the same version number; they are tracked
+ * independently by name and will both execute.
  *
- * ### Numbering scheme
+ * **version = ordering** — migrations are sorted and applied in ascending
+ * version order. Sharing a version simply means two migrations are at the
+ * same position in the dependency chain (they operate on different tables).
+ *
+ * ### Version numbering scheme
  *
  * | Range   | Purpose                              |
  * |---------|--------------------------------------|
  * | 1–99    | Core table creation + column changes |
  * | 100+    | Independent feature modules          |
  *
- * ### How to pick a version
- *
- * 1. Run `grep -rn "version:" src/migrations/ | sort -t: -k3 -n` to list all versions
- * 2. For a core column change: find the next unused number in the 1–99 range
- * 3. For a new feature module: find the next unused number in the 100+ range
- * 4. Verify your chosen version is NOT in the grep output
- *
  * @remarks
- * - Migrations are sorted and applied in version order on startup
+ * - Migrations are applied in version order on startup
  * - The version must be strictly greater than any migration it depends on
  * - Use `IF NOT EXISTS` / `IF EXISTS` in SQL to keep migrations idempotent
- * - A version number collision will cause one migration to be silently skipped
+ * - Duplicate **names** are rejected at registration time
+ * - Duplicate **versions** are allowed (different features, different tables)
  */
 
 import { Pool, PoolClient } from "pg";
@@ -33,8 +32,8 @@ import { Pool, PoolClient } from "pg";
  * Migration record stored in database
  */
 interface MigrationRecord {
-  version: number;
   name: string;
+  version: number;
   applied_at: Date;
 }
 
@@ -60,45 +59,84 @@ export class MigrationManager {
   }
 
   /**
-   * Register a migration
+   * Register a migration.
+   * Duplicate names are rejected immediately (name = identity).
+   * Duplicate versions are allowed (version = ordering only).
    */
   register(migration: Migration): void {
+    const existing = this.migrations.find((m) => m.name === migration.name);
+    if (existing) {
+      throw new Error(
+        `Migration name conflict: "${migration.name}" v${migration.version} ` +
+        `collides with existing registration at v${existing.version}`
+      );
+    }
     this.migrations.push(migration);
-    // Sort migrations by version
     this.migrations.sort((a, b) => a.version - b.version);
   }
 
   /**
-   * Initialize migrations table if it doesn't exist
+   * Ensure the tracking table exists with the correct schema.
+   * Handles migration from old schema (PK on version) to new schema (PK on name).
    */
   private async ensureMigrationsTable(client: PoolClient): Promise<void> {
-    await client.query(`
-      CREATE TABLE IF NOT EXISTS lattice_schema_migrations (
-        version INTEGER PRIMARY KEY,
-        name VARCHAR(255) NOT NULL,
-        applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+    const tableExists = await client.query(`
+      SELECT EXISTS (
+        SELECT FROM information_schema.tables
+        WHERE table_name = 'lattice_schema_migrations'
       )
     `);
+
+    if (!tableExists.rows[0].exists) {
+      await client.query(`
+        CREATE TABLE lattice_schema_migrations (
+          name VARCHAR(255) PRIMARY KEY,
+          version INTEGER NOT NULL,
+          applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+        )
+      `);
+      return;
+    }
+
+    // Check if we need to migrate PK from version to name
+    const pkCheck = await client.query<{ column_name: string; constraint_name: string }>(`
+      SELECT a.attname AS column_name, con.conname AS constraint_name
+      FROM pg_index i
+      JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
+      JOIN pg_constraint con ON con.conindid = i.indexrelid
+      WHERE i.indrelid = 'lattice_schema_migrations'::regclass
+        AND i.indisprimary
+    `);
+
+    const pkColumn = pkCheck.rows[0]?.column_name;
+    const constraintName = pkCheck.rows[0]?.constraint_name;
+    if (pkColumn === 'version' && constraintName) {
+      // Deduplicate names before changing PK. Under the old schema, the
+      // same migration name could be inserted at multiple versions by
+      // different MigrationManagers. Keep only the highest version.
+      await client.query(`
+        DELETE FROM lattice_schema_migrations a
+        USING lattice_schema_migrations b
+        WHERE a.name = b.name AND a.version < b.version
+      `);
+
+      await client.query(
+        `ALTER TABLE lattice_schema_migrations DROP CONSTRAINT ${constraintName}`
+      );
+      await client.query(
+        `ALTER TABLE lattice_schema_migrations ADD PRIMARY KEY (name)`
+      );
+    }
   }
 
   /**
-   * Get applied migrations from database
+   * Get the set of already-applied migration names from the tracking table.
    */
-  private async getAppliedMigrations(
-    client: PoolClient
-  ): Promise<MigrationRecord[]> {
-    await client.query(`
-      CREATE TABLE IF NOT EXISTS lattice_schema_migrations (
-        version INTEGER PRIMARY KEY,
-        name VARCHAR(255) NOT NULL,
-        applied_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
-      )
-    `);
-
+  private async getAppliedNames(client: PoolClient): Promise<Set<string>> {
     const result = await client.query<MigrationRecord>(
-      "SELECT version, name, applied_at FROM lattice_schema_migrations ORDER BY version"
+      "SELECT name, version, applied_at FROM lattice_schema_migrations ORDER BY version"
     );
-    return result.rows;
+    return new Set(result.rows.map((r) => r.name));
   }
 
   /**
@@ -120,12 +158,11 @@ export class MigrationManager {
         await client.query("BEGIN");
 
         await this.ensureMigrationsTable(client);
-        const appliedMigrations = await this.getAppliedMigrations(client);
-        const appliedVersions = new Set(appliedMigrations.map((m) => m.version));
+        const appliedNames = await this.getAppliedNames(client);
 
-        // Find pending migrations
+        // Find pending migrations by name (not version)
         const pendingMigrations = this.migrations.filter(
-          (m) => !appliedVersions.has(m.version)
+          (m) => !appliedNames.has(m.name)
         );
 
         if (pendingMigrations.length === 0) {
@@ -134,16 +171,15 @@ export class MigrationManager {
           return;
         }
 
-        // Apply pending migrations
+        // Apply pending migrations (already sorted by version from register)
         for (const migration of pendingMigrations) {
           console.log(
-            `Applying migration ${migration.version}: ${migration.name}`
+            `Applying migration v${migration.version}: ${migration.name}`
           );
           await migration.up(client);
-          // Use ON CONFLICT DO NOTHING to handle race conditions gracefully
           await client.query(
-            "INSERT INTO lattice_schema_migrations (version, name) VALUES ($1, $2) ON CONFLICT (version) DO NOTHING",
-            [migration.version, migration.name]
+            "INSERT INTO lattice_schema_migrations (name, version) VALUES ($1, $2) ON CONFLICT (name) DO NOTHING",
+            [migration.name, migration.version]
           );
         }
 
@@ -162,38 +198,42 @@ export class MigrationManager {
   }
 
   /**
-   * Rollback last migration
+   * Rollback the last-applied migration (highest version).
    */
   async rollback(): Promise<void> {
     const client = await this.pool.connect();
     try {
       await client.query("BEGIN");
 
-      const appliedMigrations = await this.getAppliedMigrations(client);
-      if (appliedMigrations.length === 0) {
+      await this.ensureMigrationsTable(client);
+      const result = await client.query<MigrationRecord>(
+        "SELECT name, version FROM lattice_schema_migrations ORDER BY version DESC LIMIT 1"
+      );
+
+      if (result.rows.length === 0) {
         console.log("No migrations to rollback");
         await client.query("COMMIT");
         return;
       }
 
-      const lastMigration = appliedMigrations[appliedMigrations.length - 1];
+      const lastApplied = result.rows[0];
       const migration = this.migrations.find(
-        (m) => m.version === lastMigration.version
+        (m) => m.name === lastApplied.name
       );
 
       if (!migration || !migration.down) {
         throw new Error(
-          `Migration ${lastMigration.version} does not have a down migration`
+          `Migration "${lastApplied.name}" (v${lastApplied.version}) does not have a down migration`
         );
       }
 
       console.log(
-        `Rolling back migration ${lastMigration.version}: ${lastMigration.name}`
+        `Rolling back migration v${lastApplied.version}: ${lastApplied.name}`
       );
       await migration.down(client);
       await client.query(
-        "DELETE FROM lattice_schema_migrations WHERE version = $1",
-        [lastMigration.version]
+        "DELETE FROM lattice_schema_migrations WHERE name = $1",
+        [lastApplied.name]
       );
 
       await client.query("COMMIT");
@@ -207,16 +247,16 @@ export class MigrationManager {
   }
 
   /**
-   * Get current migration version
+   * Get the highest applied migration version.
    */
   async getCurrentVersion(): Promise<number> {
     const client = await this.pool.connect();
     try {
-      const appliedMigrations = await this.getAppliedMigrations(client);
-      if (appliedMigrations.length === 0) {
-        return 0;
-      }
-      return Math.max(...appliedMigrations.map((m) => m.version));
+      await this.ensureMigrationsTable(client);
+      const result = await client.query<{ max: number }>(
+        "SELECT COALESCE(MAX(version), 0) AS max FROM lattice_schema_migrations"
+      );
+      return result.rows[0]?.max ?? 0;
     } finally {
       client.release();
     }

package/src/stores/ChannelBindingStore.ts

@@ -9,7 +9,8 @@ import { MigrationManager } from "../migrations/migration";
 import { createChannelBindingsTable } from "../migrations/channel_bindings_migration";
 
 export interface ChannelBindingStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
@@ -33,15 +34,26 @@ type BindingRow = {
 
 export class ChannelBindingStore implements BindingRegistry {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized = false;
+  private ownsPool: boolean = true;
   private initPromise: Promise<void> | null = null;
 
   constructor(options: ChannelBindingStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     this.pool =
       typeof options.poolConfig === "string"
         ? new Pool({ connectionString: options.poolConfig })
-        : new Pool(options.poolConfig);
+        : options.poolConfig
+          ? new Pool(options.poolConfig as PoolConfig)
+          : (() => { throw new Error("Either pool or poolConfig must be provided"); })();
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createChannelBindingsTable);
@@ -68,6 +80,12 @@ export class ChannelBindingStore implements BindingRegistry {
     return this.initPromise;
   }
 
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
   async resolve(params: {
     channel: string;
     senderId: string;

package/src/stores/ChannelIdentityMappingStore.ts

@@ -4,7 +4,8 @@ import { MigrationManager } from "../migrations/migration";
 import { createChannelIdentityMappingTables } from "../migrations/channel_identity_mapping_migration";
 
 export interface ChannelIdentityMappingStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
@@ -60,15 +61,26 @@ export interface ChannelIdentityMapping {
 
 export class ChannelIdentityMappingStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized = false;
+  private ownsPool: boolean = true;
   private initPromise: Promise<void> | null = null;
 
   constructor(options: ChannelIdentityMappingStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     this.pool =
       typeof options.poolConfig === "string"
         ? new Pool({ connectionString: options.poolConfig })
-        : new Pool(options.poolConfig);
+        : options.poolConfig
+          ? new Pool(options.poolConfig as PoolConfig)
+          : (() => { throw new Error("Either pool or poolConfig must be provided"); })();
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createChannelIdentityMappingTables);
@@ -102,6 +114,12 @@ export class ChannelIdentityMappingStore {
     return this.initPromise;
   }
 
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
   async createMapping(
     input: CreateChannelIdentityMappingInput,
   ): Promise<ChannelIdentityMapping> {

package/src/stores/MenuStore.ts

@@ -12,7 +12,8 @@ import { createMenuItemsTable } from "../migrations/menu_items_migration";
 import { addFileContentType } from "../migrations/menu_items_add_file_type";
 
 export interface MenuStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
@@ -33,15 +34,26 @@ type MenuItemRow = {
 
 export class MenuStore implements MenuRegistry {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized = false;
+  private ownsPool: boolean = true;
   private initPromise: Promise<void> | null = null;
 
   constructor(options: MenuStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     this.pool =
       typeof options.poolConfig === "string"
         ? new Pool({ connectionString: options.poolConfig })
-        : new Pool(options.poolConfig);
+        : options.poolConfig
+          ? new Pool(options.poolConfig as PoolConfig)
+          : (() => { throw new Error("Either pool or poolConfig must be provided"); })();
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createMenuItemsTable);
@@ -69,6 +81,12 @@ export class MenuStore implements MenuRegistry {
     return this.initPromise;
   }
 
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
   async list(params: { tenantId: string; menuTarget?: string }): Promise<MenuItem[]> {
     await this.ensureInitialized();
     const conditions: string[] = ["tenant_id = $1"];