recon + sqlglot validator + drill_down package; guard ReAct dimension picks against candidate list

api/recon/validate.py

@@ -0,0 +1,76 @@
+"""SQL validation against the recon schema.
+
+`validate_sql` parses the SQL, runs sqlglot's `qualify` optimizer pass with
+the dataset's column→type schema, and raises with a clear message when the
+LLM references a column that doesn't exist on the table it's bound to.
+
+This is the deterministic complement to the LLM prompt: prompt hints help
+the LLM get it right; the validator GUARANTEES we don't ship a schema-
+wrong query to Postgres. If validation fails, the caller can either bubble
+the error up (no silent retry) or do one *guided* re-prompt that includes
+the validator's message — which is a correction with concrete facts, not
+a blind retry.
+"""
+
+from __future__ import annotations
+
+import sqlglot
+from sqlglot.errors import OptimizeError
+from sqlglot.optimizer.qualify import qualify
+
+from api.recon import load_recon
+from api.recon.types import Recon
+
+
+class ReconValidationError(ValueError):
+    """The SQL references a table/column combination that doesn't exist
+    in the recon. The original sqlglot error is in `__cause__`."""
+
+
+def _build_sqlglot_schema(recon: Recon) -> dict[str, dict[str, str]]:
+    """Convert recon → sqlglot's expected schema shape: {table: {col: type}}."""
+    return {
+        name: {col.name: col.sql_type for col in t.columns}
+        for name, t in recon.tables.items()
+    }
+
+
+def validate_sql(sql: str, recon: Recon | None = None) -> None:
+    """Raise `ReconValidationError` if any column reference in `sql` doesn't
+    exist on the table it's bound to. Returns None on success."""
+    recon = recon or load_recon()
+    schema = _build_sqlglot_schema(recon)
+    try:
+        parsed = sqlglot.parse_one(sql, dialect="postgres")
+        qualify(parsed, schema=schema, dialect="postgres")
+    except OptimizeError as e:
+        # Try to enrich the message with hints from the recon.
+        msg = str(e)
+        hint = _column_hint(msg, recon)
+        full = f"{msg}{hint}" if hint else msg
+        raise ReconValidationError(full) from e
+
+
+_COL_PATTERNS = [
+    r"Column '([^']+)'",                       # "Column 'X' could not be resolved."
+    r"Unknown column:\s*\"?([^\"\s,]+)\"?",  # "Unknown column: X"
+    r"column \"([^\"]+)\"",                   # 'column "X" does not exist'
+]
+
+
+def _column_hint(msg: str, recon: Recon) -> str:
+    """If the error names a specific column, append a hint about which
+    table(s) actually own it. Tries the wordings sqlglot and psycopg use."""
+    import re
+    col_name: str | None = None
+    for pat in _COL_PATTERNS:
+        m = re.search(pat, msg)
+        if m:
+            col_name = m.group(1)
+            break
+    if col_name is None:
+        return ""
+    owners = recon.owning_tables(col_name)
+    if owners:
+        return f"  (Column '{col_name}' actually lives in: {', '.join(owners)}.)"
+    return f"  (No table in the {recon.schema!r} schema has a column named '{col_name}'.)"