@deepagents/text2sql 0.3.1 → 0.7.0

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

@@ -10,9 +10,8 @@ function createGroundingContext() {
 
 // packages/text2sql/src/lib/adapters/adapter.ts
 var Adapter = class {
-  async introspect() {
+  async introspect(ctx = createGroundingContext()) {
     const lines = [];
-    const ctx = createGroundingContext();
     for (const fn of this.grounding) {
       const grounding = fn(this);
       lines.push({
@@ -82,7 +81,7 @@ ${description}
   }
 };
 
-// packages/text2sql/src/lib/adapters/grounding.ticket.ts
+// packages/text2sql/src/lib/adapters/groundings/abstract.grounding.ts
 var AbstractGrounding = class {
   tag;
   constructor(tag) {
@@ -196,31 +195,108 @@ var InfoGrounding = class extends AbstractGrounding {
   }
 };
 
-// packages/text2sql/src/lib/adapters/groundings/low-cardinality.grounding.ts
-var LowCardinalityGrounding = class extends AbstractGrounding {
+// packages/text2sql/src/lib/adapters/groundings/column-values.grounding.ts
+var ColumnValuesGrounding = class extends AbstractGrounding {
+  lowCardinalityLimit;
   constructor(config = {}) {
-    super("low_cardinality");
+    super("column_values");
+    this.lowCardinalityLimit = config.lowCardinalityLimit ?? 20;
+  }
+  /**
+   * Get values for native ENUM type columns.
+   * Return undefined if column is not an ENUM type.
+   * Default implementation returns undefined (no native ENUM support).
+   */
+  async collectEnumValues(_tableName, _column) {
+    return void 0;
+  }
+  /**
+   * Parse CHECK constraint for enum-like IN clause.
+   * Extracts values from patterns like:
+   * - CHECK (status IN ('active', 'inactive'))
+   * - CHECK ((status)::text = ANY (ARRAY['a'::text, 'b'::text]))
+   * - CHECK (status = 'active' OR status = 'inactive')
+   */
+  parseCheckConstraint(constraint, columnName) {
+    if (constraint.type !== "CHECK" || !constraint.definition) {
+      return void 0;
+    }
+    if (constraint.columns && !constraint.columns.includes(columnName)) {
+      return void 0;
+    }
+    const def = constraint.definition;
+    const escapedCol = this.escapeRegex(columnName);
+    const colPattern = `(?:\\(?\\(?${escapedCol}\\)?(?:::(?:text|varchar|character varying))?\\)?)`;
+    const inMatch = def.match(
+      new RegExp(`${colPattern}\\s+IN\\s*\\(([^)]+)\\)`, "i")
+    );
+    if (inMatch) {
+      return this.extractStringValues(inMatch[1]);
+    }
+    const anyMatch = def.match(
+      new RegExp(
+        `${colPattern}\\s*=\\s*ANY\\s*\\(\\s*(?:ARRAY)?\\s*\\[([^\\]]+)\\]`,
+        "i"
+      )
+    );
+    if (anyMatch) {
+      return this.extractStringValues(anyMatch[1]);
+    }
+    const orPattern = new RegExp(
+      `\\b${this.escapeRegex(columnName)}\\b\\s*=\\s*'([^']*)'`,
+      "gi"
+    );
+    const orMatches = [...def.matchAll(orPattern)];
+    if (orMatches.length >= 2) {
+      return orMatches.map((m) => m[1]);
+    }
+    return void 0;
+  }
+  /**
+   * Extract string values from a comma-separated list.
+   */
+  extractStringValues(input) {
+    const values = [];
+    const matches = input.matchAll(/'([^']*)'/g);
+    for (const match of matches) {
+      values.push(match[1]);
+    }
+    return values.length > 0 ? values : void 0;
+  }
+  /**
+   * Escape special regex characters in a string.
+   */
+  escapeRegex(str) {
+    return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+  }
+  /**
+   * Get the table from context by name.
+   */
+  getTable(ctx, name) {
+    return ctx.tables.find((t) => t.name === name);
   }
   /**
    * Execute the grounding process.
-   * Annotates columns in ctx.tables and ctx.views with low cardinality values.
+   * Annotates columns in ctx.tables and ctx.views with values.
    */
   async execute(ctx) {
     const allContainers = [...ctx.tables, ...ctx.views];
     for (const container of allContainers) {
+      const table = this.getTable(ctx, container.name);
       for (const column of container.columns) {
         try {
-          const lowCard = await this.collectLowCardinality(
+          const result = await this.resolveColumnValues(
             container.name,
-            column
+            column,
+            table?.constraints
           );
-          if (lowCard) {
-            column.kind = lowCard.kind;
-            column.values = lowCard.values;
+          if (result) {
+            column.kind = result.kind;
+            column.values = result.values;
           }
         } catch (error) {
           console.warn(
-            "Error collecting low cardinality values for",
+            "Error collecting column values for",
             container.name,
             column.name,
             error
@@ -230,6 +306,28 @@ var LowCardinalityGrounding = class extends AbstractGrounding {
     }
     return () => this.#describe();
   }
+  /**
+   * Resolve column values from all sources in priority order.
+   */
+  async resolveColumnValues(tableName, column, constraints2) {
+    const enumValues = await this.collectEnumValues(tableName, column);
+    if (enumValues?.length) {
+      return { kind: "Enum", values: enumValues };
+    }
+    if (constraints2) {
+      for (const constraint of constraints2) {
+        const checkValues = this.parseCheckConstraint(constraint, column.name);
+        if (checkValues?.length) {
+          return { kind: "Enum", values: checkValues };
+        }
+      }
+    }
+    const lowCardValues = await this.collectLowCardinality(tableName, column);
+    if (lowCardValues?.length) {
+      return { kind: "LowCardinality", values: lowCardValues };
+    }
+    return void 0;
+  }
   #describe() {
     return null;
   }
@@ -478,18 +576,61 @@ var TableGrounding = class extends AbstractGrounding {
     }
     return tables2.map((table) => {
       const rowCountInfo = table.rowCount != null ? ` [rows: ${table.rowCount}${table.sizeHint ? `, size: ${table.sizeHint}` : ""}]` : "";
-      const pkConstraint = table.constraints?.find((c) => c.type === "PRIMARY_KEY");
+      const pkConstraint = table.constraints?.find(
+        (c) => c.type === "PRIMARY_KEY"
+      );
       const pkColumns = new Set(pkConstraint?.columns ?? []);
+      const notNullColumns = new Set(
+        table.constraints?.filter((c) => c.type === "NOT_NULL").flatMap((c) => c.columns ?? []) ?? []
+      );
+      const defaultByColumn = /* @__PURE__ */ new Map();
+      for (const c of table.constraints?.filter(
+        (c2) => c2.type === "DEFAULT"
+      ) ?? []) {
+        for (const col of c.columns ?? []) {
+          if (c.defaultValue != null) {
+            defaultByColumn.set(col, c.defaultValue);
+          }
+        }
+      }
+      const uniqueColumns = new Set(
+        table.constraints?.filter((c) => c.type === "UNIQUE" && c.columns?.length === 1).flatMap((c) => c.columns ?? []) ?? []
+      );
+      const fkByColumn = /* @__PURE__ */ new Map();
+      for (const c of table.constraints?.filter(
+        (c2) => c2.type === "FOREIGN_KEY"
+      ) ?? []) {
+        const cols = c.columns ?? [];
+        const refCols = c.referencedColumns ?? [];
+        for (let i = 0; i < cols.length; i++) {
+          const refCol = refCols[i] ?? refCols[0] ?? cols[i];
+          fkByColumn.set(cols[i], `${c.referencedTable}.${refCol}`);
+        }
+      }
       const columns = table.columns.map((column) => {
         const annotations = [];
         const isPrimaryKey = pkColumns.has(column.name);
         if (isPrimaryKey) {
           annotations.push("PK");
         }
+        if (fkByColumn.has(column.name)) {
+          annotations.push(`FK -> ${fkByColumn.get(column.name)}`);
+        }
+        if (uniqueColumns.has(column.name)) {
+          annotations.push("UNIQUE");
+        }
+        if (notNullColumns.has(column.name)) {
+          annotations.push("NOT NULL");
+        }
+        if (defaultByColumn.has(column.name)) {
+          annotations.push(`DEFAULT: ${defaultByColumn.get(column.name)}`);
+        }
         if (column.isIndexed && !isPrimaryKey) {
           annotations.push("Indexed");
         }
-        if (column.kind === "LowCardinality" && column.values?.length) {
+        if (column.kind === "Enum" && column.values?.length) {
+          annotations.push(`Enum: ${column.values.join(", ")}`);
+        } else if (column.kind === "LowCardinality" && column.values?.length) {
           annotations.push(`LowCardinality: ${column.values.join(", ")}`);
         }
         if (column.stats) {
@@ -524,9 +665,19 @@ ${table.indexes.map((index) => {
         const columnsText = index.columns?.length ? index.columns.join(", ") : "expression";
         return `    - ${index.name}${propsText}: ${columnsText}`;
       }).join("\n")}` : "";
+      const multiColumnUniques = table.constraints?.filter(
+        (c) => c.type === "UNIQUE" && (c.columns?.length ?? 0) > 1
+      ) ?? [];
+      const uniqueConstraints = multiColumnUniques.length ? `
+  Unique Constraints:
+${multiColumnUniques.map((c) => `    - ${c.name}: (${c.columns?.join(", ")})`).join("\n")}` : "";
+      const checkConstraints = table.constraints?.filter((c) => c.type === "CHECK") ?? [];
+      const checks = checkConstraints.length ? `
+  Check Constraints:
+${checkConstraints.map((c) => `    - ${c.name}: ${c.definition}`).join("\n")}` : "";
       return `- Table: ${table.name}${rowCountInfo}
   Columns:
-${columns}${indexes2}`;
+${columns}${indexes2}${uniqueConstraints}${checks}`;
     }).join("\n\n");
   }
   #formatTableLabel = (tableName) => {
@@ -613,7 +764,7 @@ var PostgresColumnStatsGrounding = class extends ColumnStatsGrounding {
       return false;
     }
     const normalized = type.toLowerCase();
-    return /int|real|numeric|double|float|decimal|date|time|bool|serial/.test(
+    return /int|real|numeric|double|float|decimal|date|time|serial/.test(
       normalized
     );
   }
@@ -808,19 +959,71 @@ var PostgresInfoGrounding = class extends InfoGrounding {
   }
 };
 
-// packages/text2sql/src/lib/adapters/postgres/low-cardinality.postgres.grounding.ts
-var LOW_CARDINALITY_LIMIT = 20;
-var PostgresLowCardinalityGrounding = class extends LowCardinalityGrounding {
+// packages/text2sql/src/lib/adapters/postgres/column-values.postgres.grounding.ts
+var PostgresColumnValuesGrounding = class extends ColumnValuesGrounding {
   #adapter;
+  #enumCache = /* @__PURE__ */ new Map();
+  #enumCacheLoaded = false;
   constructor(adapter, config = {}) {
     super(config);
     this.#adapter = adapter;
   }
+  /**
+   * Load all ENUM types and their values into cache.
+   * This is more efficient than querying per-column.
+   */
+  async #loadEnumCache() {
+    if (this.#enumCacheLoaded) {
+      return;
+    }
+    const rows = await this.#adapter.runQuery(`
+      SELECT
+        t.typname AS type_name,
+        n.nspname AS type_schema,
+        e.enumlabel AS enum_value
+      FROM pg_type t
+      JOIN pg_enum e ON t.oid = e.enumtypid
+      JOIN pg_namespace n ON n.oid = t.typnamespace
+      ORDER BY t.typname, e.enumsortorder
+    `);
+    for (const row of rows) {
+      const key = `${row.type_schema}.${row.type_name}`;
+      const existing = this.#enumCache.get(key) ?? [];
+      existing.push(row.enum_value);
+      this.#enumCache.set(key, existing);
+      const simpleKey = row.type_name;
+      const simpleExisting = this.#enumCache.get(simpleKey) ?? [];
+      simpleExisting.push(row.enum_value);
+      this.#enumCache.set(simpleKey, simpleExisting);
+    }
+    this.#enumCacheLoaded = true;
+  }
+  async collectEnumValues(tableName, column) {
+    if (column.type.toLowerCase() !== "user-defined") {
+      return void 0;
+    }
+    await this.#loadEnumCache();
+    const { schema, table } = this.#adapter.parseTableName(tableName);
+    const rows = await this.#adapter.runQuery(`
+      SELECT udt_name, udt_schema
+      FROM information_schema.columns
+      WHERE table_schema = '${this.#adapter.escapeString(schema)}'
+        AND table_name = '${this.#adapter.escapeString(table)}'
+        AND column_name = '${this.#adapter.escapeString(column.name)}'
+    `);
+    if (!rows.length) {
+      return void 0;
+    }
+    const { udt_name, udt_schema } = rows[0];
+    const fullKey = `${udt_schema}.${udt_name}`;
+    const values = this.#enumCache.get(fullKey) ?? this.#enumCache.get(udt_name);
+    return values?.length ? values : void 0;
+  }
   async collectLowCardinality(tableName, column) {
     const { schema, table } = this.#adapter.parseTableName(tableName);
     const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table)}`;
     const columnIdentifier = this.#adapter.quoteIdentifier(column.name);
-    const limit = LOW_CARDINALITY_LIMIT + 1;
+    const limit = this.lowCardinalityLimit + 1;
     const sql = `
       SELECT DISTINCT ${columnIdentifier}::text AS value
       FROM ${tableIdentifier}
@@ -828,7 +1031,7 @@ var PostgresLowCardinalityGrounding = class extends LowCardinalityGrounding {
       LIMIT ${limit}
     `;
     const rows = await this.#adapter.runQuery(sql);
-    if (!rows.length || rows.length > LOW_CARDINALITY_LIMIT) {
+    if (!rows.length || rows.length > this.lowCardinalityLimit) {
       return void 0;
     }
     const values = [];
@@ -838,10 +1041,7 @@ var PostgresLowCardinalityGrounding = class extends LowCardinalityGrounding {
       }
       values.push(row.value);
     }
-    if (!values.length) {
-      return void 0;
-    }
-    return { kind: "LowCardinality", values };
+    return values.length ? values : void 0;
   }
 };
 
@@ -868,7 +1068,7 @@ var POSTGRES_ERROR_MAP = {
     hint: "The function or operator you used is not recognized. Confirm its name and argument types."
   }
 };
-var LOW_CARDINALITY_LIMIT2 = 20;
+var LOW_CARDINALITY_LIMIT = 20;
 function isPostgresError(error) {
   return typeof error === "object" && error !== null && "code" in error && typeof error.code === "string";
 }
@@ -1168,7 +1368,7 @@ var Postgres = class extends Adapter {
       });
       for (const column of table.columns) {
         const columnIdentifier = this.#quoteIdentifier(column.name);
-        const limit = LOW_CARDINALITY_LIMIT2 + 1;
+        const limit = LOW_CARDINALITY_LIMIT + 1;
         const sql = `
           SELECT DISTINCT ${columnIdentifier} AS value
           FROM ${tableIdentifier}
@@ -1181,7 +1381,7 @@ var Postgres = class extends Adapter {
         } catch {
           continue;
         }
-        if (!rows.length || rows.length > LOW_CARDINALITY_LIMIT2) {
+        if (!rows.length || rows.length > LOW_CARDINALITY_LIMIT) {
           continue;
         }
         const values = [];
@@ -1594,9 +1794,9 @@ function columnStats(config = {}) {
     return new PostgresColumnStatsGrounding(adapter, config);
   };
 }
-function lowCardinality(config = {}) {
+function columnValues(config = {}) {
   return (adapter) => {
-    return new PostgresLowCardinalityGrounding(adapter, config);
+    return new PostgresColumnValuesGrounding(adapter, config);
   };
 }
 function indexes(config = {}) {
@@ -1622,7 +1822,7 @@ var postgres_default = {
   info,
   views,
   columnStats,
-  lowCardinality,
+  columnValues,
   indexes,
   rowCount,
   constraints,
@@ -1632,12 +1832,12 @@ var postgres_default = {
 export {
   Postgres,
   columnStats,
+  columnValues,
   constraints,
   postgres_default as default,
   formatPostgresError,
   indexes,
   info,
-  lowCardinality,
   report,
   rowCount,
   tables,