@deepagents/text2sql 7.0.2 → 7.1.0

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

@@ -1,5 +1,6 @@
 // packages/text2sql/src/lib/adapters/adapter.ts
 import { format as formatSql } from "sql-formatter";
+import { z } from "zod";
 
 // packages/text2sql/src/lib/agents/exceptions.ts
 var sqlValidationMarker = /* @__PURE__ */ Symbol("SQLValidationError");
@@ -172,6 +173,16 @@ function createGroundingContext(options = {}) {
 }
 
 // packages/text2sql/src/lib/adapters/adapter.ts
+function groundingFor(adapterType, create) {
+  return (adapter) => {
+    if (!(adapter instanceof adapterType)) {
+      throw new Error(
+        `A ${adapterType.name} grounding needs a ${adapterType.name} adapter, not ${adapter.constructor.name}.`
+      );
+    }
+    return create(adapter);
+  };
+}
 var Adapter = class {
   #policyAnalyzer;
   /**
@@ -432,6 +443,40 @@ var Adapter = class {
     }
     return decoded;
   }
+  /**
+   * Run an introspection query and check every row against `row`, the shape
+   * the driver returns for that query. A row of any other shape throws.
+   */
+  async runQuery(sql, row) {
+    const rows = z.array(row).safeParse(await this.queryRows(sql));
+    if (!rows.success) {
+      throw new Error(
+        `${this.constructor.name} introspection query returned rows of an unexpected shape:
+${z.prettifyError(rows.error)}
+SQL: ${sql.trim()}`,
+        { cause: rows.error }
+      );
+    }
+    return rows.data;
+  }
+  /**
+   * The rows in a driver's execute() result: the result itself when it is an
+   * array, otherwise the first array found at one of `paths` (dot-separated,
+   * so `recordsets.0` reads `result.recordsets[0]`).
+   */
+  rowsFrom(result, paths) {
+    if (Array.isArray(result)) return result;
+    for (const path of paths) {
+      const rows = path.split(".").reduce(
+        (value, key) => typeof value === "object" && value !== null ? Reflect.get(value, key) : void 0,
+        result
+      );
+      if (Array.isArray(rows)) return rows;
+    }
+    throw new Error(
+      `${this.constructor.name} execute() must return an array of rows or an object holding one at ${paths.map((path) => `result.${path}`).join(", ")}.`
+    );
+  }
   /**
    * Convert unknown database value to number.
    * Handles number, bigint, and string types.
@@ -1141,7 +1186,7 @@ var BigQuerySqlPolicyAnalyzer = class extends ParserSqlPolicyAnalyzer {
 
 // packages/text2sql/src/lib/adapters/bigquery/bigquery.ts
 function formatBigQueryError(sql, error) {
-  const errorMessage = error instanceof Error ? error.message : typeof error === "string" ? error : typeof error === "object" && error !== null ? error.message ?? JSON.stringify(error) : "Unknown error occurred";
+  const errorMessage = error instanceof Error ? error.message : typeof error === "string" ? error : typeof error === "object" && error !== null ? "message" in error && typeof error.message === "string" ? error.message : JSON.stringify(error) : "Unknown error occurred";
   return {
     error: errorMessage,
     error_type: "BIGQUERY_ERROR",
@@ -1204,17 +1249,8 @@ var BigQuery = class extends Adapter {
       return JSON.stringify(formatBigQueryError(this.format(sql), error));
     }
   }
-  async runQuery(sql) {
-    const result = await this.#options.execute(sql);
-    if (Array.isArray(result)) {
-      return result;
-    }
-    if (result && typeof result === "object" && "rows" in result && Array.isArray(result.rows)) {
-      return result.rows;
-    }
-    throw new Error(
-      "BigQuery adapter execute() must return an array of rows or an object with a rows array when introspecting."
-    );
+  async queryRows(sql) {
+    return this.rowsFrom(await this.#options.execute(sql), ["rows"]);
   }
   quoteIdentifier(name) {
     return name.split(".").map((part) => `\`${part.replace(/`/g, "``")}\``).join(".");
@@ -1239,12 +1275,46 @@ var BigQuery = class extends Adapter {
   }
 };
 
+// packages/text2sql/src/lib/adapters/bigquery/constraint.bigquery.grounding.ts
+import { z as z4 } from "zod";
+
+// packages/text2sql/src/lib/adapters/bigquery/bigquery-fk.ts
+import { z as z3 } from "zod";
+
+// packages/text2sql/src/lib/adapters/bigquery/bigquery-rows.ts
+import { z as z2 } from "zod";
+var int64 = z2.number();
+var tableNameRow = z2.object({ table_name: z2.string().nullable() });
+var constraintNameRow = z2.object({
+  constraint_name: z2.string().nullable()
+});
+var keyColumnUsageRow = z2.object({
+  constraint_name: z2.string().nullable(),
+  column_name: z2.string().nullable(),
+  ordinal_position: int64.nullable(),
+  position_in_unique_constraint: int64.nullable()
+});
+
 // packages/text2sql/src/lib/adapters/bigquery/bigquery-fk.ts
+var referencedTableRow = z3.object({
+  table_schema: z3.string().nullable(),
+  table_name: z3.string().nullable()
+});
+var primaryKeyColumnRow = z3.object({
+  column_name: z3.string().nullable(),
+  ordinal_position: int64.nullable()
+});
+var fkResolution = z3.object({
+  referencedDataset: z3.string(),
+  referencedTable: z3.string(),
+  referencedColumns: z3.array(z3.string()),
+  childColumns: z3.array(z3.string())
+});
 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);
+    const cached = fkResolution.optional().parse(cache.get(cacheKey));
     if (!cached) return void 0;
     return {
       ...cached,
@@ -1261,11 +1331,14 @@ async function resolveForeignKey(adapter, constraintDataset, constraintName, chi
   return result;
 }
 async function resolveFK(adapter, constraintDataset, constraintName, childColumns) {
-  const refRows = await adapter.runQuery(`
+  const refRows = await adapter.runQuery(
+    `
     SELECT DISTINCT table_schema, table_name
     FROM ${adapter.infoSchemaView(constraintDataset, "CONSTRAINT_COLUMN_USAGE")}
     WHERE constraint_name = '${adapter.escapeString(constraintName)}'
-  `);
+  `,
+    referencedTableRow
+  );
   const referenced = refRows.find((r) => r.table_schema && r.table_name);
   if (!referenced?.table_schema || !referenced.table_name) {
     return void 0;
@@ -1275,22 +1348,28 @@ async function resolveFK(adapter, constraintDataset, constraintName, childColumn
   if (!adapter.isDatasetAllowed(referencedDataset)) {
     return void 0;
   }
-  const pkConstraintRows = await adapter.runQuery(`
+  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
-  `);
+  `,
+    constraintNameRow
+  );
   const pkConstraintName = pkConstraintRows[0]?.constraint_name;
   if (!pkConstraintName) return void 0;
-  const pkColumnRows = await adapter.runQuery(`
+  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
-  `);
+  `,
+    primaryKeyColumnRow
+  );
   const pkByOrdinal = /* @__PURE__ */ new Map();
   for (const row of pkColumnRows) {
     if (!row.column_name || row.ordinal_position == null) continue;
@@ -1309,6 +1388,16 @@ async function resolveFK(adapter, constraintDataset, constraintName, childColumn
 }
 
 // packages/text2sql/src/lib/adapters/bigquery/constraint.bigquery.grounding.ts
+var columnMetadataRow = z4.object({
+  table_name: z4.string().nullable(),
+  column_name: z4.string().nullable(),
+  is_nullable: z4.string().nullable(),
+  column_default: z4.string().nullable()
+});
+var tableKeyColumnRow = keyColumnUsageRow.extend({
+  table_name: z4.string().nullable(),
+  constraint_type: z4.string().nullable()
+});
 var BigQueryConstraintGrounding = class extends ConstraintGrounding {
   #adapter;
   #cache;
@@ -1367,12 +1456,15 @@ var BigQueryConstraintGrounding = class extends ConstraintGrounding {
     }
   }
   async #batchColumnMetadata(dataset, inList, constraintsByTable) {
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT table_name, column_name, is_nullable, column_default
       FROM ${this.#adapter.infoSchemaView(dataset, "COLUMNS")}
       WHERE table_name IN (${inList})
       ORDER BY table_name, ordinal_position
-    `);
+    `,
+      columnMetadataRow
+    );
     for (const row of rows) {
       if (!row.table_name || !row.column_name) continue;
       const constraints2 = constraintsByTable.get(row.table_name);
@@ -1395,7 +1487,8 @@ var BigQueryConstraintGrounding = class extends ConstraintGrounding {
     }
   }
   async #batchKeyConstraints(dataset, inList, constraintsByTable) {
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT
         tc.table_name,
         tc.constraint_name,
@@ -1410,7 +1503,9 @@ var BigQueryConstraintGrounding = class extends ConstraintGrounding {
       WHERE tc.table_name IN (${inList})
         AND tc.constraint_type IN ('PRIMARY KEY', 'FOREIGN KEY')
       ORDER BY tc.table_name, tc.constraint_name, kcu.ordinal_position
-    `);
+    `,
+      tableKeyColumnRow
+    );
     const pkByTable = /* @__PURE__ */ new Map();
     const fkByTable = /* @__PURE__ */ new Map();
     for (const row of rows) {
@@ -1477,6 +1572,13 @@ var BigQueryConstraintGrounding = class extends ConstraintGrounding {
 };
 
 // packages/text2sql/src/lib/adapters/bigquery/indexes.bigquery.grounding.ts
+import { z as z5 } from "zod";
+var indexHintRow = z5.object({
+  table_name: z5.string().nullable(),
+  column_name: z5.string().nullable(),
+  is_partitioning_column: z5.string().nullable(),
+  clustering_ordinal_position: int64.nullable()
+});
 var BigQueryIndexesGrounding = class extends IndexesGrounding {
   #adapter;
   constructor(adapter, config = {}) {
@@ -1516,13 +1618,16 @@ var BigQueryIndexesGrounding = class extends IndexesGrounding {
       (t) => this.#adapter.parseTableName(t.name).table
     );
     const inList = tableNames.map((n) => `'${this.#adapter.escapeString(n)}'`).join(", ");
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT table_name, column_name, is_partitioning_column, clustering_ordinal_position
       FROM ${this.#adapter.infoSchemaView(dataset, "COLUMNS")}
       WHERE table_name IN (${inList})
         AND (is_partitioning_column = 'YES' OR clustering_ordinal_position IS NOT NULL)
       ORDER BY table_name, clustering_ordinal_position
-    `);
+    `,
+      indexHintRow
+    );
     const byTable = /* @__PURE__ */ new Map();
     for (const row of rows) {
       if (!row.table_name || !row.column_name) continue;
@@ -1603,6 +1708,29 @@ var BigQueryInfoGrounding = class extends InfoGrounding {
 };
 
 // packages/text2sql/src/lib/adapters/bigquery/row-count.bigquery.grounding.ts
+import { z as z7 } from "zod";
+
+// packages/text2sql/src/lib/adapters/groundings/rows.ts
+import { z as z6 } from "zod";
+var nameRow = z6.object({ name: z6.string() });
+var versionRow = z6.object({ version: z6.string() });
+var databaseRow = z6.object({ db: z6.string() });
+var viewDefinitionRow = z6.object({
+  definition: z6.string().nullable()
+});
+var distinctValueRow = z6.object({ value: z6.unknown() });
+var textValueRow = z6.object({ value: z6.string() });
+var numericValue = z6.union([z6.number(), z6.bigint(), z6.string()]);
+
+// packages/text2sql/src/lib/adapters/bigquery/row-count.bigquery.grounding.ts
+var tableStorageRow = z7.object({
+  table_name: z7.string().nullable(),
+  total_rows: numericValue.nullable()
+});
+var legacyRowCountRow = z7.object({
+  table_name: z7.string().nullable(),
+  row_count: numericValue.nullable()
+});
 var BigQueryRowCountGrounding = class extends RowCountGrounding {
   #adapter;
   constructor(adapter, config = {}) {
@@ -1658,11 +1786,14 @@ var BigQueryRowCountGrounding = class extends RowCountGrounding {
     return /* @__PURE__ */ new Map();
   }
   async #fromTableStorage(dataset, inList) {
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT table_name, total_rows
       FROM ${this.#adapter.infoSchemaView(dataset, "TABLE_STORAGE")}
       WHERE table_name IN (${inList})
-    `);
+    `,
+      tableStorageRow
+    );
     const result = /* @__PURE__ */ new Map();
     for (const row of rows) {
       if (!row.table_name) continue;
@@ -1673,11 +1804,14 @@ var BigQueryRowCountGrounding = class extends RowCountGrounding {
   }
   async #fromLegacyTables(dataset, inList) {
     const projectPrefix = this.#adapter.projectId ? `\`${this.#adapter.projectId}\`.` : "";
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT table_id AS table_name, row_count
       FROM ${projectPrefix}\`${dataset}\`.__TABLES__
       WHERE table_id IN (${inList})
-    `);
+    `,
+      legacyRowCountRow
+    );
     const result = /* @__PURE__ */ new Map();
     for (const row of rows) {
       if (!row.table_name) continue;
@@ -1699,6 +1833,15 @@ var BigQueryRowCountGrounding = class extends RowCountGrounding {
 };
 
 // packages/text2sql/src/lib/adapters/bigquery/table.bigquery.grounding.ts
+import { z as z8 } from "zod";
+var columnFieldPathRow = z8.object({
+  field_path: z8.string().nullable(),
+  data_type: z8.string().nullable(),
+  ordinal_position: int64.nullable()
+});
+var childKeyColumnRow = keyColumnUsageRow.extend({
+  child_table_name: z8.string().nullable()
+});
 var BigQueryTableGrounding = class extends TableGrounding {
   #adapter;
   #cache;
@@ -1722,12 +1865,15 @@ var BigQueryTableGrounding = class extends TableGrounding {
   async getAllTableNames() {
     const names = [];
     for (const dataset of this.#adapter.datasets) {
-      const rows = await this.#adapter.runQuery(`
+      const rows = await this.#adapter.runQuery(
+        `
         SELECT table_name
         FROM ${this.#adapter.infoSchemaView(dataset, "TABLES")}
         WHERE table_type = 'BASE TABLE'
         ORDER BY table_name
-      `);
+      `,
+        tableNameRow
+      );
       for (const row of rows) {
         if (!row.table_name) continue;
         names.push(`${dataset}.${row.table_name}`);
@@ -1737,7 +1883,8 @@ var BigQueryTableGrounding = class extends TableGrounding {
   }
   async getTable(tableName) {
     const { schema: dataset, table: table2 } = this.#adapter.parseTableName(tableName);
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT
         f.field_path,
         f.data_type,
@@ -1748,7 +1895,9 @@ var BigQueryTableGrounding = class extends TableGrounding {
         AND f.column_name = c.column_name
       WHERE f.table_name = '${this.#adapter.escapeString(table2)}'
       ORDER BY c.ordinal_position, f.field_path
-    `);
+    `,
+      columnFieldPathRow
+    );
     const seen = /* @__PURE__ */ new Set();
     const columns = rows.map((r) => ({
       name: r.field_path ?? "unknown",
@@ -1769,7 +1918,8 @@ var BigQueryTableGrounding = class extends TableGrounding {
   }
   async findOutgoingRelations(tableName) {
     const { schema: dataset, table: table2 } = this.#adapter.parseTableName(tableName);
-    const rows = await this.#adapter.runQuery(`
+    const rows = await this.#adapter.runQuery(
+      `
       SELECT
         kcu.constraint_name,
         kcu.column_name,
@@ -1782,7 +1932,9 @@ var BigQueryTableGrounding = class extends TableGrounding {
       WHERE tc.constraint_type = 'FOREIGN KEY'
         AND tc.table_name = '${this.#adapter.escapeString(table2)}'
       ORDER BY kcu.constraint_name, kcu.ordinal_position
-    `);
+    `,
+      keyColumnUsageRow
+    );
     const byConstraint = /* @__PURE__ */ new Map();
     for (const row of rows) {
       if (!row.constraint_name || !row.column_name) continue;
@@ -1818,12 +1970,15 @@ var BigQueryTableGrounding = class extends TableGrounding {
     const { schema: referencedDataset, table: referencedTable } = this.#adapter.parseTableName(tableName);
     const rels = [];
     for (const constraintDataset of this.#adapter.datasets) {
-      const rows = await this.#adapter.runQuery(`
+      const rows = await this.#adapter.runQuery(
+        `
         SELECT DISTINCT constraint_name
         FROM ${this.#adapter.infoSchemaView(constraintDataset, "CONSTRAINT_COLUMN_USAGE")}
         WHERE table_schema = '${this.#adapter.escapeString(referencedDataset)}'
           AND table_name = '${this.#adapter.escapeString(referencedTable)}'
-      `);
+      `,
+        constraintNameRow
+      );
       for (const row of rows) {
         if (!row.constraint_name) continue;
         const rel = await this.#resolveIncomingRelationship(
@@ -1838,7 +1993,8 @@ var BigQueryTableGrounding = class extends TableGrounding {
     return rels;
   }
   async #resolveIncomingRelationship(constraintDataset, constraintName, expectedReferencedDataset, expectedReferencedTable) {
-    const keyRows = await this.#adapter.runQuery(`
+    const keyRows = await this.#adapter.runQuery(
+      `
       SELECT
         kcu.constraint_name,
         tc.table_name AS child_table_name,
@@ -1852,7 +2008,9 @@ var BigQueryTableGrounding = class extends TableGrounding {
       WHERE tc.constraint_type = 'FOREIGN KEY'
         AND tc.constraint_name = '${this.#adapter.escapeString(constraintName)}'
       ORDER BY kcu.ordinal_position
-    `);
+    `,
+      childKeyColumnRow
+    );
     if (keyRows.length === 0) return void 0;
     const childTable = keyRows[0]?.child_table_name;
     if (!childTable) return void 0;
@@ -1884,6 +2042,9 @@ var BigQueryTableGrounding = class extends TableGrounding {
   }
 };
 
+// packages/text2sql/src/lib/adapters/bigquery/view.bigquery.grounding.ts
+import { z as z9 } from "zod";
+
 // packages/text2sql/src/lib/adapters/groundings/view.grounding.ts
 var ViewGrounding = class extends AbstractGrounding {
   #filter;
@@ -1940,6 +2101,11 @@ var ViewGrounding = class extends AbstractGrounding {
 };
 
 // packages/text2sql/src/lib/adapters/bigquery/view.bigquery.grounding.ts
+var viewDefinitionRow2 = z9.object({ ddl: z9.string().nullable() });
+var columnRow = z9.object({
+  column_name: z9.string().nullable(),
+  data_type: z9.string().nullable()
+});
 var BigQueryViewGrounding = class extends ViewGrounding {
   #adapter;
   constructor(adapter, config = {}) {
@@ -1953,12 +2119,15 @@ var BigQueryViewGrounding = class extends ViewGrounding {
   async getAllViewNames() {
     const names = [];
     for (const dataset of this.#adapter.datasets) {
-      const rows = await this.#adapter.runQuery(`
+      const rows = await this.#adapter.runQuery(
+        `
         SELECT table_name
         FROM ${this.#adapter.infoSchemaView(dataset, "TABLES")}
         WHERE table_type IN ('VIEW', 'MATERIALIZED VIEW')
         ORDER BY table_name
-      `);
+      `,
+        tableNameRow
+      );
       for (const row of rows) {
         if (!row.table_name) continue;
         names.push(`${dataset}.${row.table_name}`);
@@ -1970,21 +2139,27 @@ var BigQueryViewGrounding = class extends ViewGrounding {
     const { schema: dataset, table: table2 } = this.#adapter.parseTableName(viewName);
     let definition;
     if (this.includeDefinition) {
-      const defRows = await this.#adapter.runQuery(`
+      const defRows = await this.#adapter.runQuery(
+        `
         SELECT ddl
         FROM ${this.#adapter.infoSchemaView(dataset, "TABLES")}
         WHERE table_name = '${this.#adapter.escapeString(table2)}'
           AND table_type IN ('VIEW', 'MATERIALIZED VIEW')
         LIMIT 1
-      `);
+      `,
+        viewDefinitionRow2
+      );
       definition = defRows[0]?.ddl ?? void 0;
     }
-    const columns = await this.#adapter.runQuery(`
+    const columns = await this.#adapter.runQuery(
+      `
       SELECT column_name, data_type
       FROM ${this.#adapter.infoSchemaView(dataset, "COLUMNS")}
       WHERE table_name = '${this.#adapter.escapeString(table2)}'
       ORDER BY ordinal_position
-    `);
+    `,
+      columnRow
+    );
     return {
       name: `${dataset}.${table2}`,
       schema: dataset,
@@ -2004,22 +2179,40 @@ var BigQueryViewGrounding = class extends ViewGrounding {
 
 // packages/text2sql/src/lib/adapters/bigquery/index.ts
 function tables(config = {}) {
-  return (adapter) => new BigQueryTableGrounding(adapter, config);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryTableGrounding(adapter, config)
+  );
 }
 function info(config = {}) {
-  return (adapter) => new BigQueryInfoGrounding(adapter);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryInfoGrounding(adapter)
+  );
 }
 function views(config = {}) {
-  return (adapter) => new BigQueryViewGrounding(adapter, config);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryViewGrounding(adapter, config)
+  );
 }
 function indexes(config = {}) {
-  return (adapter) => new BigQueryIndexesGrounding(adapter, config);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryIndexesGrounding(adapter, config)
+  );
 }
 function rowCount(config = {}) {
-  return (adapter) => new BigQueryRowCountGrounding(adapter, config);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryRowCountGrounding(adapter, config)
+  );
 }
 function constraints(config = {}) {
-  return (adapter) => new BigQueryConstraintGrounding(adapter, config);
+  return groundingFor(
+    BigQuery,
+    (adapter) => new BigQueryConstraintGrounding(adapter, config)
+  );
 }
 var bigquery_default = {
   tables,