prazy1208/text2sql

1"""2Rule-based validation for generated SQL (Gen-SQL pipeline).3 4Goals:5- Single read-oriented statement (SELECT or WITH ... SELECT; optional EXPLAIN prefix).6- Reject obvious DML/DDL/session/control keywords via bounded scan (not a full SQL parser).7- Optional: referenced schema.table in FROM/JOIN must be subset of selected_tables when provided.8 9Returns a flat dict suitable for API and DB columns (bools + strings).10"""11 12from __future__ import annotations13 14import re15from typing import Any16 17# Whole-word matches on a comment-stripped, uppercased scan string.18_FORBIDDEN_RE = re.compile(19    r"\b("20    r"INSERT|UPDATE|DELETE|MERGE|TRUNCATE|"21    r"DROP|CREATE|ALTER|RENAME|REPLACE|"22    r"GRANT|REVOKE|COPY|"23    r"CALL|EXECUTE|EXEC|"24    r"LISTEN|NOTIFY|LOAD|CLUSTER|"25    r"VACUUM|REINDEX|REFRESH\s+MATERIALIZED\s+VIEW|"26    r"DISCARD|RESET\b"27    r")\b",28    re.IGNORECASE,29)30 31# SELECT INTO ... TABLE is DDL-like; allow plain INTO in SELECT lists only when not this pattern.32_SELECT_INTO_TABLE_RE = re.compile(33    r"\bINTO\s+(TEMPORARY|TEMP|UNLOGGED|TABLE)\b",34    re.IGNORECASE,35)36 37_FROM_JOIN_TABLE_RE = re.compile(38    r"(?:\bFROM|\bJOIN)\s+([a-z_][a-z0-9_]*\.[a-z_][a-z0-9_]*)\b",39    re.IGNORECASE,40)41 42 43def _mask_single_quoted_strings(sql: str) -> str:44    """45    Replace characters inside single-quoted string literals with spaces46    (PostgreSQL-style '' escape). Reduces false positives when scanning for keywords.47    """48    out: list[str] = []49    i = 050    n = len(sql)51    while i < n:52        ch = sql[i]53        if ch != "'":54            out.append(ch)55            i += 156            continue57        out.append("'")58        i += 159        while i < n:60            if sql[i] == "'":61                if i + 1 < n and sql[i + 1] == "'":62                    out.append("''")63                    i += 264                    continue65                out.append("'")66                i += 167                break68            out.append(" ")69            i += 170    return "".join(out)71 72 73def _strip_sql_comments(sql: str) -> str:74    """Remove /* */ and -- line comments (best-effort; does not handle quotes inside comments)."""75    s = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)76    lines: list[str] = []77    for line in s.splitlines():78        lines.append(line.split("--", 1)[0])79    return " ".join(lines)80 81 82def _trim_trailing_semicolons(sql: str) -> str:83    s = sql.strip()84    while s.endswith(";"):85        s = s[:-1].strip()86    return s87 88 89def _is_single_statement(sql: str) -> bool:90    s = _mask_single_quoted_strings(_strip_sql_comments(sql))91    s = _trim_trailing_semicolons(s)92    return ";" not in s93 94 95def _starts_with_select_family(sql: str) -> bool:96    """97    True if SQL (after comments) begins with EXPLAIN?, then WITH or SELECT.98 99    EXPLAIN options are stripped explicitly so ``EXPLAIN ANALYZE SELECT`` does not100    treat the main ``SELECT`` as an EXPLAIN keyword (``(\\s+\\w+)+`` would consume it).101    """102    head = _mask_single_quoted_strings(_strip_sql_comments(sql))103    head = head.strip()104    if re.match(r"^\s*EXPLAIN\b", head, re.IGNORECASE):105        head = re.sub(r"^\s*EXPLAIN\s+", "", head, count=1, flags=re.IGNORECASE).lstrip()106        # Known EXPLAIN / planner tokens only (repeat until none match).107        explain_opt = re.compile(108            r"^\s*("109            r"ANALYZE|VERBOSE|BUFFERS|TIMING|SUMMARY|WAL|SETTINGS|COSTS|"110            r"FORMAT\s+\w+"111            r")\s+",112            re.IGNORECASE,113        )114        while True:115            m = explain_opt.match(head)116            if not m:117                break118            head = head[m.end() :].lstrip()119    return bool(re.match(r"^\s*(WITH|SELECT)\b", head, re.IGNORECASE))120 121 122def _forbidden_hits(scan_upper: str) -> list[str]:123    hits: list[str] = []124    for m in _FORBIDDEN_RE.finditer(scan_upper):125        kw = re.sub(r"\s+", " ", m.group(1).upper()).strip()126        if kw not in hits:127            hits.append(kw)128    return hits129 130 131def _referenced_fqns_from_from_join(sql: str) -> list[str]:132    """Extract schema.table tokens immediately after FROM or JOIN (best-effort)."""133    s = _mask_single_quoted_strings(_strip_sql_comments(sql))134    found: list[str] = []135    seen: set[str] = set()136    for m in _FROM_JOIN_TABLE_RE.finditer(s):137        fqn = m.group(1).lower()138        if fqn not in seen:139            seen.add(fqn)140            found.append(fqn)141    return found142 143 144def validate_generated_sql(145    sql: str | None,146    *,147    selected_tables: list[str] | None = None,148    selected_columns: dict[str, list[str]] | None = None,149) -> dict[str, Any]:150    """151    Validate generated SQL for safe single-statement read queries.152 153    selected_columns is accepted for API symmetry; table/column-level proof154    without a parser is noisy, so only selected_tables participates in subset checks.155 156    Returns dict keys:157      validation_passed (bool)158      validation_error_codes (str)159      validation_error_message (str)160      blocked_keywords (str)161      is_single_statement (bool)162      is_select_only (bool)163    """164    raw = (sql or "").strip()165    codes: list[str] = []166    messages: list[str] = []167 168    if not raw:169        return {170            "validation_passed": False,171            "validation_error_codes": "EMPTY_SQL",172            "validation_error_message": "Generated SQL is empty.",173            "blocked_keywords": "",174            "is_single_statement": False,175            "is_select_only": False,176        }177 178    single = _is_single_statement(raw)179    if not single:180        codes.append("MULTIPLE_STATEMENTS")181        messages.append("Multiple statements are not allowed (found ';').")182 183    scan = _mask_single_quoted_strings(_strip_sql_comments(raw))184    scan_upper = scan.upper()185 186    if _SELECT_INTO_TABLE_RE.search(scan_upper):187        codes.append("SELECT_INTO_DDL")188        messages.append("SELECT INTO / CREATE TABLE forms are not allowed.")189 190    forbidden = _forbidden_hits(scan_upper)191    blocked_csv = ", ".join(forbidden) if forbidden else ""192    if forbidden:193        codes.append("FORBIDDEN_KEYWORD")194        messages.append(f"Disallowed keyword(s): {blocked_csv}.")195 196    select_family = _starts_with_select_family(raw)197    if not select_family:198        codes.append("NOT_SELECT")199        messages.append("Query must start with SELECT or WITH (optional EXPLAIN prefix).")200 201    # Optional: FROM/JOIN schema.table must be in selected_tables (lowercased set).202    if selected_tables:203        allowed = {t.strip().lower() for t in selected_tables if t and str(t).strip()}204        if allowed:205            refs = _referenced_fqns_from_from_join(raw)206            unknown = [r for r in refs if r not in allowed]207            if unknown:208                codes.append("TABLE_NOT_IN_SELECTION")209                messages.append(210                    "FROM/JOIN references table(s) not in the selected table list: "211                    + ", ".join(unknown)212                )213 214    # selected_columns reserved for future stricter checks215    _ = selected_columns216 217    passed = not codes218    select_into_hit = bool(_SELECT_INTO_TABLE_RE.search(scan_upper))219    is_select_shape = (220        single221        and select_family222        and not forbidden223        and not select_into_hit224    )225    return {226        "validation_passed": passed,227        "validation_error_codes": ";".join(codes) if codes else "",228        "validation_error_message": " ".join(messages).strip() if messages else "",229        "blocked_keywords": blocked_csv,230        "is_single_statement": single,231        "is_select_only": is_select_shape,232    }233