@axiom-lattice/pg-stores 1.0.83 → 1.0.85

package/src/stores/PostgreSQLUserStore.ts

@@ -6,20 +6,32 @@ import { createUsersTable } from "../migrations/user_migrations";
 import { addUserStatusColumn } from "../migrations/user_status_migration";
 
 export interface PostgreSQLUserStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
 export class PostgreSQLUserStore implements UserStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized: boolean = false;
+  private ownsPool: boolean = true;
 
   constructor(options: PostgreSQLUserStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
     } else {
-      this.pool = new Pool(options.poolConfig);
+      throw new Error("Either pool or poolConfig must be provided");
     }
 
     this.migrationManager = new MigrationManager(this.pool);
@@ -35,7 +47,9 @@ export class PostgreSQLUserStore implements UserStore {
   }
 
   async dispose(): Promise<void> {
-    await this.pool.end();
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
   }
 
   async initialize(): Promise<void> {

package/src/stores/PostgreSQLUserTenantLinkStore.ts

@@ -11,20 +11,32 @@ import { MigrationManager } from "../migrations/migration";
 import { createUserTenantLinkTable, migrateRemoveTenantIdFromUsers } from "../migrations/user_tenant_link_migrations";
 
 export interface PostgreSQLUserTenantLinkStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
 export class PostgreSQLUserTenantLinkStore implements UserTenantLinkStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized: boolean = false;
+  private ownsPool: boolean = true;
 
   constructor(options: PostgreSQLUserTenantLinkStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
     } else {
-      this.pool = new Pool(options.poolConfig);
+      throw new Error("Either pool or poolConfig must be provided");
     }
 
     this.migrationManager = new MigrationManager(this.pool);
@@ -40,7 +52,9 @@ export class PostgreSQLUserTenantLinkStore implements UserTenantLinkStore {
   }
 
   async dispose(): Promise<void> {
-    await this.pool.end();
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
   }
 
   async initialize(): Promise<void> {

package/src/stores/PostgreSQLWorkflowTrackingStore.ts

@@ -19,21 +19,33 @@ import { MigrationManager } from "../migrations/migration";
 import { createWorkflowTrackingTables, addStepThreadId } from "../migrations/workflow_tracking_migrations";
 
 export interface PostgreSQLWorkflowTrackingStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
 export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized: boolean = false;
+  private ownsPool: boolean = true;
   private initPromise: Promise<void> | null = null;
 
   constructor(options: PostgreSQLWorkflowTrackingStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
     } else {
-      this.pool = new Pool(options.poolConfig);
+      throw new Error("Either pool or poolConfig must be provided");
     }
 
     this.migrationManager = new MigrationManager(this.pool);
@@ -49,7 +61,9 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
   }
 
   async dispose(): Promise<void> {
-    await this.pool.end();
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
   }
 
   async initialize(): Promise<void> {

package/src/stores/PostgreSQLWorkspaceStore.ts

@@ -17,11 +17,17 @@ import { createWorkspacesTable } from "../migrations/workspace_migrations";
  * PostgreSQL WorkspaceStore options
  */
 export interface PostgreSQLWorkspaceStoreOptions {
+  /**
+   * External pool instance. When provided, poolConfig is not required and the
+   * caller is responsible for pool lifecycle and migrations.
+   */
+  pool?: Pool;
+
   /**
    * PostgreSQL connection pool configuration
    * Can be a connection string or PoolConfig object
    */
-  poolConfig: string | PoolConfig;
+  poolConfig?: string | PoolConfig;
 
   /**
    * Whether to run migrations automatically on initialization
@@ -35,16 +41,26 @@ export interface PostgreSQLWorkspaceStoreOptions {
  */
 export class PostgreSQLWorkspaceStore implements WorkspaceStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized: boolean = false;
   private ownsPool: boolean = true;
 
   constructor(options: PostgreSQLWorkspaceStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     // Create Pool from config
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
     } else {
-      this.pool = new Pool(options.poolConfig);
+      throw new Error("Either pool or poolConfig must be provided");
     }
 
     this.migrationManager = new MigrationManager(this.pool);

package/src/stores/PostgresSharedResourceStore.ts

@@ -6,7 +6,8 @@ import { MigrationManager } from "../migrations/migration";
 import { createSharedResourcesTable } from "../migrations/shared_resources_migration";
 
 export interface PostgresSharedResourceStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
@@ -32,15 +33,26 @@ type ShareRow = {
 
 export class PostgresSharedResourceStore implements SharedResourceStore {
   private pool: Pool;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
   private initialized = false;
+  private ownsPool: boolean = true;
   private initPromise: Promise<void> | null = null;
 
   constructor(options: PostgresSharedResourceStoreOptions) {
+    // 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(createSharedResourcesTable);
@@ -70,6 +82,12 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
     return this.initPromise;
   }
 
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
   async findByToken(token: string): Promise<ShareRecord | null> {
     await this.ensureInitialized();
     const { rows } = await this.pool.query<ShareRow>(

package/src/stores/ThreadMessageQueueStore.ts

@@ -22,7 +22,8 @@ import { addWorkspaceProjectToQueue } from "../migrations/add_workspace_project_
 export type { PendingMessage, AddMessageParams, ThreadInfo };
 
 export interface PostgreSQLThreadMessageQueueStoreOptions {
-  poolConfig: string | PoolConfig;
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
   autoMigrate?: boolean;
 }
 
@@ -31,14 +32,24 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   private ownsPool: boolean = true;
   private initialized: boolean = false;
   private initPromise: Promise<void> | null = null;
-  private migrationManager: MigrationManager;
+  private migrationManager!: MigrationManager;
 
   constructor(options: PostgreSQLThreadMessageQueueStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
     // Create Pool from config
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
     } else {
-      this.pool = new Pool(options.poolConfig);
+      throw new Error("Either pool or poolConfig must be provided");
     }
 
     this.migrationManager = new MigrationManager(this.pool);