@kiyeonjeon21/datacontext 0.3.2 → 0.4.0

package/src/knowledge/types.ts

@@ -1,6 +1,40 @@
 /**
  * Knowledge Types
- * Type definitions for the knowledge store (descriptions, rules, examples)
+ * 
+ * Type definitions for the Knowledge Store / Knowledge Graph.
+ * This module defines the core data structures for storing business context
+ * about database schemas, including:
+ * 
+ * - **TableDescription**: Business context for tables and columns
+ * - **TableRelationship**: Graph edges representing table relationships (FK, implicit, manual)
+ * - **QueryExample**: Verified SQL patterns with intent
+ * - **BusinessRule**: Domain constraints and conventions
+ * - **BusinessTerm**: Glossary entries for natural language → SQL translation
+ * 
+ * ## Knowledge Graph Architecture
+ * 
+ * The Knowledge Store is designed to evolve into a Knowledge Graph:
+ * 
+ * ```
+ * ┌─────────────────────────────────────────────────────────────────┐
+ * │                     Knowledge Graph                             │
+ * │                                                                 │
+ * │   [Table A] ──JOINS_WITH──> [Table B] ──JOINS_WITH──> [Table C]│
+ * │       │                         │                               │
+ * │       │                         │                               │
+ * │   [Description]             [Description]                       │
+ * │   [Rules]                   [Rules]                             │
+ * │   [Examples]                [Terms]                             │
+ * │                                                                 │
+ * └─────────────────────────────────────────────────────────────────┘
+ * ```
+ * 
+ * ## Multi-Database Support
+ * 
+ * All references use `TableRef` which includes optional database identifier,
+ * enabling future cross-database relationship tracking.
+ * 
+ * @module knowledge/types
  */
 
 /** Source of the knowledge entry */
@@ -26,6 +60,127 @@ export interface KnowledgeMeta {
   owner?: string;
 }
 
+// ============================================================
+// Graph Types (Relationships)
+// ============================================================
+
+/**
+ * Fully qualified table reference.
+ * 
+ * Enables cross-database and cross-schema relationship tracking.
+ * The `database` field is optional for single-database scenarios.
+ * 
+ * @example
+ * ```typescript
+ * // Single database (most common)
+ * const ref: TableRef = { schema: 'public', table: 'users' };
+ * 
+ * // Multi-database (future)
+ * const ref: TableRef = { database: 'analytics', schema: 'public', table: 'events' };
+ * ```
+ */
+export interface TableRef {
+  /** Database identifier (optional for single-DB mode) */
+  database?: string;
+  /** Schema name */
+  schema: string;
+  /** Table name */
+  table: string;
+}
+
+/**
+ * Relationship type between tables.
+ * 
+ * - `foreign_key`: Relationship derived from database FK constraint
+ * - `implicit_join`: Relationship inferred from query patterns (no FK)
+ * - `manual`: Relationship defined manually by user
+ */
+export type RelationshipType = 'foreign_key' | 'implicit_join' | 'manual';
+
+/**
+ * Cardinality of a table relationship.
+ * 
+ * - `one-to-one`: Each row in A maps to exactly one row in B
+ * - `one-to-many`: Each row in A maps to multiple rows in B (most common for FK)
+ * - `many-to-one`: Multiple rows in A map to one row in B
+ * - `many-to-many`: Requires junction table
+ */
+export type RelationshipCardinality = 'one-to-one' | 'one-to-many' | 'many-to-one' | 'many-to-many';
+
+/**
+ * Table relationship representing a graph edge.
+ * 
+ * This is the core type for the Knowledge Graph, representing how tables
+ * are connected. Relationships can be:
+ * - Auto-harvested from foreign key constraints
+ * - Inferred from query patterns (implicit joins)
+ * - Manually defined by users
+ * 
+ * ## Graph Traversal
+ * 
+ * Relationships enable graph queries like:
+ * - "Find all tables related to users"
+ * - "Find the shortest join path from orders to products"
+ * - "Which tables reference this column?"
+ * 
+ * @example
+ * ```typescript
+ * // FK-based relationship (auto-harvested)
+ * const relationship: TableRelationship = {
+ *   type: 'table_relationship',
+ *   from: { schema: 'public', table: 'orders' },
+ *   to: { schema: 'public', table: 'users' },
+ *   relationshipType: 'foreign_key',
+ *   joinCondition: 'orders.user_id = users.id',
+ *   cardinality: 'many-to-one',
+ *   fromColumns: ['user_id'],
+ *   toColumns: ['id'],
+ *   isPreferred: true,
+ *   // ... metadata fields
+ * };
+ * 
+ * // Implicit join (learned from queries)
+ * const implicit: TableRelationship = {
+ *   type: 'table_relationship',
+ *   from: { schema: 'public', table: 'events' },
+ *   to: { schema: 'public', table: 'users' },
+ *   relationshipType: 'implicit_join',
+ *   joinCondition: 'events.actor_id = users.id',
+ *   cardinality: 'many-to-one',
+ *   isPreferred: false,
+ *   notes: 'Inferred from query patterns',
+ *   // ... metadata fields
+ * };
+ * ```
+ */
+export interface TableRelationship extends KnowledgeMeta {
+  type: 'table_relationship';
+  /** Source table (FK side for foreign_key type) */
+  from: TableRef;
+  /** Target table (referenced side for foreign_key type) */
+  to: TableRef;
+  /** How this relationship was discovered */
+  relationshipType: RelationshipType;
+  /** SQL join condition */
+  joinCondition: string;
+  /** Relationship cardinality */
+  cardinality?: RelationshipCardinality;
+  /** Columns in the source table involved in the join */
+  fromColumns?: string[];
+  /** Columns in the target table involved in the join */
+  toColumns?: string[];
+  /** Whether this is the preferred join path (when multiple exist) */
+  isPreferred: boolean;
+  /** FK constraint name (if from database) */
+  constraintName?: string;
+  /** Human-readable notes about this relationship */
+  notes?: string;
+}
+
+// ============================================================
+// Entity Types (Nodes)
+// ============================================================
+
 /** Column description */
 export interface ColumnDescription {
   /** Column name */
@@ -135,9 +290,42 @@ export interface BusinessTerm extends KnowledgeMeta {
 }
 
 /** Union type for all knowledge entries */
-export type KnowledgeEntry = TableDescription | QueryExample | BusinessRule | BusinessTerm;
+export type KnowledgeEntry = 
+  | TableDescription 
+  | TableRelationship 
+  | QueryExample 
+  | BusinessRule 
+  | BusinessTerm;
 
-/** Complete knowledge data structure */
+/**
+ * Complete knowledge data structure.
+ * 
+ * This is the root data structure for the Knowledge Store, containing:
+ * - **Nodes**: Tables, columns (within TableDescription)
+ * - **Edges**: Relationships between tables (TableRelationship)
+ * - **Annotations**: Rules, examples, terms
+ * 
+ * ## Data Model (Graph Perspective)
+ * 
+ * ```
+ * Nodes:
+ *   - TableDescription (with embedded ColumnDescription[])
+ * 
+ * Edges:
+ *   - TableRelationship (from → to with join condition)
+ * 
+ * Annotations (attached to nodes):
+ *   - QueryExample (references tables)
+ *   - BusinessRule (applies to tables/columns)
+ *   - BusinessTerm (applies to tables/columns)
+ * ```
+ * 
+ * ## Version History
+ * 
+ * - `1.0.0`: Initial version (tables, examples, rules)
+ * - `1.1.0`: Added businessTerms (glossary)
+ * - `2.0.0`: Added tableRelationships (graph edges)
+ */
 export interface KnowledgeData {
   /** Version of the knowledge format */
   version: string;
@@ -147,8 +335,16 @@ export interface KnowledgeData {
   schemaHash: string;
   /** Last sync timestamp */
   lastSyncAt: string;
-  /** Table descriptions */
+  
+  // === Nodes ===
+  /** Table descriptions (graph nodes) */
   tableDescriptions: TableDescription[];
+  
+  // === Edges ===
+  /** Table relationships (graph edges) - NEW in v2.0 */
+  tableRelationships: TableRelationship[];
+  
+  // === Annotations ===
   /** Query examples */
   queryExamples: QueryExample[];
   /** Business rules */
@@ -157,6 +353,10 @@ export interface KnowledgeData {
   businessTerms: BusinessTerm[];
 }
 
+// ============================================================
+// Factory Functions
+// ============================================================
+
 /** Generate a unique ID */
 export function generateId(): string {
   return `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
@@ -181,3 +381,120 @@ export function createKnowledgeMeta(
   };
 }
 
+/**
+ * Create a TableRef from table name and optional schema/database.
+ * 
+ * @param table - Table name
+ * @param schema - Schema name (defaults to 'public')
+ * @param database - Database name (optional)
+ * @returns TableRef object
+ * 
+ * @example
+ * ```typescript
+ * const ref = createTableRef('users');
+ * // { schema: 'public', table: 'users' }
+ * 
+ * const ref2 = createTableRef('orders', 'sales', 'analytics_db');
+ * // { database: 'analytics_db', schema: 'sales', table: 'orders' }
+ * ```
+ */
+export function createTableRef(
+  table: string,
+  schema: string = 'public',
+  database?: string
+): TableRef {
+  const ref: TableRef = { schema, table };
+  if (database) {
+    ref.database = database;
+  }
+  return ref;
+}
+
+/**
+ * Convert TableRef to a fully qualified string representation.
+ * 
+ * @param ref - TableRef to stringify
+ * @returns String like "database.schema.table" or "schema.table"
+ * 
+ * @example
+ * ```typescript
+ * tableRefToString({ schema: 'public', table: 'users' })
+ * // "public.users"
+ * 
+ * tableRefToString({ database: 'analytics', schema: 'public', table: 'events' })
+ * // "analytics.public.events"
+ * ```
+ */
+export function tableRefToString(ref: TableRef): string {
+  if (ref.database) {
+    return `${ref.database}.${ref.schema}.${ref.table}`;
+  }
+  return `${ref.schema}.${ref.table}`;
+}
+
+/**
+ * Check if two TableRefs refer to the same table.
+ * 
+ * @param a - First TableRef
+ * @param b - Second TableRef
+ * @returns true if they refer to the same table
+ */
+export function tableRefsEqual(a: TableRef, b: TableRef): boolean {
+  return (
+    a.table === b.table &&
+    a.schema === b.schema &&
+    (a.database ?? '') === (b.database ?? '')
+  );
+}
+
+/**
+ * Create a TableRelationship from foreign key information.
+ * 
+ * This is the standard way to create relationships from harvested FK data.
+ * 
+ * @param params - Relationship parameters
+ * @returns Complete TableRelationship object
+ * 
+ * @example
+ * ```typescript
+ * const relationship = createRelationshipFromFK({
+ *   fromTable: 'orders',
+ *   fromSchema: 'public',
+ *   fromColumn: 'user_id',
+ *   toTable: 'users',
+ *   toSchema: 'public',
+ *   toColumn: 'id',
+ *   constraintName: 'orders_user_id_fkey',
+ *   schemaHash: 'abc123',
+ * });
+ * ```
+ */
+export function createRelationshipFromFK(params: {
+  fromTable: string;
+  fromSchema?: string;
+  fromColumn: string;
+  toTable: string;
+  toSchema?: string;
+  toColumn: string;
+  constraintName?: string;
+  schemaHash: string;
+  database?: string;
+}): TableRelationship {
+  const fromSchema = params.fromSchema ?? 'public';
+  const toSchema = params.toSchema ?? 'public';
+  
+  return {
+    ...createKnowledgeMeta('auto', params.schemaHash),
+    type: 'table_relationship',
+    from: createTableRef(params.fromTable, fromSchema, params.database),
+    to: createTableRef(params.toTable, toSchema, params.database),
+    relationshipType: 'foreign_key',
+    joinCondition: `${params.fromTable}.${params.fromColumn} = ${params.toTable}.${params.toColumn}`,
+    cardinality: 'many-to-one', // Default for FK
+    fromColumns: [params.fromColumn],
+    toColumns: [params.toColumn],
+    isPreferred: true,
+    constraintName: params.constraintName,
+  };
+}
+

package/src/mcp/tools.ts

@@ -355,6 +355,54 @@ export function getMcpTools(): Tool[] {
         required: ['query'],
       },
     },
+    // === Knowledge Graph Tools ===
+    {
+      name: 'find_join_path',
+      description: 'Find the optimal join path between two tables using the Knowledge Graph. Automatically discovers relationships and suggests JOIN conditions. This is useful when you need to join tables that are not directly related.',
+      inputSchema: {
+        type: 'object',
+        properties: {
+          fromTable: {
+            type: 'string',
+            description: 'Source table name to start from.',
+          },
+          toTable: {
+            type: 'string',
+            description: 'Destination table name to reach.',
+          },
+          schema: {
+            type: 'string',
+            description: 'Database schema. Defaults to "public" for PostgreSQL, "main" for SQLite.',
+          },
+        },
+        required: ['fromTable', 'toTable'],
+      },
+    },
+    {
+      name: 'get_relationships',
+      description: 'Get all table relationships (foreign keys, joins) from the Knowledge Graph. Returns information about how tables are connected.',
+      inputSchema: {
+        type: 'object',
+        properties: {
+          table: {
+            type: 'string',
+            description: 'Optional: Filter to relationships involving this table.',
+          },
+          schema: {
+            type: 'string',
+            description: 'Database schema. Defaults to "public".',
+          },
+        },
+      },
+    },
+    {
+      name: 'get_graph_summary',
+      description: 'Get a summary of the Knowledge Graph including total tables, relationships, and connectivity statistics.',
+      inputSchema: {
+        type: 'object',
+        properties: {},
+      },
+    },
   ];
 }
 
@@ -401,6 +449,13 @@ export async function handleToolCall(
       return handleSearchTerms(args, context);
     case 'enhance_query':
       return handleEnhanceQuery(args, context);
+    // === Knowledge Graph Tools ===
+    case 'find_join_path':
+      return handleFindJoinPath(args, context);
+    case 'get_relationships':
+      return handleGetRelationships(args, context);
+    case 'get_graph_summary':
+      return handleGetGraphSummary(args, context);
     default:
       throw new Error(`Unknown tool: ${name}`);
   }
@@ -589,6 +644,9 @@ async function handleDescribeTable(
 
 /**
  * Handle get_context tool
+ * 
+ * Enhanced to include table relationships from the Knowledge Graph.
+ * This helps AI understand how tables connect and generate accurate JOINs.
  */
 async function handleGetContext(
   args: Record<string, unknown>,
@@ -600,14 +658,41 @@ async function handleGetContext(
     throw new Error('Tables array is required');
   }
 
-  const contextText = context.knowledge.buildContext(tables);
+  // Use buildContextWithRelationships to include relationship information
+  const contextText = context.knowledge.buildContextWithRelationships(tables);
   
   // Also get query examples
   const examples = context.knowledge.getQueryExamplesForTables(tables);
 
+  // Get relationships involving these tables
+  const allRelationships = context.knowledge.getRelationships();
+  const relevantRelationships = allRelationships.filter(
+    rel => tables.includes(rel.from.table) || tables.includes(rel.to.table)
+  );
+
+  // Find related tables (1-hop connections)
+  const relatedTables = new Set<string>();
+  for (const table of tables) {
+    const rels = context.knowledge.getRelatedTables(table);
+    for (const rel of rels) {
+      relatedTables.add(rel.from.table);
+      relatedTables.add(rel.to.table);
+    }
+  }
+  // Remove tables already in the query
+  tables.forEach(t => relatedTables.delete(t));
+
   return {
     context: contextText,
     tables,
+    relatedTables: Array.from(relatedTables),
+    relationships: relevantRelationships.map(rel => ({
+      from: `${rel.from.schema}.${rel.from.table}`,
+      to: `${rel.to.schema}.${rel.to.table}`,
+      type: rel.relationshipType,
+      joinCondition: rel.joinCondition,
+      cardinality: rel.cardinality,
+    })),
     queryExamples: examples.slice(0, 5).map(ex => ({
       intent: ex.intent,
       sql: ex.sql,
@@ -709,13 +794,15 @@ async function handleHarvestMetadata(
     }
 
     const result = await harvester.syncToKnowledge(schema);
+    const relMsg = result.relationshipsAdded ? `, ${result.relationshipsAdded} relationships` : '';
     return {
       success: true,
       schema,
       tablesProcessed: result.tablesProcessed,
       descriptionsAdded: result.descriptionsAdded,
       columnsAdded: result.columnsAdded,
-      message: `Harvested ${result.tablesProcessed} tables. Added ${result.descriptionsAdded} table descriptions and ${result.columnsAdded} column descriptions.`,
+      relationshipsAdded: result.relationshipsAdded ?? 0,
+      message: `Harvested ${result.tablesProcessed} tables. Added ${result.descriptionsAdded} table descriptions, ${result.columnsAdded} column descriptions${relMsg}.`,
     };
   }
 
@@ -730,13 +817,15 @@ async function handleHarvestMetadata(
   }
 
   const result = await context.harvester.syncToKnowledge(schema);
+  const relMsg = result.relationshipsAdded ? `, ${result.relationshipsAdded} relationships` : '';
   return {
     success: true,
     schema,
     tablesProcessed: result.tablesProcessed,
     descriptionsAdded: result.descriptionsAdded,
     columnsAdded: result.columnsAdded,
-    message: `Harvested ${result.tablesProcessed} tables. Added ${result.descriptionsAdded} table descriptions and ${result.columnsAdded} column descriptions.`,
+    relationshipsAdded: result.relationshipsAdded ?? 0,
+    message: `Harvested ${result.tablesProcessed} tables. Added ${result.descriptionsAdded} table descriptions, ${result.columnsAdded} column descriptions${relMsg}.`,
   };
 }
 
@@ -1235,6 +1324,187 @@ function calculateTermRelevance(query: string, term: { term: string; synonyms: s
   return 0.5;
 }
 
+// ============================================================
+// Knowledge Graph Tools
+// ============================================================
+
+/**
+ * Handle find_join_path tool
+ * 
+ * Uses the Knowledge Graph to find the optimal path between two tables.
+ * This is particularly useful for complex schemas where tables are connected
+ * through intermediate tables (e.g., order_items → orders → users).
+ * 
+ * @example
+ * // Request:
+ * { fromTable: "order_items", toTable: "users" }
+ * 
+ * // Response:
+ * {
+ *   path: [
+ *     { from: "order_items", to: "orders", join: "order_items.order_id = orders.id" },
+ *     { from: "orders", to: "users", join: "orders.user_id = users.id" }
+ *   ],
+ *   sqlHint: "JOIN orders ON ... JOIN users ON ..."
+ * }
+ */
+async function handleFindJoinPath(
+  args: Record<string, unknown>,
+  context: ToolContext
+): Promise<unknown> {
+  const fromTable = args.fromTable as string;
+  const toTable = args.toTable as string;
+  const schema = (args.schema as string) ?? 'public';
+
+  if (!fromTable || !toTable) {
+    throw new Error('Both fromTable and toTable are required');
+  }
+
+  // Find path using BFS in the Knowledge Graph
+  const path = context.knowledge.findJoinPath(fromTable, toTable, schema);
+
+  if (path.length === 0) {
+    // Try with 'main' schema for SQLite
+    const pathMain = context.knowledge.findJoinPath(fromTable, toTable, 'main');
+    if (pathMain.length === 0) {
+      return {
+        found: false,
+        message: `No path found between '${fromTable}' and '${toTable}'. These tables may not be connected via foreign keys. Consider using implicit joins or check if the tables exist.`,
+        suggestion: 'Run harvest_metadata to discover relationships, or manually define a relationship.',
+      };
+    }
+    return formatJoinPathResult(fromTable, toTable, pathMain);
+  }
+
+  return formatJoinPathResult(fromTable, toTable, path);
+}
+
+/**
+ * Format join path result for AI consumption
+ */
+function formatJoinPathResult(
+  fromTable: string,
+  toTable: string,
+  path: Array<{ from: { schema: string; table: string }; to: { schema: string; table: string }; joinCondition: string; cardinality?: string }>
+): unknown {
+  const steps = path.map((rel, idx) => ({
+    step: idx + 1,
+    from: `${rel.from.schema}.${rel.from.table}`,
+    to: `${rel.to.schema}.${rel.to.table}`,
+    joinCondition: rel.joinCondition,
+    cardinality: rel.cardinality || 'unknown',
+  }));
+
+  // Generate SQL hint
+  const sqlParts = path.map(rel => 
+    `JOIN ${rel.to.table} ON ${rel.joinCondition}`
+  );
+  const sqlHint = `FROM ${fromTable}\n${sqlParts.join('\n')}`;
+
+  // Generate human-readable explanation
+  const pathDescription = path.map(rel => 
+    `${rel.from.table} → ${rel.to.table}`
+  ).join(' → ');
+
+  return {
+    found: true,
+    from: fromTable,
+    to: toTable,
+    hops: path.length,
+    path: steps,
+    pathDescription,
+    sqlHint,
+    message: `Found a ${path.length}-hop path: ${pathDescription}`,
+  };
+}
+
+/**
+ * Handle get_relationships tool
+ * 
+ * Returns all relationships from the Knowledge Graph, optionally filtered by table.
+ */
+async function handleGetRelationships(
+  args: Record<string, unknown>,
+  context: ToolContext
+): Promise<unknown> {
+  const table = args.table as string | undefined;
+  const schema = (args.schema as string) ?? 'public';
+
+  let relationships = context.knowledge.getRelationships();
+
+  // Filter by table if specified
+  if (table) {
+    relationships = relationships.filter(
+      rel => rel.from.table === table || rel.to.table === table
+    );
+  }
+
+  if (relationships.length === 0) {
+    return {
+      count: 0,
+      relationships: [],
+      message: table 
+        ? `No relationships found for table '${table}'. Run harvest_metadata to discover foreign keys.`
+        : 'No relationships in Knowledge Graph. Run harvest_metadata to discover foreign keys.',
+    };
+  }
+
+  return {
+    count: relationships.length,
+    relationships: relationships.map(rel => ({
+      id: rel.id,
+      from: `${rel.from.schema}.${rel.from.table}`,
+      to: `${rel.to.schema}.${rel.to.table}`,
+      type: rel.relationshipType,
+      joinCondition: rel.joinCondition,
+      cardinality: rel.cardinality || 'unknown',
+      isPreferred: rel.isPreferred,
+      notes: rel.notes,
+    })),
+    message: `Found ${relationships.length} relationship(s)${table ? ` involving '${table}'` : ''}.`,
+  };
+}
+
+/**
+ * Handle get_graph_summary tool
+ * 
+ * Returns summary statistics about the Knowledge Graph.
+ */
+async function handleGetGraphSummary(
+  args: Record<string, unknown>,
+  context: ToolContext
+): Promise<unknown> {
+  const summary = context.knowledge.getGraphSummary();
+  const relationships = context.knowledge.getRelationships();
+
+  // Build a quick connectivity summary
+  const tableConnections = new Map<string, number>();
+  for (const rel of relationships) {
+    const fromKey = rel.from.table;
+    const toKey = rel.to.table;
+    tableConnections.set(fromKey, (tableConnections.get(fromKey) || 0) + 1);
+    tableConnections.set(toKey, (tableConnections.get(toKey) || 0) + 1);
+  }
+
+  // Sort by connections
+  const sortedTables = Array.from(tableConnections.entries())
+    .sort((a, b) => b[1] - a[1])
+    .slice(0, 10)
+    .map(([table, connections]) => ({ table, connections }));
+
+  return {
+    nodeCount: summary.nodeCount,
+    edgeCount: summary.edgeCount,
+    tablesWithRelationships: summary.tablesWithRelationships,
+    isolatedTables: summary.isolatedTables,
+    relationshipsByType: summary.relationshipsByType,
+    mostConnectedTables: sortedTables,
+    message: summary.edgeCount > 0
+      ? `Knowledge Graph has ${summary.edgeCount} relationships connecting ${summary.tablesWithRelationships} tables.`
+      : 'Knowledge Graph is empty. Run harvest_metadata to discover relationships.',
+  };
+}
+
 // ============================================================
 // Utility Functions
 // ============================================================