@axiom-lattice/pg-stores 1.0.21 → 1.0.23

package/src/stores/PostgreSQLAssistantStore.ts

@@ -10,6 +10,8 @@ import {
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
 import { createAssistantsTable } from "../migrations/assistant_migrations";
+import { addAssistantTenantId } from "../migrations/assistant_tenant_migration";
+import { changeAssistantPrimaryKey } from "../migrations/assistant_pk_migration";
 
 /**
  * PostgreSQL AssistantStore options
@@ -30,12 +32,16 @@ export interface PostgreSQLAssistantStoreOptions {
 
 /**
  * PostgreSQL implementation of AssistantStore
+ * 
+ * Features:
+ * - Multi-tenant isolation via tenant_id
  */
 export class PostgreSQLAssistantStore implements AssistantStore {
   private pool: Pool;
   private migrationManager: MigrationManager;
   private initialized: boolean = false;
   private ownsPool: boolean = true;
+  private initPromise: Promise<void> | null = null;
 
   constructor(options: PostgreSQLAssistantStoreOptions) {
     // Create Pool from config
@@ -47,6 +53,8 @@ export class PostgreSQLAssistantStore implements AssistantStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createAssistantsTable);
+    this.migrationManager.register(addAssistantTenantId);
+    this.migrationManager.register(changeAssistantPrimaryKey);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -59,24 +67,38 @@ export class PostgreSQLAssistantStore implements AssistantStore {
 
   /**
    * Initialize the store and run migrations
+   * Uses a promise-based lock to prevent concurrent initialization
    */
   async initialize(): Promise<void> {
     if (this.initialized) {
       return;
     }
 
-    await this.migrationManager.migrate();
-    this.initialized = true;
+    if (this.initPromise) {
+      return this.initPromise;
+    }
+
+    this.initPromise = (async () => {
+      try {
+        await this.migrationManager.migrate();
+        this.initialized = true;
+      } finally {
+        this.initPromise = null;
+      }
+    })();
+
+    return this.initPromise;
   }
 
   /**
-   * Get all assistants
+   * Get all assistants for a tenant
    */
-  async getAllAssistants(): Promise<Assistant[]> {
+  async getAllAssistants(tenantId: string): Promise<Assistant[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string | null;
       graph_definition: any;
@@ -84,10 +106,12 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, graph_definition, created_at, updated_at
+      SELECT id, tenant_id, name, description, graph_definition, created_at, updated_at
       FROM lattice_assistants
+      WHERE tenant_id = $1
       ORDER BY created_at DESC
-    `
+    `,
+      [tenantId]
     );
 
     return result.rows.map(this.mapRowToAssistant);
@@ -96,11 +120,12 @@ export class PostgreSQLAssistantStore implements AssistantStore {
   /**
    * Get assistant by ID
    */
-  async getAssistantById(id: string): Promise<Assistant | null> {
+  async getAssistantById(tenantId: string, id: string): Promise<Assistant | null> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string | null;
       graph_definition: any;
@@ -108,11 +133,11 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, graph_definition, created_at, updated_at
+      SELECT id, tenant_id, name, description, graph_definition, created_at, updated_at
       FROM lattice_assistants
-      WHERE id = $1
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [id]
+      [tenantId, id]
     );
 
     if (result.rows.length === 0) {
@@ -126,6 +151,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
    * Create a new assistant
    */
   async createAssistant(
+    tenantId: string,
     id: string,
     data: CreateAssistantRequest
   ): Promise<Assistant> {
@@ -135,9 +161,10 @@ export class PostgreSQLAssistantStore implements AssistantStore {
 
     await this.pool.query(
       `
-      INSERT INTO lattice_assistants (id, name, description, graph_definition, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5, $6)
+      INSERT INTO lattice_assistants (id, tenant_id, name, description, graph_definition, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7)
       ON CONFLICT (id) DO UPDATE SET
+        tenant_id = EXCLUDED.tenant_id,
         name = EXCLUDED.name,
         description = EXCLUDED.description,
         graph_definition = EXCLUDED.graph_definition,
@@ -145,6 +172,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
     `,
       [
         id,
+        tenantId,
         data.name,
         data.description || null,
         JSON.stringify(data.graphDefinition),
@@ -155,6 +183,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
 
     return {
       id,
+      tenantId,
       name: data.name,
       description: data.description,
       graphDefinition: data.graphDefinition,
@@ -167,13 +196,14 @@ export class PostgreSQLAssistantStore implements AssistantStore {
    * Update an existing assistant
    */
   async updateAssistant(
+    tenantId: string,
     id: string,
     updates: Partial<CreateAssistantRequest>
   ): Promise<Assistant | null> {
     await this.ensureInitialized();
 
     // Get existing assistant
-    const existing = await this.getAssistantById(id);
+    const existing = await this.getAssistantById(tenantId, id);
     if (!existing) {
       return null;
     }
@@ -207,34 +237,35 @@ export class PostgreSQLAssistantStore implements AssistantStore {
     updateFields.push(`updated_at = $${paramIndex++}`);
     updateValues.push(new Date());
 
-    // Add id for WHERE clause
+    // Add tenant_id and id for WHERE clause
+    updateValues.push(tenantId);
     updateValues.push(id);
 
     await this.pool.query(
       `
       UPDATE lattice_assistants
       SET ${updateFields.join(", ")}
-      WHERE id = $${paramIndex}
+      WHERE tenant_id = $${paramIndex++} AND id = $${paramIndex}
     `,
       updateValues
     );
 
     // Return updated assistant
-    return await this.getAssistantById(id);
+    return await this.getAssistantById(tenantId, id);
   }
 
   /**
    * Delete an assistant by ID
    */
-  async deleteAssistant(id: string): Promise<boolean> {
+  async deleteAssistant(tenantId: string, id: string): Promise<boolean> {
     await this.ensureInitialized();
 
     const result = await this.pool.query(
       `
       DELETE FROM lattice_assistants
-      WHERE id = $1
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [id]
+      [tenantId, id]
     );
 
     return result.rowCount !== null && result.rowCount > 0;
@@ -243,16 +274,16 @@ export class PostgreSQLAssistantStore implements AssistantStore {
   /**
    * Check if assistant exists
    */
-  async hasAssistant(id: string): Promise<boolean> {
+  async hasAssistant(tenantId: string, id: string): Promise<boolean> {
     await this.ensureInitialized();
 
     const result = await this.pool.query(
       `
       SELECT 1 FROM lattice_assistants
-      WHERE id = $1
+      WHERE tenant_id = $1 AND id = $2
       LIMIT 1
     `,
-      [id]
+      [tenantId, id]
     );
 
     return result.rows.length > 0;
@@ -282,6 +313,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
    */
   private mapRowToAssistant(row: {
     id: string;
+    tenant_id: string;
     name: string;
     description: string | null;
     graph_definition: any;
@@ -290,6 +322,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
   }): Assistant {
     return {
       id: row.id,
+      tenantId: row.tenant_id,
       name: row.name,
       description: row.description || undefined,
       graphDefinition:

package/src/stores/PostgreSQLScheduleStorage.ts

@@ -14,6 +14,7 @@ import {
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
 import { createScheduledTasksTable } from "../migrations/schedule_migrations";
+import { addScheduleTenantId } from "../migrations/schedule_tenant_migration";
 
 /**
  * PostgreSQL ScheduleStorage options
@@ -78,6 +79,7 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createScheduledTasksTable);
+    this.migrationManager.register(addScheduleTenantId);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -424,6 +426,7 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
    * Get all tasks with optional filters
    */
   async getAllTasks(filters?: {
+    tenantId?: string;
     status?: ScheduledTaskStatus;
     executionType?: ScheduleExecutionType;
     taskType?: string;
@@ -438,6 +441,12 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
     const values: any[] = [];
     let paramIndex = 1;
 
+    // Always filter by tenantId if provided
+    if (filters?.tenantId !== undefined) {
+      whereClauses.push(`tenant_id = $${paramIndex++}`);
+      values.push(filters.tenantId);
+    }
+
     if (filters?.status !== undefined) {
       whereClauses.push(`status = $${paramIndex++}`);
       values.push(filters.status);
@@ -483,6 +492,7 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
    * Count tasks with optional filters
    */
   async countTasks(filters?: {
+    tenantId?: string;
     status?: ScheduledTaskStatus;
     executionType?: ScheduleExecutionType;
     taskType?: string;
@@ -495,6 +505,12 @@ export class PostgreSQLScheduleStorage implements ScheduleStorage {
     const values: any[] = [];
     let paramIndex = 1;
 
+    // Always filter by tenantId if provided
+    if (filters?.tenantId !== undefined) {
+      whereClauses.push(`tenant_id = $${paramIndex++}`);
+      values.push(filters.tenantId);
+    }
+
     if (filters?.status !== undefined) {
       whereClauses.push(`status = $${paramIndex++}`);
       values.push(filters.status);