@deepagents/text2sql 0.30.0 → 0.32.0

package/dist/lib/adapters/postgres/index.js

@@ -328,7 +328,18 @@ var Adapter = class {
       fragments.push(this.#viewToFragment(v));
     }
     const tableMap = new Map(ctx.tables.map((t) => [t.name, t]));
+    const tableColumnSets = new Map(
+      ctx.tables.map((t) => [t.name, new Set(t.columns.map((c) => c.name))])
+    );
     for (const rel of ctx.relationships) {
+      const sourceColumns = tableColumnSets.get(rel.table);
+      const targetColumns = tableColumnSets.get(rel.referenced_table);
+      if (sourceColumns && rel.from.some((column2) => !sourceColumns.has(column2))) {
+        continue;
+      }
+      if (targetColumns && rel.to.some((column2) => !targetColumns.has(column2))) {
+        continue;
+      }
       const sourceTable = tableMap.get(rel.table);
       const targetTable = tableMap.get(rel.referenced_table);
       fragments.push(
@@ -383,7 +394,10 @@ var Adapter = class {
         stats: col.stats
       })
     );
-    const indexFragments = (t.indexes ?? []).map(
+    const presentColumns = new Set(t.columns.map((c) => c.name));
+    const indexFragments = (t.indexes ?? []).filter(
+      (idx) => idx.columns.every((column2) => presentColumns.has(column2))
+    ).map(
       (idx) => index({
         name: idx.name,
         columns: idx.columns,
@@ -393,6 +407,8 @@ var Adapter = class {
     );
     const constraintFragments = (t.constraints ?? []).filter(
       (c) => c.type === "CHECK" || c.type === "UNIQUE" && (c.columns?.length ?? 0) > 1
+    ).filter(
+      (c) => !c.columns?.length || c.columns.every((column2) => presentColumns.has(column2))
     ).map(
       (c) => constraint({
         name: c.name,
@@ -586,6 +602,21 @@ var Adapter = class {
 };
 
 // packages/text2sql/src/lib/adapters/groundings/abstract.grounding.ts
+function filterColumns(columns, filter) {
+  if (Array.isArray(filter)) {
+    return columns.filter((col) => filter.includes(col.name));
+  }
+  if (filter instanceof RegExp) {
+    return columns.filter((col) => filter.test(col.name));
+  }
+  return columns.filter((col) => filter(col.name));
+}
+function applyColumnFilter(entity, columnsConfig) {
+  if (!columnsConfig) return entity;
+  const filter = columnsConfig[entity.name];
+  if (!filter) return entity;
+  return { ...entity, columns: filterColumns(entity.columns, filter) };
+}
 var AbstractGrounding = class {
   /**
    * Grounding identifier for debugging/logging.
@@ -957,11 +988,13 @@ var RowCountGrounding = class extends AbstractGrounding {
 // packages/text2sql/src/lib/adapters/groundings/table.grounding.ts
 var TableGrounding = class extends AbstractGrounding {
   #filter;
+  #columns;
   #forward;
   #backward;
   constructor(config = {}) {
     super("table");
     this.#filter = config.filter;
+    this.#columns = config.columns;
     this.#forward = config.forward;
     this.#backward = config.backward;
   }
@@ -975,7 +1008,9 @@ var TableGrounding = class extends AbstractGrounding {
     const backward = this.#backward;
     if (!forward && !backward) {
       const tables3 = await Promise.all(
-        seedTables.map((name) => this.getTable(name))
+        seedTables.map(
+          async (name) => applyColumnFilter(await this.getTable(name), this.#columns)
+        )
       );
       ctx.tables.push(...tables3);
       return;
@@ -999,11 +1034,15 @@ var TableGrounding = class extends AbstractGrounding {
       if (forwardVisited.has(name)) continue;
       forwardVisited.add(name);
       if (!tables2[name]) {
-        tables2[name] = await this.getTable(name);
+        tables2[name] = applyColumnFilter(
+          await this.getTable(name),
+          this.#columns
+        );
       }
       if (depth < forwardLimit) {
         const rels = await this.findOutgoingRelations(name);
         for (const rel of rels) {
+          if (!this.isRelationshipVisible(rel)) continue;
           this.addRelationship(rel, allRelationships, seenRelationships);
           if (!forwardVisited.has(rel.referenced_table)) {
             forwardQueue.push({ name: rel.referenced_table, depth: depth + 1 });
@@ -1019,11 +1058,15 @@ var TableGrounding = class extends AbstractGrounding {
       if (backwardVisited.has(name)) continue;
       backwardVisited.add(name);
       if (!tables2[name]) {
-        tables2[name] = await this.getTable(name);
+        tables2[name] = applyColumnFilter(
+          await this.getTable(name),
+          this.#columns
+        );
       }
       if (depth < backwardLimit) {
         const rels = await this.findIncomingRelations(name);
         for (const rel of rels) {
+          if (!this.isRelationshipVisible(rel)) continue;
           this.addRelationship(rel, allRelationships, seenRelationships);
           if (!backwardVisited.has(rel.table)) {
             backwardQueue.push({ name: rel.table, depth: depth + 1 });
@@ -1063,6 +1106,16 @@ var TableGrounding = class extends AbstractGrounding {
       all.push(rel);
     }
   }
+  isRelationshipVisible(rel) {
+    return this.areColumnsVisible(rel.from, this.#columns?.[rel.table]) && this.areColumnsVisible(rel.to, this.#columns?.[rel.referenced_table]);
+  }
+  areColumnsVisible(names, filter) {
+    if (!filter) return true;
+    return filterColumns(
+      names.map((name) => ({ name })),
+      filter
+    ).length === names.length;
+  }
 };
 
 // packages/text2sql/src/lib/adapters/postgres/column-stats.postgres.grounding.ts
@@ -2125,10 +2178,12 @@ var PostgresTableGrounding = class extends TableGrounding {
 // packages/text2sql/src/lib/adapters/groundings/view.grounding.ts
 var ViewGrounding = class extends AbstractGrounding {
   #filter;
+  #columns;
   includeDefinition;
   constructor(config = {}) {
     super("view");
     this.#filter = config.filter;
+    this.#columns = config.columns;
     this.includeDefinition = config.includeDefinition ?? true;
   }
   /**
@@ -2138,7 +2193,9 @@ var ViewGrounding = class extends AbstractGrounding {
   async execute(ctx) {
     const viewNames = await this.applyFilter();
     const views2 = await Promise.all(
-      viewNames.map((name) => this.getView(name))
+      viewNames.map(
+        async (name) => applyColumnFilter(await this.getView(name), this.#columns)
+      )
     );
     ctx.views.push(...views2);
   }