Spaces:
Build error
Build error
| """ | |
| Validates LLM-generated SQL before it touches the database. | |
| Rules enforced: | |
| - Exactly one statement, and it must be a SELECT. | |
| - No DDL/DML keywords anywhere (belt and suspenders alongside the read-only | |
| DB connection used in app.py). | |
| - Every table referenced must be in the whitelist. | |
| - A LIMIT is added if the query doesn't already have one and isn't an | |
| aggregate-only query. | |
| """ | |
| import re | |
| ALLOWED_TABLES = { | |
| "devices", | |
| "interfaces", | |
| "traffic_flows", | |
| "alerts", | |
| "metrics_timeseries", | |
| } | |
| FORBIDDEN_KEYWORDS = { | |
| "insert", "update", "delete", "drop", "alter", "create", | |
| "truncate", "replace", "attach", "detach", "pragma", "vacuum", | |
| } | |
| DEFAULT_LIMIT = 50 | |
| class SQLGuardrailError(Exception): | |
| pass | |
| def _extract_tables(sql: str) -> set: | |
| # crude but effective for this schema: look for tokens after FROM/JOIN | |
| tables = set() | |
| for m in re.finditer(r"\b(?:from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)", sql, re.IGNORECASE): | |
| tables.add(m.group(1).lower()) | |
| return tables | |
| def validate_and_sanitize(sql: str) -> str: | |
| """Raises SQLGuardrailError on anything unsafe. Returns a safe SQL string.""" | |
| if sql is None: | |
| raise SQLGuardrailError("Model returned no SQL.") | |
| cleaned = sql.strip().strip("`").strip() | |
| # strip a leading "sql" language tag if the model added a fenced block anyway | |
| cleaned = re.sub(r"^sql\s*", "", cleaned, flags=re.IGNORECASE) | |
| cleaned = cleaned.rstrip(";").strip() | |
| if not cleaned: | |
| raise SQLGuardrailError("Empty SQL after cleaning.") | |
| if ";" in cleaned: | |
| raise SQLGuardrailError("Multiple statements are not allowed.") | |
| lowered = cleaned.lower() | |
| if not lowered.startswith("select"): | |
| raise SQLGuardrailError("Only SELECT statements are allowed.") | |
| tokens = set(re.findall(r"[a-zA-Z_]+", lowered)) | |
| forbidden_hit = tokens & FORBIDDEN_KEYWORDS | |
| if forbidden_hit: | |
| raise SQLGuardrailError(f"Forbidden keyword(s) detected: {forbidden_hit}") | |
| referenced = _extract_tables(cleaned) | |
| unknown = referenced - ALLOWED_TABLES | |
| if unknown: | |
| raise SQLGuardrailError(f"Unknown table(s) referenced: {unknown}") | |
| if "limit" not in lowered: | |
| cleaned = f"{cleaned} LIMIT {DEFAULT_LIMIT}" | |
| return cleaned + ";" | |