@axiom-lattice/pg-stores 1.0.21 → 1.0.23

package/src/stores/PostgreSQLSkillStore.ts

@@ -1,5 +1,5 @@
 /**
- * PostgreSQL implementation of SkillStore
+ * PostgreSQL implementation of SkillStore with tenant isolation
  */
 
 import { Pool, PoolConfig } from "pg";
@@ -10,35 +10,28 @@ import {
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
 import { createSkillsTable } from "../migrations/skill_migrations";
+import { addSkillTenantId } from "../migrations/skill_tenant_migration";
+import { changeSkillPrimaryKey } from "../migrations/skill_pk_migration";
 
-/**
- * PostgreSQL SkillStore options
- */
 export interface PostgreSQLSkillStoreOptions {
-  /**
-   * PostgreSQL connection pool configuration
-   * Can be a connection string or PoolConfig object
-   */
   poolConfig: string | PoolConfig;
-
-  /**
-   * Whether to run migrations automatically on initialization
-   * @default true
-   */
   autoMigrate?: boolean;
 }
 
 /**
- * PostgreSQL implementation of SkillStore
+ * PostgreSQL implementation of SkillStore with tenant isolation
+ * 
+ * Features:
+ * - Multi-tenant isolation via tenant_id
  */
 export class PostgreSQLSkillStore implements SkillStore {
   private pool: Pool;
   private migrationManager: MigrationManager;
   private initialized: boolean = false;
   private ownsPool: boolean = true;
+  private initPromise: Promise<void> | null = null;
 
   constructor(options: PostgreSQLSkillStoreOptions) {
-    // Create Pool from config
     if (typeof options.poolConfig === "string") {
       this.pool = new Pool({ connectionString: options.poolConfig });
     } else {
@@ -47,8 +40,9 @@ export class PostgreSQLSkillStore implements SkillStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createSkillsTable);
+    this.migrationManager.register(addSkillTenantId);
+    this.migrationManager.register(changeSkillPrimaryKey);
 
-    // Auto-migrate by default
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
         console.error("Failed to initialize PostgreSQLSkillStore:", error);
@@ -57,31 +51,37 @@ export class PostgreSQLSkillStore implements SkillStore {
     }
   }
 
-  /**
-   * Dispose resources and close the connection pool
-   * Should be called when the store is no longer needed
-   */
-  async dispose(): Promise<void> {
-    if (this.ownsPool && this.pool) {
-      await this.pool.end();
-    }
-  }
-
   /**
    * 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;
+  }
+
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
   }
 
-  /**
-   * Ensure store is initialized
-   */
   private async ensureInitialized(): Promise<void> {
     if (!this.initialized) {
       await this.initialize();
@@ -93,6 +93,7 @@ export class PostgreSQLSkillStore implements SkillStore {
    */
   private mapRowToSkill(row: {
     id: string;
+    tenant_id: string;
     name: string;
     description: string;
     license: string | null;
@@ -105,28 +106,28 @@ export class PostgreSQLSkillStore implements SkillStore {
   }): Skill {
     return {
       id: row.id,
+      tenantId: row.tenant_id,
       name: row.name,
       description: row.description,
       license: row.license || undefined,
       compatibility: row.compatibility || undefined,
       metadata: row.metadata || {},
       content: row.content || undefined,
-      subSkills: Array.isArray(row.sub_skills) && row.sub_skills.length > 0 
-        ? row.sub_skills 
-        : undefined,
+      subSkills: Array.isArray(row.sub_skills) ? row.sub_skills : undefined,
       createdAt: row.created_at,
       updatedAt: row.updated_at,
     };
   }
 
   /**
-   * Get all skills
+   * Get all skills for a tenant
    */
-  async getAllSkills(): Promise<Skill[]> {
+  async getAllSkills(tenantId: string): Promise<Skill[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -138,23 +139,26 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
+      WHERE tenant_id = $1
       ORDER BY created_at DESC
-    `
+    `,
+      [tenantId]
     );
 
     return result.rows.map(this.mapRowToSkill);
   }
 
   /**
-   * Get skill by ID
+   * Get skill by ID for a tenant
    */
-  async getSkillById(id: string): Promise<Skill | null> {
+  async getSkillById(tenantId: string, id: string): Promise<Skill | null> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -166,11 +170,11 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
-      WHERE id = $1
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [id]
+      [tenantId, id]
     );
 
     if (result.rows.length === 0) {
@@ -181,19 +185,23 @@ export class PostgreSQLSkillStore implements SkillStore {
   }
 
   /**
-   * Create a new skill
+   * Create a new skill for a tenant
    */
-  async createSkill(id: string, data: CreateSkillRequest): Promise<Skill> {
+  async createSkill(
+    tenantId: string,
+    id: string,
+    data: CreateSkillRequest
+  ): Promise<Skill> {
     await this.ensureInitialized();
 
     const now = new Date();
-    const metadata = data.metadata || {};
 
     await this.pool.query(
       `
-      INSERT INTO lattice_skills (id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
+      INSERT INTO lattice_skills (id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
       ON CONFLICT (id) DO UPDATE SET
+        tenant_id = EXCLUDED.tenant_id,
         name = EXCLUDED.name,
         description = EXCLUDED.description,
         license = EXCLUDED.license,
@@ -205,11 +213,12 @@ export class PostgreSQLSkillStore implements SkillStore {
     `,
       [
         id,
+        tenantId,
         data.name,
         data.description,
         data.license || null,
         data.compatibility || null,
-        JSON.stringify(metadata),
+        JSON.stringify(data.metadata || {}),
         data.content || null,
         JSON.stringify(data.subSkills || []),
         now,
@@ -217,116 +226,143 @@ export class PostgreSQLSkillStore implements SkillStore {
       ]
     );
 
-    return this.getSkillById(id) as Promise<Skill>;
+    return {
+      id,
+      tenantId,
+      name: data.name,
+      description: data.description,
+      license: data.license,
+      compatibility: data.compatibility,
+      metadata: data.metadata || {},
+      content: data.content,
+      subSkills: data.subSkills,
+      createdAt: now,
+      updatedAt: now,
+    };
   }
 
   /**
-   * Update an existing skill
+   * Update an existing skill for a tenant
    */
   async updateSkill(
+    tenantId: string,
     id: string,
     updates: Partial<CreateSkillRequest>
   ): Promise<Skill | null> {
     await this.ensureInitialized();
 
+    // Get existing skill
+    const existing = await this.getSkillById(tenantId, id);
+    if (!existing) {
+      return null;
+    }
+
+    // Build update query dynamically based on provided fields
     const updateFields: string[] = [];
-    const values: any[] = [];
+    const updateValues: any[] = [];
     let paramIndex = 1;
 
     if (updates.name !== undefined) {
       updateFields.push(`name = $${paramIndex++}`);
-      values.push(updates.name);
+      updateValues.push(updates.name);
     }
 
     if (updates.description !== undefined) {
       updateFields.push(`description = $${paramIndex++}`);
-      values.push(updates.description);
+      updateValues.push(updates.description);
     }
 
     if (updates.license !== undefined) {
       updateFields.push(`license = $${paramIndex++}`);
-      values.push(updates.license || null);
+      updateValues.push(updates.license || null);
     }
 
     if (updates.compatibility !== undefined) {
       updateFields.push(`compatibility = $${paramIndex++}`);
-      values.push(updates.compatibility || null);
+      updateValues.push(updates.compatibility || null);
     }
 
     if (updates.metadata !== undefined) {
       updateFields.push(`metadata = $${paramIndex++}`);
-      values.push(JSON.stringify(updates.metadata || {}));
+      updateValues.push(JSON.stringify(updates.metadata || {}));
     }
 
     if (updates.content !== undefined) {
       updateFields.push(`content = $${paramIndex++}`);
-      values.push(updates.content || null);
+      updateValues.push(updates.content || null);
     }
 
     if (updates.subSkills !== undefined) {
       updateFields.push(`sub_skills = $${paramIndex++}`);
-      values.push(JSON.stringify(updates.subSkills || []));
+      updateValues.push(JSON.stringify(updates.subSkills || []));
     }
 
     if (updateFields.length === 0) {
-      return this.getSkillById(id);
+      // No fields to update
+      return existing;
     }
 
+    // Always update updated_at
     updateFields.push(`updated_at = $${paramIndex++}`);
-    values.push(new Date());
-    values.push(id);
+    updateValues.push(new Date());
+
+    // Add tenant_id and id for WHERE clause
+    updateValues.push(tenantId);
+    updateValues.push(id);
 
     await this.pool.query(
       `
       UPDATE lattice_skills
       SET ${updateFields.join(", ")}
-      WHERE id = $${paramIndex}
+      WHERE tenant_id = $${paramIndex++} AND id = $${paramIndex}
     `,
-      values
+      updateValues
     );
 
-    return this.getSkillById(id);
+    // Return updated skill
+    return await this.getSkillById(tenantId, id);
   }
 
   /**
-   * Delete a skill by ID
+   * Delete a skill by ID for a tenant
    */
-  async deleteSkill(id: string): Promise<boolean> {
+  async deleteSkill(tenantId: string, id: string): Promise<boolean> {
     await this.ensureInitialized();
 
     const result = await this.pool.query(
       `
       DELETE FROM lattice_skills
-      WHERE id = $1
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [id]
+      [tenantId, id]
     );
 
     return result.rowCount !== null && result.rowCount > 0;
   }
 
   /**
-   * Check if skill exists
+   * Check if skill exists for a tenant
    */
-  async hasSkill(id: string): Promise<boolean> {
+  async hasSkill(tenantId: string, id: string): Promise<boolean> {
     await this.ensureInitialized();
 
-    const result = await this.pool.query<{ count: string }>(
+    const result = await this.pool.query(
       `
-      SELECT COUNT(*) as count
-      FROM lattice_skills
-      WHERE id = $1
+      SELECT 1 FROM lattice_skills
+      WHERE tenant_id = $1 AND id = $2
+      LIMIT 1
     `,
-      [id]
+      [tenantId, id]
     );
 
-    return parseInt(result.rows[0].count, 10) > 0;
+    return result.rows.length > 0;
   }
 
   /**
-   * Search skills by metadata
+   * Search skills by metadata within a tenant
    */
   async searchByMetadata(
+    tenantId: string,
     metadataKey: string,
     metadataValue: string
   ): Promise<Skill[]> {
@@ -334,6 +370,7 @@ export class PostgreSQLSkillStore implements SkillStore {
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -345,25 +382,29 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
-      WHERE metadata->>$1 = $2
+      WHERE tenant_id = $1 AND metadata ->> $2 = $3
       ORDER BY created_at DESC
     `,
-      [metadataKey, metadataValue]
+      [tenantId, metadataKey, metadataValue]
     );
 
     return result.rows.map(this.mapRowToSkill);
   }
 
   /**
-   * Filter skills by compatibility
+   * Filter skills by compatibility within a tenant
    */
-  async filterByCompatibility(compatibility: string): Promise<Skill[]> {
+  async filterByCompatibility(
+    tenantId: string,
+    compatibility: string
+  ): Promise<Skill[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -375,25 +416,29 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
-      WHERE compatibility = $1
+      WHERE tenant_id = $1 AND compatibility = $2
       ORDER BY created_at DESC
     `,
-      [compatibility]
+      [tenantId, compatibility]
     );
 
     return result.rows.map(this.mapRowToSkill);
   }
 
   /**
-   * Filter skills by license
+   * Filter skills by license within a tenant
    */
-  async filterByLicense(license: string): Promise<Skill[]> {
+  async filterByLicense(
+    tenantId: string,
+    license: string
+  ): Promise<Skill[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -405,32 +450,30 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
-      WHERE license = $1
+      WHERE tenant_id = $1 AND license = $2
       ORDER BY created_at DESC
     `,
-      [license]
+      [tenantId, license]
     );
 
     return result.rows.map(this.mapRowToSkill);
   }
 
   /**
-   * Get sub-skills of a parent skill
+   * Get sub-skills of a parent skill within a tenant
+   * Note: This searches for skills that have the parent skill name in their subSkills array
    */
-  async getSubSkills(parentSkillName: string): Promise<Skill[]> {
+  async getSubSkills(
+    tenantId: string,
+    parentSkillName: string
+  ): Promise<Skill[]> {
     await this.ensureInitialized();
 
-    const parentSkill = await this.getSkillById(parentSkillName);
-    if (!parentSkill || !parentSkill.subSkills || parentSkill.subSkills.length === 0) {
-      return [];
-    }
-
-    // Query all sub-skills in one query
-    const placeholders = parentSkill.subSkills.map((_, index) => `$${index + 1}`).join(", ");
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       name: string;
       description: string;
       license: string | null;
@@ -442,12 +485,12 @@ export class PostgreSQLSkillStore implements SkillStore {
       updated_at: Date;
     }>(
       `
-      SELECT id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
+      SELECT id, tenant_id, name, description, license, compatibility, metadata, content, sub_skills, created_at, updated_at
       FROM lattice_skills
-      WHERE name IN (${placeholders})
+      WHERE tenant_id = $1 AND sub_skills @> to_jsonb($2::text)
       ORDER BY created_at DESC
     `,
-      parentSkill.subSkills
+      [tenantId, JSON.stringify([parentSkillName])]
     );
 
     return result.rows.map(this.mapRowToSkill);