text2sql-demo / src /app /sql_validator.py
minimew's picture
Upload src/app/sql_validator.py with huggingface_hub
fee3c9f verified
Raw
History Blame Contribute Delete
1.49 kB
"""
SQL Validator -- chi cho phep cau SELECT (read-only), chan DROP/DELETE/INSERT/
UPDATE/ALTER/... Day la lop bao ve AN TOAN (khong cho sinh lenh nguy hiem),
KHONG kiem tra SQL co tra loi DUNG cau hoi khong (model van co the sinh SELECT
hop le nhung sai logic -- xem demo_app_pipeline_notes.md).
"""
import re
import sqlglot
import sqlglot.expressions as exp
_FORBIDDEN_KEYWORDS = re.compile(
r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|ATTACH|DETACH|"
r"PRAGMA|REPLACE|VACUUM|REINDEX)\b",
re.IGNORECASE,
)
def validate_sql(sql: str):
"""Return (is_safe: bool, message: str). Chi cho 1 cau SELECT/UNION duy nhat."""
if not sql or not sql.strip():
return False, "Model generated an empty SQL string."
stripped = sql.strip().rstrip(";").strip()
# Chan nhieu cau lenh xep chuoi (stacked queries qua ';') -- phong thu
# ngay ca khi sqlglot co the parse duoc 1 phan cua chuoi.
if ";" in stripped:
return False, "Blocked: multiple statements separated by ';' are not allowed."
if _FORBIDDEN_KEYWORDS.search(stripped):
return False, "Blocked: only read-only SELECT queries are allowed."
try:
tree = sqlglot.parse_one(stripped, read="sqlite")
except Exception as e:
return False, f"Generated SQL failed to parse: {e}"
if not isinstance(tree, (exp.Select, exp.Union)):
return False, f"Blocked: only SELECT statements are allowed (got {type(tree).__name__})."
return True, "OK"