@kiyeonjeon21/datacontext 0.3.2 → 0.4.0

package/src/knowledge/store.ts

@@ -1,6 +1,44 @@
 /**
  * Knowledge Store
- * Local JSON-based storage for table descriptions, query examples, and business rules
+ * 
+ * Local JSON-based storage for the Knowledge Graph, including:
+ * - Table descriptions (nodes)
+ * - Table relationships (edges)
+ * - Query examples, business rules, and terms (annotations)
+ * 
+ * ## Architecture
+ * 
+ * The Knowledge Store implements a lightweight graph model:
+ * 
+ * ```
+ * ┌─────────────────────────────────────────────────────────────────┐
+ * │                     Knowledge Store                             │
+ * │                                                                 │
+ * │   Nodes:                                                        │
+ * │   - TableDescription (with embedded columns)                    │
+ * │                                                                 │
+ * │   Edges:                                                        │
+ * │   - TableRelationship (from → to with join condition)           │
+ * │                                                                 │
+ * │   Annotations:                                                  │
+ * │   - QueryExample, BusinessRule, BusinessTerm                    │
+ * │                                                                 │
+ * └─────────────────────────────────────────────────────────────────┘
+ * ```
+ * 
+ * ## Graph Traversal
+ * 
+ * Provides methods for:
+ * - Finding related tables (1-hop)
+ * - Finding shortest join path (multi-hop BFS)
+ * - Detecting relationship types
+ * 
+ * ## Future: SQLite Migration
+ * 
+ * For larger graphs, consider migrating to SQLite embedded storage
+ * for better indexing and query performance.
+ * 
+ * @module knowledge/store
  */
 
 import { readFile, writeFile, mkdir } from 'fs/promises';
@@ -10,16 +48,27 @@ import { homedir } from 'os';
 import type {
   KnowledgeData,
   TableDescription,
+  TableRelationship,
+  TableRef,
   QueryExample,
   BusinessRule,
   BusinessTerm,
   KnowledgeEntry,
   ColumnDescription,
   TermCategory,
+  RelationshipType,
+  RelationshipCardinality,
+} from './types.js';
+import { 
+  generateId, 
+  createKnowledgeMeta, 
+  createTableRef,
+  tableRefToString,
+  tableRefsEqual,
 } from './types.js';
-import { generateId, createKnowledgeMeta } from './types.js';
 
-const KNOWLEDGE_VERSION = '1.0.0';
+/** Current knowledge format version */
+const KNOWLEDGE_VERSION = '2.0.0';
 
 export class KnowledgeStore {
   private dataPath: string;
@@ -42,7 +91,9 @@ export class KnowledgeStore {
   }
 
   /**
-   * Load knowledge data from disk
+   * Load knowledge data from disk.
+   * 
+   * Includes automatic migration for older data formats.
    */
   async load(): Promise<void> {
     if (existsSync(this.dataPath)) {
@@ -50,10 +101,18 @@ export class KnowledgeStore {
         const content = await readFile(this.dataPath, 'utf-8');
         const parsed = JSON.parse(content) as KnowledgeData;
         
-        // Migration: ensure businessTerms array exists (for old data files)
+        // Migration: ensure arrays exist (for old data files)
         if (!parsed.businessTerms) {
           parsed.businessTerms = [];
         }
+        if (!parsed.tableRelationships) {
+          parsed.tableRelationships = [];
+        }
+        
+        // Update version if migrating
+        if (parsed.version !== KNOWLEDGE_VERSION) {
+          parsed.version = KNOWLEDGE_VERSION;
+        }
         
         this.data = parsed;
       } catch (error) {
@@ -224,6 +283,345 @@ export class KnowledgeStore {
     await this.save();
   }
 
+  // ============ Table Relationships (Graph Edges) ============
+
+  /**
+   * Get all table relationships.
+   * 
+   * @returns All relationships in the knowledge store
+   */
+  getRelationships(): TableRelationship[] {
+    return this.data?.tableRelationships ?? [];
+  }
+
+  /**
+   * Get tables related to a given table (1-hop traversal).
+   * 
+   * Returns all relationships where the given table is either
+   * the source (from) or target (to).
+   * 
+   * @param tableName - Table name to find relationships for
+   * @param schema - Schema name (defaults to 'public')
+   * @returns Array of relationships involving this table
+   * 
+   * @example
+   * ```typescript
+   * const related = store.getRelatedTables('orders', 'public');
+   * // Returns relationships like:
+   * // - orders → users (FK: orders.user_id → users.id)
+   * // - orders → products (FK: orders.product_id → products.id)
+   * // - order_items → orders (FK: order_items.order_id → orders.id)
+   * ```
+   */
+  getRelatedTables(tableName: string, schema: string = 'public'): TableRelationship[] {
+    return this.data?.tableRelationships.filter(
+      rel => (rel.from.table === tableName && rel.from.schema === schema) ||
+             (rel.to.table === tableName && rel.to.schema === schema)
+    ) ?? [];
+  }
+
+  /**
+   * Get outgoing relationships from a table (where table is the source).
+   * 
+   * @param tableName - Source table name
+   * @param schema - Schema name
+   * @returns Relationships where this table is the 'from' side
+   */
+  getOutgoingRelationships(tableName: string, schema: string = 'public'): TableRelationship[] {
+    return this.data?.tableRelationships.filter(
+      rel => rel.from.table === tableName && rel.from.schema === schema
+    ) ?? [];
+  }
+
+  /**
+   * Get incoming relationships to a table (where table is the target).
+   * 
+   * @param tableName - Target table name
+   * @param schema - Schema name
+   * @returns Relationships where this table is the 'to' side
+   */
+  getIncomingRelationships(tableName: string, schema: string = 'public'): TableRelationship[] {
+    return this.data?.tableRelationships.filter(
+      rel => rel.to.table === tableName && rel.to.schema === schema
+    ) ?? [];
+  }
+
+  /**
+   * Find the shortest join path between two tables.
+   * 
+   * Uses BFS (Breadth-First Search) to find the shortest path
+   * through the relationship graph.
+   * 
+   * @param fromTable - Starting table name
+   * @param toTable - Destination table name
+   * @param schema - Schema name (defaults to 'public')
+   * @returns Array of relationships forming the path, or empty if no path exists
+   * 
+   * @example
+   * ```typescript
+   * const path = store.findJoinPath('order_items', 'customers');
+   * // Returns path like:
+   * // [
+   * //   { from: 'order_items', to: 'orders', joinCondition: '...' },
+   * //   { from: 'orders', to: 'customers', joinCondition: '...' }
+   * // ]
+   * ```
+   */
+  findJoinPath(
+    fromTable: string, 
+    toTable: string, 
+    schema: string = 'public'
+  ): TableRelationship[] {
+    if (fromTable === toTable) return [];
+
+    const visited = new Set<string>();
+    const queue: Array<{ table: string; path: TableRelationship[] }> = [
+      { table: fromTable, path: [] }
+    ];
+
+    while (queue.length > 0) {
+      const { table, path } = queue.shift()!;
+      const tableKey = `${schema}.${table}`;
+
+      if (visited.has(tableKey)) continue;
+      visited.add(tableKey);
+
+      // Get all relationships for current table
+      const relationships = this.getRelatedTables(table, schema);
+
+      for (const rel of relationships) {
+        // Determine the next table in the path
+        const isFromCurrent = rel.from.table === table && rel.from.schema === schema;
+        const nextTable = isFromCurrent ? rel.to.table : rel.from.table;
+        const nextSchema = isFromCurrent ? rel.to.schema : rel.from.schema;
+
+        // Skip if we've already visited
+        if (visited.has(`${nextSchema}.${nextTable}`)) continue;
+
+        const newPath = [...path, rel];
+
+        // Found the destination
+        if (nextTable === toTable && nextSchema === schema) {
+          return newPath;
+        }
+
+        // Add to queue for further exploration
+        queue.push({ table: nextTable, path: newPath });
+      }
+    }
+
+    // No path found
+    return [];
+  }
+
+  /**
+   * Get the relationship between two specific tables (if it exists).
+   * 
+   * @param fromTable - Source table
+   * @param toTable - Target table
+   * @param schema - Schema name
+   * @returns Relationship if found, undefined otherwise
+   */
+  getRelationshipBetween(
+    fromTable: string,
+    toTable: string,
+    schema: string = 'public'
+  ): TableRelationship | undefined {
+    return this.data?.tableRelationships.find(
+      rel => 
+        (rel.from.table === fromTable && rel.to.table === toTable &&
+         rel.from.schema === schema && rel.to.schema === schema) ||
+        (rel.from.table === toTable && rel.to.table === fromTable &&
+         rel.from.schema === schema && rel.to.schema === schema)
+    );
+  }
+
+  /**
+   * Add a table relationship.
+   * 
+   * @param relationship - Complete relationship object
+   * @returns The added relationship
+   */
+  async addRelationship(relationship: TableRelationship): Promise<TableRelationship> {
+    if (!this.data) {
+      throw new Error('Knowledge store not loaded');
+    }
+
+    // Initialize if needed (migration)
+    if (!this.data.tableRelationships) {
+      this.data.tableRelationships = [];
+    }
+
+    // Check for duplicates (same from/to tables)
+    const existingIndex = this.data.tableRelationships.findIndex(
+      r => tableRefsEqual(r.from, relationship.from) && 
+           tableRefsEqual(r.to, relationship.to)
+    );
+
+    if (existingIndex >= 0) {
+      // Update existing
+      this.data.tableRelationships[existingIndex] = {
+        ...this.data.tableRelationships[existingIndex],
+        ...relationship,
+        updatedAt: new Date().toISOString(),
+      };
+      await this.save();
+      return this.data.tableRelationships[existingIndex];
+    }
+
+    // Add new
+    this.data.tableRelationships.push(relationship);
+    await this.save();
+    return relationship;
+  }
+
+  /**
+   * Add a relationship from FK information (convenience method).
+   * 
+   * @param params - FK parameters
+   * @returns The added relationship
+   * 
+   * @example
+   * ```typescript
+   * await store.addRelationshipFromFK({
+   *   fromTable: 'orders',
+   *   fromColumn: 'user_id',
+   *   toTable: 'users',
+   *   toColumn: 'id',
+   * });
+   * ```
+   */
+  async addRelationshipFromFK(params: {
+    fromTable: string;
+    fromSchema?: string;
+    fromColumn: string;
+    toTable: string;
+    toSchema?: string;
+    toColumn: string;
+    constraintName?: string;
+  }): Promise<TableRelationship> {
+    if (!this.data) {
+      throw new Error('Knowledge store not loaded');
+    }
+
+    const fromSchema = params.fromSchema ?? 'public';
+    const toSchema = params.toSchema ?? 'public';
+
+    const relationship: TableRelationship = {
+      ...createKnowledgeMeta('auto', this.data.schemaHash),
+      type: 'table_relationship',
+      from: createTableRef(params.fromTable, fromSchema),
+      to: createTableRef(params.toTable, toSchema),
+      relationshipType: 'foreign_key',
+      joinCondition: `${params.fromTable}.${params.fromColumn} = ${params.toTable}.${params.toColumn}`,
+      cardinality: 'many-to-one',
+      fromColumns: [params.fromColumn],
+      toColumns: [params.toColumn],
+      isPreferred: true,
+      constraintName: params.constraintName,
+    };
+
+    return this.addRelationship(relationship);
+  }
+
+  /**
+   * Add multiple relationships at once (batch operation).
+   * 
+   * @param relationships - Array of relationships to add
+   * @returns Number of relationships added
+   */
+  async addRelationships(relationships: TableRelationship[]): Promise<number> {
+    if (!this.data) {
+      throw new Error('Knowledge store not loaded');
+    }
+
+    if (!this.data.tableRelationships) {
+      this.data.tableRelationships = [];
+    }
+
+    let addedCount = 0;
+
+    for (const rel of relationships) {
+      const existingIndex = this.data.tableRelationships.findIndex(
+        r => tableRefsEqual(r.from, rel.from) && tableRefsEqual(r.to, rel.to)
+      );
+
+      if (existingIndex >= 0) {
+        // Update existing
+        this.data.tableRelationships[existingIndex] = {
+          ...this.data.tableRelationships[existingIndex],
+          ...rel,
+          updatedAt: new Date().toISOString(),
+        };
+      } else {
+        // Add new
+        this.data.tableRelationships.push(rel);
+        addedCount++;
+      }
+    }
+
+    await this.save();
+    return addedCount;
+  }
+
+  /**
+   * Delete a relationship by ID.
+   * 
+   * @param id - Relationship ID to delete
+   */
+  async deleteRelationship(id: string): Promise<void> {
+    if (!this.data) {
+      throw new Error('Knowledge store not loaded');
+    }
+
+    const index = this.data.tableRelationships.findIndex(r => r.id === id);
+    if (index >= 0) {
+      this.data.tableRelationships.splice(index, 1);
+      await this.save();
+    }
+  }
+
+  /**
+   * Get a summary of the relationship graph.
+   * 
+   * @returns Graph statistics
+   */
+  getGraphSummary(): {
+    nodeCount: number;
+    edgeCount: number;
+    tablesWithRelationships: number;
+    isolatedTables: number;
+    relationshipsByType: Record<RelationshipType, number>;
+  } {
+    const relationships = this.data?.tableRelationships ?? [];
+    const tables = this.data?.tableDescriptions ?? [];
+    
+    // Find unique tables in relationships
+    const tablesInRelationships = new Set<string>();
+    for (const rel of relationships) {
+      tablesInRelationships.add(tableRefToString(rel.from));
+      tablesInRelationships.add(tableRefToString(rel.to));
+    }
+
+    // Count by type
+    const byType: Record<RelationshipType, number> = {
+      'foreign_key': 0,
+      'implicit_join': 0,
+      'manual': 0,
+    };
+    for (const rel of relationships) {
+      byType[rel.relationshipType]++;
+    }
+
+    return {
+      nodeCount: tables.length,
+      edgeCount: relationships.length,
+      tablesWithRelationships: tablesInRelationships.size,
+      isolatedTables: tables.length - Math.min(tablesInRelationships.size, tables.length),
+      relationshipsByType: byType,
+    };
+  }
+
   // ============ Query Examples ============
 
   /**
@@ -542,6 +940,12 @@ export class KnowledgeStore {
       }
     }
 
+    for (const rel of this.data.tableRelationships ?? []) {
+      if (rel.schemaHash !== currentSchemaHash) {
+        outdated.push(rel);
+      }
+    }
+
     return outdated;
   }
 
@@ -656,6 +1060,7 @@ export class KnowledgeStore {
       schemaHash: '',
       lastSyncAt: new Date().toISOString(),
       tableDescriptions: [],
+      tableRelationships: [],
       queryExamples: [],
       businessRules: [],
       businessTerms: [],
@@ -672,12 +1077,19 @@ export class KnowledgeStore {
       description?: string;
       columns: Array<{ name: string; description?: string }>;
     }>;
+    relationships: TableRelationship[];
     queryExamples: QueryExample[];
     businessRules: BusinessRule[];
     businessTerms: BusinessTerm[];
   } {
     if (!this.data) {
-      return { tables: [], queryExamples: [], businessRules: [], businessTerms: [] };
+      return { 
+        tables: [], 
+        relationships: [],
+        queryExamples: [], 
+        businessRules: [], 
+        businessTerms: [] 
+      };
     }
 
     return {
@@ -687,11 +1099,73 @@ export class KnowledgeStore {
         description: td.description,
         columns: td.columns,
       })),
+      relationships: this.data.tableRelationships ?? [],
       queryExamples: this.data.queryExamples,
       businessRules: this.data.businessRules,
       businessTerms: this.data.businessTerms ?? [],
     };
   }
+
+  /**
+   * Build context including relationships for AI.
+   * 
+   * Enhanced version that includes table relationship information
+   * to help AI understand how tables connect.
+   * 
+   * @param tables - Tables to include
+   * @param userQuery - Optional user query for term matching
+   * @param schema - Optional schema (if not provided, searches all schemas)
+   * @returns Formatted context string
+   */
+  buildContextWithRelationships(tables: string[], userQuery?: string, schema?: string): string {
+    let context = this.buildContext(tables, userQuery);
+
+    // Add relationship information
+    // If no schema specified, find relationships across all schemas
+    const relevantRelationships: TableRelationship[] = [];
+    for (const table of tables) {
+      // Search relationships that involve this table (in any schema if not specified)
+      const allRelationships = this.data?.tableRelationships ?? [];
+      const rels = schema 
+        ? this.getRelatedTables(table, schema)
+        : allRelationships.filter(
+            rel => rel.from.table === table || rel.to.table === table
+          );
+      for (const rel of rels) {
+        // Avoid duplicates
+        if (!relevantRelationships.find(r => r.id === rel.id)) {
+          relevantRelationships.push(rel);
+        }
+      }
+    }
+
+    if (relevantRelationships.length > 0) {
+      const parts: string[] = [
+        '',
+        '## Table Relationships',
+        'Use these relationships when joining tables:',
+        '',
+      ];
+
+      for (const rel of relevantRelationships) {
+        const fromStr = `${rel.from.schema}.${rel.from.table}`;
+        const toStr = `${rel.to.schema}.${rel.to.table}`;
+        parts.push(`- **${fromStr}** → **${toStr}**`);
+        parts.push(`  Join: \`${rel.joinCondition}\``);
+        if (rel.cardinality) {
+          parts.push(`  Cardinality: ${rel.cardinality}`);
+        }
+        if (rel.notes) {
+          parts.push(`  Note: ${rel.notes}`);
+        }
+        parts.push('');
+      }
+
+      context += '\n' + parts.join('\n');
+    }
+
+    return context;
+  }
 }
 
 /**