77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
"""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}'.)"
|