@deepagents/text2sql 0.22.0 → 0.23.0

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

@@ -98,7 +98,8 @@ function createGroundingContext() {
     tables: [],
     views: [],
     relationships: [],
-    info: void 0
+    info: void 0,
+    cache: /* @__PURE__ */ new Map()
   };
 }
 
@@ -737,43 +738,152 @@ var BigQuery = class extends Adapter {
   }
 };
 
+// packages/text2sql/src/lib/adapters/bigquery/bigquery-fk.ts
+var FK_CACHE_PREFIX = "fk:";
+async function resolveForeignKey(adapter, constraintDataset, constraintName, childColumns, cache) {
+  const cacheKey = `${FK_CACHE_PREFIX}${constraintDataset}:${constraintName}`;
+  if (cache?.has(cacheKey)) {
+    const cached = cache.get(cacheKey);
+    if (!cached) return void 0;
+    return {
+      ...cached,
+      childColumns: [...childColumns].sort((a, b) => a.ordinal - b.ordinal).map((c) => c.column)
+    };
+  }
+  const result = await resolveFK(
+    adapter,
+    constraintDataset,
+    constraintName,
+    childColumns
+  );
+  cache?.set(cacheKey, result);
+  return result;
+}
+async function resolveFK(adapter, constraintDataset, constraintName, childColumns) {
+  const refRows = await adapter.runQuery(`
+    SELECT DISTINCT table_schema, table_name
+    FROM ${adapter.infoSchemaView(constraintDataset, "CONSTRAINT_COLUMN_USAGE")}
+    WHERE constraint_name = '${adapter.escapeString(constraintName)}'
+  `);
+  const referenced = refRows.find((r) => r.table_schema && r.table_name);
+  if (!referenced?.table_schema || !referenced.table_name) {
+    return void 0;
+  }
+  const referencedDataset = referenced.table_schema;
+  const referencedTable = referenced.table_name;
+  if (!adapter.isDatasetAllowed(referencedDataset)) {
+    return void 0;
+  }
+  const pkConstraintRows = await adapter.runQuery(`
+    SELECT constraint_name
+    FROM ${adapter.infoSchemaView(referencedDataset, "TABLE_CONSTRAINTS")}
+    WHERE constraint_type = 'PRIMARY KEY'
+      AND table_name = '${adapter.escapeString(referencedTable)}'
+    LIMIT 1
+  `);
+  const pkConstraintName = pkConstraintRows[0]?.constraint_name;
+  if (!pkConstraintName) return void 0;
+  const pkColumnRows = await adapter.runQuery(`
+    SELECT column_name, ordinal_position
+    FROM ${adapter.infoSchemaView(referencedDataset, "KEY_COLUMN_USAGE")}
+    WHERE constraint_name = '${adapter.escapeString(pkConstraintName)}'
+      AND table_name = '${adapter.escapeString(referencedTable)}'
+    ORDER BY ordinal_position
+  `);
+  const pkByOrdinal = /* @__PURE__ */ new Map();
+  for (const row of pkColumnRows) {
+    if (!row.column_name || row.ordinal_position == null) continue;
+    pkByOrdinal.set(row.ordinal_position, row.column_name);
+  }
+  const ordered = [...childColumns].sort((a, b) => a.ordinal - b.ordinal);
+  return {
+    referencedDataset,
+    referencedTable,
+    referencedColumns: ordered.map((c) => {
+      const pkOrdinal = c.pkOrdinal ?? c.ordinal;
+      return pkByOrdinal.get(pkOrdinal) ?? "unknown";
+    }),
+    childColumns: ordered.map((c) => c.column)
+  };
+}
+
 // packages/text2sql/src/lib/adapters/bigquery/constraint.bigquery.grounding.ts
 var BigQueryConstraintGrounding = class extends ConstraintGrounding {
   #adapter;
+  #cache;
   constructor(adapter, config = {}) {
     super(config);
     this.#adapter = adapter;
   }
-  async getConstraints(tableName) {
-    const { schema: dataset, table: table2 } = this.#adapter.parseTableName(tableName);
-    const constraints2 = [];
-    const columnRows = await this.#adapter.runQuery(`
-      SELECT column_name, is_nullable, column_default
+  async execute(ctx) {
+    this.#cache = ctx.cache;
+    const byDataset = /* @__PURE__ */ new Map();
+    for (const table2 of ctx.tables) {
+      const { schema: dataset } = this.#adapter.parseTableName(table2.name);
+      const list = byDataset.get(dataset) ?? [];
+      list.push(table2);
+      byDataset.set(dataset, list);
+    }
+    for (const [dataset, tables2] of byDataset) {
+      try {
+        await this.#batchConstraints(dataset, tables2);
+      } catch (error) {
+        console.warn(
+          "Error collecting constraints for dataset",
+          dataset,
+          error
+        );
+      }
+    }
+  }
+  async #batchConstraints(dataset, tables2) {
+    const tableNames = tables2.map(
+      (t) => this.#adapter.parseTableName(t.name).table
+    );
+    const inList = tableNames.map((n) => `'${this.#adapter.escapeString(n)}'`).join(", ");
+    const constraintsByTable = /* @__PURE__ */ new Map();
+    for (const name of tableNames) {
+      constraintsByTable.set(name, []);
+    }
+    await this.#batchColumnMetadata(dataset, inList, constraintsByTable);
+    await this.#batchKeyConstraints(dataset, inList, constraintsByTable);
+    for (const table2 of tables2) {
+      const rawName = this.#adapter.parseTableName(table2.name).table;
+      table2.constraints = constraintsByTable.get(rawName) ?? [];
+    }
+  }
+  async #batchColumnMetadata(dataset, inList, constraintsByTable) {
+    const rows = await this.#adapter.runQuery(`
+      SELECT table_name, column_name, is_nullable, column_default
       FROM ${this.#adapter.infoSchemaView(dataset, "COLUMNS")}
-      WHERE table_name = '${this.#adapter.escapeString(table2)}'
-      ORDER BY ordinal_position
+      WHERE table_name IN (${inList})
+      ORDER BY table_name, ordinal_position
     `);
-    for (const row of columnRows) {
-      const col = row.column_name;
-      if (!col) continue;
+    for (const row of rows) {
+      if (!row.table_name || !row.column_name) continue;
+      const constraints2 = constraintsByTable.get(row.table_name);
+      if (!constraints2) continue;
       if ((row.is_nullable ?? "").toUpperCase() === "NO") {
         constraints2.push({
-          name: `${tableName}.${col}.NOT_NULL`,
+          name: `${dataset}.${row.table_name}.${row.column_name}.NOT_NULL`,
           type: "NOT_NULL",
-          columns: [col]
+          columns: [row.column_name]
         });
       }
       if (row.column_default != null && row.column_default !== "") {
         constraints2.push({
-          name: `${tableName}.${col}.DEFAULT`,
+          name: `${dataset}.${row.table_name}.${row.column_name}.DEFAULT`,
           type: "DEFAULT",
-          columns: [col],
+          columns: [row.column_name],
           defaultValue: row.column_default
         });
       }
     }
-    const keyRows = await this.#adapter.runQuery(`
+  }
+  async #batchKeyConstraints(dataset, inList, constraintsByTable) {
+    const rows = await this.#adapter.runQuery(`
       SELECT
+        tc.table_name,
         tc.constraint_name,
         tc.constraint_type,
         kcu.column_name,
@@ -783,100 +893,72 @@ var BigQueryConstraintGrounding = class extends ConstraintGrounding {
       JOIN ${this.#adapter.infoSchemaView(dataset, "KEY_COLUMN_USAGE")} AS kcu
         ON tc.constraint_name = kcu.constraint_name
         AND tc.constraint_schema = kcu.constraint_schema
-      WHERE tc.table_name = '${this.#adapter.escapeString(table2)}'
+      WHERE tc.table_name IN (${inList})
         AND tc.constraint_type IN ('PRIMARY KEY', 'FOREIGN KEY')
-      ORDER BY tc.constraint_name, kcu.ordinal_position
+      ORDER BY tc.table_name, tc.constraint_name, kcu.ordinal_position
     `);
-    const pkByName = /* @__PURE__ */ new Map();
-    const fkByName = /* @__PURE__ */ new Map();
-    for (const row of keyRows) {
-      if (!row.constraint_name || !row.column_name) continue;
+    const pkByTable = /* @__PURE__ */ new Map();
+    const fkByTable = /* @__PURE__ */ new Map();
+    for (const row of rows) {
+      if (!row.table_name || !row.constraint_name || !row.column_name) continue;
       const type = (row.constraint_type ?? "").toUpperCase();
       if (type === "PRIMARY KEY") {
-        const cols = pkByName.get(row.constraint_name) ?? [];
+        const tableMap = pkByTable.get(row.table_name) ?? /* @__PURE__ */ new Map();
+        const cols = tableMap.get(row.constraint_name) ?? [];
         cols.push(row.column_name);
-        pkByName.set(row.constraint_name, cols);
-        continue;
-      }
-      if (type === "FOREIGN KEY") {
-        const cols = fkByName.get(row.constraint_name) ?? [];
+        tableMap.set(row.constraint_name, cols);
+        pkByTable.set(row.table_name, tableMap);
+      } else if (type === "FOREIGN KEY") {
+        const tableMap = fkByTable.get(row.table_name) ?? /* @__PURE__ */ new Map();
+        const cols = tableMap.get(row.constraint_name) ?? [];
         cols.push({
           column: row.column_name,
           ordinal: row.ordinal_position ?? 0,
           pkOrdinal: row.position_in_unique_constraint
         });
-        fkByName.set(row.constraint_name, cols);
+        tableMap.set(row.constraint_name, cols);
+        fkByTable.set(row.table_name, tableMap);
       }
     }
-    for (const [name, cols] of pkByName.entries()) {
-      constraints2.push({
-        name,
-        type: "PRIMARY_KEY",
-        columns: cols
-      });
+    for (const [tableName, pkMap] of pkByTable) {
+      const constraints2 = constraintsByTable.get(tableName);
+      if (!constraints2) continue;
+      for (const [name, cols] of pkMap) {
+        constraints2.push({ name, type: "PRIMARY_KEY", columns: cols });
+      }
     }
-    for (const [constraintName, cols] of fkByName.entries()) {
-      const fk = await this.#buildForeignKeyConstraint({
-        constraintDataset: dataset,
-        constraintName,
-        childTableName: `${dataset}.${table2}`,
-        childColumns: cols
-      });
-      if (fk) constraints2.push(fk);
+    for (const [tableName, fkMap] of fkByTable) {
+      const constraints2 = constraintsByTable.get(tableName);
+      if (!constraints2) continue;
+      for (const [constraintName, childColumns] of fkMap) {
+        try {
+          const resolution = await resolveForeignKey(
+            this.#adapter,
+            dataset,
+            constraintName,
+            childColumns,
+            this.#cache
+          );
+          if (resolution) {
+            constraints2.push({
+              name: constraintName,
+              type: "FOREIGN_KEY",
+              columns: resolution.childColumns,
+              referencedTable: `${resolution.referencedDataset}.${resolution.referencedTable}`,
+              referencedColumns: resolution.referencedColumns
+            });
+          }
+        } catch (err) {
+          console.warn(
+            `Failed to resolve FK constraint ${constraintName}:`,
+            err
+          );
+        }
+      }
     }
-    return constraints2;
   }
-  async #buildForeignKeyConstraint(args) {
-    const refRows = await this.#adapter.runQuery(`
-      SELECT DISTINCT table_schema, table_name
-      FROM ${this.#adapter.infoSchemaView(args.constraintDataset, "CONSTRAINT_COLUMN_USAGE")}
-      WHERE constraint_name = '${this.#adapter.escapeString(args.constraintName)}'
-    `);
-    const referenced = refRows.find((r) => r.table_schema && r.table_name);
-    if (!referenced?.table_schema || !referenced.table_name) {
-      return void 0;
-    }
-    const referencedDataset = referenced.table_schema;
-    const referencedTable = referenced.table_name;
-    if (!this.#adapter.isDatasetAllowed(referencedDataset)) {
-      return void 0;
-    }
-    const pkConstraintRows = await this.#adapter.runQuery(`
-      SELECT constraint_name
-      FROM ${this.#adapter.infoSchemaView(referencedDataset, "TABLE_CONSTRAINTS")}
-      WHERE constraint_type = 'PRIMARY KEY'
-        AND table_name = '${this.#adapter.escapeString(referencedTable)}'
-      LIMIT 1
-    `);
-    const pkConstraintName = pkConstraintRows[0]?.constraint_name;
-    if (!pkConstraintName) return void 0;
-    const pkColumns = await this.#adapter.runQuery(`
-      SELECT column_name, ordinal_position
-      FROM ${this.#adapter.infoSchemaView(referencedDataset, "KEY_COLUMN_USAGE")}
-      WHERE constraint_name = '${this.#adapter.escapeString(pkConstraintName)}'
-        AND table_name = '${this.#adapter.escapeString(referencedTable)}'
-      ORDER BY ordinal_position
-    `);
-    const pkByOrdinal = /* @__PURE__ */ new Map();
-    for (const row of pkColumns) {
-      if (!row.column_name || row.ordinal_position == null) continue;
-      pkByOrdinal.set(row.ordinal_position, row.column_name);
-    }
-    const orderedChild = [...args.childColumns].sort(
-      (a, b) => a.ordinal - b.ordinal
-    );
-    const columns = orderedChild.map((c) => c.column);
-    const referencedColumns = orderedChild.map((c) => {
-      const pkOrdinal = c.pkOrdinal ?? c.ordinal;
-      return pkByOrdinal.get(pkOrdinal) ?? "unknown";
-    });
-    return {
-      name: args.constraintName,
-      type: "FOREIGN_KEY",
-      columns,
-      referencedTable: `${referencedDataset}.${referencedTable}`,
-      referencedColumns
-    };
+  async getConstraints(_tableName) {
+    return [];
   }
 };
 
@@ -887,46 +969,82 @@ var BigQueryIndexesGrounding = class extends IndexesGrounding {
     super(config);
     this.#adapter = adapter;
   }
-  async getIndexes(tableName) {
-    const { schema: dataset, table: table2 } = this.#adapter.parseTableName(tableName);
+  async execute(ctx) {
+    const byDataset = /* @__PURE__ */ new Map();
+    for (const table2 of ctx.tables) {
+      const { schema: dataset } = this.#adapter.parseTableName(table2.name);
+      const list = byDataset.get(dataset) ?? [];
+      list.push(table2);
+      byDataset.set(dataset, list);
+    }
+    for (const [dataset, tables2] of byDataset) {
+      try {
+        await this.#batchIndexes(dataset, tables2);
+      } catch {
+      }
+    }
+  }
+  async #batchIndexes(dataset, tables2) {
+    const tableNames = tables2.map(
+      (t) => this.#adapter.parseTableName(t.name).table
+    );
+    const inList = tableNames.map((n) => `'${this.#adapter.escapeString(n)}'`).join(", ");
     const rows = await this.#adapter.runQuery(`
-      SELECT column_name, is_partitioning_column, clustering_ordinal_position
+      SELECT table_name, column_name, is_partitioning_column, clustering_ordinal_position
       FROM ${this.#adapter.infoSchemaView(dataset, "COLUMNS")}
-      WHERE table_name = '${this.#adapter.escapeString(table2)}'
+      WHERE table_name IN (${inList})
         AND (is_partitioning_column = 'YES' OR clustering_ordinal_position IS NOT NULL)
-      ORDER BY clustering_ordinal_position
+      ORDER BY table_name, clustering_ordinal_position
     `);
-    const partitionColumns = [];
-    const clusteringColumns = [];
+    const byTable = /* @__PURE__ */ new Map();
     for (const row of rows) {
-      if (!row.column_name) continue;
+      if (!row.table_name || !row.column_name) continue;
+      const entry = byTable.get(row.table_name) ?? {
+        partition: [],
+        clustering: []
+      };
       if ((row.is_partitioning_column ?? "").toUpperCase() === "YES") {
-        partitionColumns.push(row.column_name);
+        entry.partition.push(row.column_name);
       }
       if (row.clustering_ordinal_position != null) {
-        clusteringColumns.push({
+        entry.clustering.push({
           name: row.column_name,
           pos: row.clustering_ordinal_position
         });
       }
+      byTable.set(row.table_name, entry);
+    }
+    for (const table2 of tables2) {
+      const rawName = this.#adapter.parseTableName(table2.name).table;
+      const entry = byTable.get(rawName);
+      if (!entry) continue;
+      const indexes2 = [];
+      if (entry.partition.length > 0) {
+        indexes2.push({
+          name: `${rawName}_partition`,
+          columns: entry.partition,
+          type: "PARTITION"
+        });
+      }
+      if (entry.clustering.length > 0) {
+        entry.clustering.sort((a, b) => a.pos - b.pos);
+        indexes2.push({
+          name: `${rawName}_clustering`,
+          columns: entry.clustering.map((c) => c.name),
+          type: "CLUSTERING"
+        });
+      }
+      table2.indexes = indexes2;
+      for (const idx of indexes2) {
+        for (const colName of idx.columns) {
+          const column2 = table2.columns.find((c) => c.name === colName);
+          if (column2) column2.isIndexed = true;
+        }
+      }
     }
-    const indexes2 = [];
-    if (partitionColumns.length > 0) {
-      indexes2.push({
-        name: `${table2}_partition`,
-        columns: partitionColumns,
-        type: "PARTITION"
-      });
-    }
-    if (clusteringColumns.length > 0) {
-      clusteringColumns.sort((a, b) => a.pos - b.pos);
-      indexes2.push({
-        name: `${table2}_clustering`,
-        columns: clusteringColumns.map((c) => c.name),
-        type: "CLUSTERING"
-      });
-    }
-    return indexes2;
+  }
+  async getIndexes(_tableName) {
+    return [];
   }
 };
 
@@ -938,13 +1056,14 @@ var BigQueryInfoGrounding = class extends InfoGrounding {
     this.#adapter = adapter;
   }
   async collectInfo() {
+    const qualifiedTable = this.#adapter.projectId ? "project.dataset.table" : "dataset.table";
     return {
       dialect: "bigquery",
       database: this.#adapter.projectId,
       details: {
         identifierQuote: "`",
         identifiers: {
-          qualifiedTable: "dataset.table",
+          qualifiedTable,
           nestedFieldPath: "col.path.to.field"
         },
         parameters: {
@@ -963,26 +1082,99 @@ var BigQueryRowCountGrounding = class extends RowCountGrounding {
     super(config);
     this.#adapter = adapter;
   }
-  async getRowCount(tableName) {
-    const { schema: dataset, table: table2 } = this.#adapter.parseTableName(tableName);
+  async execute(ctx) {
+    const byDataset = /* @__PURE__ */ new Map();
+    for (const table2 of ctx.tables) {
+      const { schema: dataset } = this.#adapter.parseTableName(table2.name);
+      const list = byDataset.get(dataset) ?? [];
+      list.push(table2);
+      byDataset.set(dataset, list);
+    }
+    for (const [dataset, tables2] of byDataset) {
+      const tableNames = tables2.map(
+        (t) => this.#adapter.parseTableName(t.name).table
+      );
+      const counts = await this.#fetchRowCounts(dataset, tableNames);
+      for (const table2 of tables2) {
+        const rawName = this.#adapter.parseTableName(table2.name).table;
+        const count = counts.get(rawName);
+        if (count != null) {
+          table2.rowCount = count;
+          table2.sizeHint = this.#classifyRowCount(count);
+        }
+      }
+    }
+  }
+  async #fetchRowCounts(dataset, tableNames) {
+    const inList = tableNames.map((n) => `'${this.#adapter.escapeString(n)}'`).join(", ");
+    try {
+      return await this.#fromTableStorage(dataset, inList);
+    } catch {
+    }
+    try {
+      return await this.#fromLegacyTables(dataset, inList);
+    } catch {
+    }
+    return /* @__PURE__ */ new Map();
+  }
+  async #fromTableStorage(dataset, inList) {
     const rows = await this.#adapter.runQuery(`
-      SELECT total_rows
+      SELECT table_name, total_rows
       FROM ${this.#adapter.infoSchemaView(dataset, "TABLE_STORAGE")}
-      WHERE table_name = '${this.#adapter.escapeString(table2)}'
-      LIMIT 1
+      WHERE table_name IN (${inList})
     `);
-    const value = rows[0]?.total_rows;
-    return this.#adapter.toNumber(value);
+    const result = /* @__PURE__ */ new Map();
+    for (const row of rows) {
+      if (!row.table_name) continue;
+      const count = this.#adapter.toNumber(row.total_rows);
+      if (count != null) result.set(row.table_name, count);
+    }
+    return result;
+  }
+  async #fromLegacyTables(dataset, inList) {
+    const projectPrefix = this.#adapter.projectId ? `\`${this.#adapter.projectId}\`.` : "";
+    const rows = await this.#adapter.runQuery(`
+      SELECT table_id AS table_name, row_count
+      FROM ${projectPrefix}\`${dataset}\`.__TABLES__
+      WHERE table_id IN (${inList})
+    `);
+    const result = /* @__PURE__ */ new Map();
+    for (const row of rows) {
+      if (!row.table_name) continue;
+      const count = this.#adapter.toNumber(row.row_count);
+      if (count != null) result.set(row.table_name, count);
+    }
+    return result;
+  }
+  #classifyRowCount(count) {
+    if (count < 100) return "tiny";
+    if (count < 1e3) return "small";
+    if (count < 1e4) return "medium";
+    if (count < 1e5) return "large";
+    return "huge";
+  }
+  async getRowCount(_tableName) {
+    return void 0;
   }
 };
 
 // packages/text2sql/src/lib/adapters/bigquery/table.bigquery.grounding.ts
 var BigQueryTableGrounding = class extends TableGrounding {
   #adapter;
+  #cache;
   constructor(adapter, config = {}) {
     super(config);
     this.#adapter = adapter;
   }
+  async execute(ctx) {
+    this.#cache = ctx.cache;
+    await super.execute(ctx);
+    ctx.tables = ctx.tables.filter((t) => t.columns.length > 0);
+    const tableNames = new Set(ctx.tables.map((t) => t.name));
+    ctx.relationships = ctx.relationships.filter(
+      (r) => tableNames.has(r.table) && tableNames.has(r.referenced_table)
+    );
+  }
   async applyFilter() {
     const names = await super.applyFilter();
     return names.filter((name) => this.#isTableInScope(name));
@@ -1064,14 +1256,21 @@ var BigQueryTableGrounding = class extends TableGrounding {
     }
     const rels = [];
     for (const [constraintName, columns] of byConstraint.entries()) {
-      const rel = await this.#buildForeignKeyRelationship({
-        constraintDataset: dataset,
-        childDataset: dataset,
-        childTable: table2,
+      const resolution = await resolveForeignKey(
+        this.#adapter,
+        dataset,
         constraintName,
-        childColumns: columns
-      });
-      if (rel) rels.push(rel);
+        columns,
+        this.#cache
+      );
+      if (resolution) {
+        rels.push({
+          table: `${dataset}.${table2}`,
+          from: resolution.childColumns,
+          referenced_table: `${resolution.referencedDataset}.${resolution.referencedTable}`,
+          to: resolution.referencedColumns
+        });
+      }
     }
     return rels;
   }
@@ -1087,22 +1286,18 @@ var BigQueryTableGrounding = class extends TableGrounding {
       `);
       for (const row of rows) {
         if (!row.constraint_name) continue;
-        const rel = await this.#buildForeignKeyRelationshipFromConstraintName(
+        const rel = await this.#resolveIncomingRelationship(
           constraintDataset,
-          row.constraint_name
+          row.constraint_name,
+          referencedDataset,
+          referencedTable
         );
-        if (rel && rel.referenced_table === `${referencedDataset}.${referencedTable}`) {
-          rels.push(rel);
-        }
+        if (rel) rels.push(rel);
       }
     }
     return rels;
   }
-  #isTableInScope(tableName) {
-    const { schema } = this.#adapter.parseTableName(tableName);
-    return this.#adapter.isDatasetAllowed(schema);
-  }
-  async #buildForeignKeyRelationshipFromConstraintName(constraintDataset, constraintName) {
+  async #resolveIncomingRelationship(constraintDataset, constraintName, expectedReferencedDataset, expectedReferencedTable) {
     const keyRows = await this.#adapter.runQuery(`
       SELECT
         kcu.constraint_name,
@@ -1126,67 +1321,27 @@ var BigQueryTableGrounding = class extends TableGrounding {
       ordinal: r.ordinal_position ?? 0,
       pkOrdinal: r.position_in_unique_constraint
     }));
-    return this.#buildForeignKeyRelationship({
+    const resolution = await resolveForeignKey(
+      this.#adapter,
       constraintDataset,
-      childDataset: constraintDataset,
-      childTable,
       constraintName,
-      childColumns
-    });
-  }
-  async #buildForeignKeyRelationship(args) {
-    const refTableRows = await this.#adapter.runQuery(`
-      SELECT DISTINCT table_schema, table_name
-      FROM ${this.#adapter.infoSchemaView(args.constraintDataset, "CONSTRAINT_COLUMN_USAGE")}
-      WHERE constraint_name = '${this.#adapter.escapeString(args.constraintName)}'
-    `);
-    const referenced = refTableRows.find((r) => r.table_schema && r.table_name);
-    if (!referenced?.table_schema || !referenced.table_name) {
-      return void 0;
-    }
-    const referencedDataset = referenced.table_schema;
-    const referencedTable = referenced.table_name;
-    if (!this.#adapter.isDatasetAllowed(referencedDataset)) {
-      return void 0;
-    }
-    const pkConstraintRows = await this.#adapter.runQuery(`
-      SELECT constraint_name
-      FROM ${this.#adapter.infoSchemaView(referencedDataset, "TABLE_CONSTRAINTS")}
-      WHERE constraint_type = 'PRIMARY KEY'
-        AND table_name = '${this.#adapter.escapeString(referencedTable)}'
-      LIMIT 1
-    `);
-    const pkConstraintName = pkConstraintRows[0]?.constraint_name;
-    if (!pkConstraintName) {
+      childColumns,
+      this.#cache
+    );
+    if (!resolution || resolution.referencedDataset !== expectedReferencedDataset || resolution.referencedTable !== expectedReferencedTable) {
       return void 0;
     }
-    const pkColumnRows = await this.#adapter.runQuery(`
-      SELECT column_name, ordinal_position
-      FROM ${this.#adapter.infoSchemaView(referencedDataset, "KEY_COLUMN_USAGE")}
-      WHERE constraint_name = '${this.#adapter.escapeString(pkConstraintName)}'
-        AND table_name = '${this.#adapter.escapeString(referencedTable)}'
-      ORDER BY ordinal_position
-    `);
-    const pkByOrdinal = /* @__PURE__ */ new Map();
-    for (const row of pkColumnRows) {
-      if (!row.column_name || row.ordinal_position == null) continue;
-      pkByOrdinal.set(row.ordinal_position, row.column_name);
-    }
-    const orderedChild = [...args.childColumns].sort(
-      (a, b) => a.ordinal - b.ordinal
-    );
-    const from = orderedChild.map((c) => c.column);
-    const to = orderedChild.map((c) => {
-      const pkOrdinal = c.pkOrdinal ?? c.ordinal;
-      return pkByOrdinal.get(pkOrdinal) ?? "unknown";
-    });
     return {
-      table: `${args.childDataset}.${args.childTable}`,
-      from,
-      referenced_table: `${referencedDataset}.${referencedTable}`,
-      to
+      table: `${constraintDataset}.${childTable}`,
+      from: resolution.childColumns,
+      referenced_table: `${resolution.referencedDataset}.${resolution.referencedTable}`,
+      to: resolution.referencedColumns
     };
   }
+  #isTableInScope(tableName) {
+    const { schema } = this.#adapter.parseTableName(tableName);
+    return this.#adapter.isDatasetAllowed(schema);
+  }
 };
 
 // packages/text2sql/src/lib/adapters/groundings/view.grounding.ts