| """SQL safety: AST allow-list (SELECT-only, no DDL/DML, no risky functions).
|
|
|
| Uses sqlglot to parse and validate any LLM-generated SQL before execution.
|
| """
|
| from __future__ import annotations
|
|
|
| from dataclasses import dataclass
|
| from typing import Iterable
|
|
|
| import sqlglot
|
| from sqlglot import exp
|
|
|
| ALLOWED_TABLES_PREFIXES: tuple[str, ...] = ("raw.", "dim.", "sem.")
|
| DENY_FUNCTIONS = {"load_file", "sleep", "benchmark", "user", "version"}
|
| MAX_LIMIT = 5000
|
|
|
|
|
| @dataclass(frozen=True)
|
| class SqlValidation:
|
| ok: bool
|
| reason: str = ""
|
| rewritten: str = ""
|
|
|
|
|
| def validate_sql(sql: str, dialect: str = "postgres") -> SqlValidation:
|
| try:
|
| parsed = sqlglot.parse_one(sql, read=dialect)
|
| except Exception as e:
|
| return SqlValidation(False, f"parse_error: {e}")
|
|
|
| if not isinstance(parsed, (exp.Select, exp.Subqueryable, exp.With, exp.Union)):
|
| return SqlValidation(False, f"not_a_select: {type(parsed).__name__}")
|
|
|
|
|
| forbidden = (
|
| exp.Insert,
|
| exp.Update,
|
| exp.Delete,
|
| exp.Drop,
|
| exp.Create,
|
| exp.AlterTable,
|
| exp.TruncateTable,
|
| exp.Merge,
|
| )
|
| for node in parsed.walk():
|
| n = node[0] if isinstance(node, tuple) else node
|
| if isinstance(n, forbidden):
|
| return SqlValidation(False, f"forbidden_statement:{type(n).__name__}")
|
| if isinstance(n, exp.Anonymous) and (n.name or "").lower() in DENY_FUNCTIONS:
|
| return SqlValidation(False, f"forbidden_function:{n.name}")
|
|
|
|
|
| for tbl in parsed.find_all(exp.Table):
|
| fq = ".".join(p for p in [tbl.db, tbl.name] if p).lower()
|
| if not any(fq.startswith(p) or fq == p.rstrip(".") for p in ALLOWED_TABLES_PREFIXES):
|
|
|
| if tbl.name.lower() in {
|
| "extended_bva_v5", "bva", "item_mb", "open_po", "expense_mb",
|
| "opening_budget_fy_25_26", "opening_underspents", "training_blocked_budget",
|
| "funders", "projects", "categories", "vendors", "employees",
|
| }:
|
| continue
|
| return SqlValidation(False, f"forbidden_table:{fq or tbl.name}")
|
|
|
|
|
| if isinstance(parsed, exp.Select):
|
| lim = parsed.args.get("limit")
|
| if lim is None:
|
| parsed.set("limit", exp.Limit(expression=exp.Literal.number(MAX_LIMIT)))
|
| else:
|
| try:
|
| n = int(lim.expression.name)
|
| if n > MAX_LIMIT:
|
| parsed.set("limit", exp.Limit(expression=exp.Literal.number(MAX_LIMIT)))
|
| except Exception:
|
| pass
|
|
|
| return SqlValidation(True, rewritten=parsed.sql(dialect=dialect))
|
|
|
