# coding=utf-8 """ @Time :2026/6/24 @Author :tao.chen SQL safety guard for PySpark job submissions. Rule: PySpark source submitted to this server may only contain SELECT and INSERT statements. Anything else (DROP, DELETE, UPDATE, TRUNCATE, ALTER, CREATE, REPLACE, MERGE, GRANT, REVOKE, etc.) is rejected with a clear error. How it works: 1. Extract string literals from the Python source that look like SQL (literal starts with a SQL keyword). This catches `spark.sql("...")`, `f"SELECT * FROM {x}"` (literal still starts with SELECT), and any other string with SQL-looking content. 2. For each extracted literal, sqlparse splits the text into statements and we look at each statement's first keyword. 3. WITH (CTE) is allowed at the top level, but we recurse into the CTE body to check the first DML verb. This catches the obvious `WITH cte AS (DROP TABLE x) DROP TABLE y`. Limitations (documented; not security theater): - f-strings where the SQL itself is built at runtime can't be fully validated. `f"SELECT * FROM {user_input}"` is treated as a SELECT because the literal starts with SELECT. If the substituted value itself is dangerous, that bypasses the guard. This is a fundamental limitation of static analysis on f-strings. - `pyspark.sql.functions.expr("...")` accepts SQL expressions inline and is NOT caught by this regex. (Future work.) """ import re import sqlparse from sqlparse.sql import Statement from sqlparse.tokens import Keyword, DDL, DML # --- Policy: what is allowed? --- # Top-level allowed verbs. WITH is handled separately so we can recurse # into CTE bodies to check the inner statement kind. ALLOWED_TOP_LEVEL = frozenset({"SELECT", "INSERT", "WITH"}) # All DML/DDL kinds we want to flag. The list is explicit (not "anything # not in ALLOWED") so adding new ops later is a conscious decision. FORBIDDEN_KINDS = frozenset({ # DDL "DROP", "CREATE", "ALTER", "TRUNCATE", "RENAME", "COMMENT", "REPLACE", "MERGE", "GRANT", "REVOKE", "USE", # DML "DELETE", "UPDATE", # Other "SET", "SHOW", "KILL", "EXEC", "EXECUTE", "CALL", "EXPLAIN", }) # --- Heuristic: which string literals look like SQL? --- # Match Python string literals. Verbatim so escapes work. (?:\\.|(?!(?P=q)).)* # matches an escaped char OR any char that is not the matching quote. _LITERAL_RE = re.compile( r'''(?P['"]) (?P(?:\\.|(?!(?P=q)).)*) (?P=q) ''', re.DOTALL | re.VERBOSE, ) # Strip leading whitespace and a possible "f"/"r"/"b" prefix on the literal # before peeking at the first keyword. (We don't need to interpolate f-strings # here — we only care about the literal text.) _PREFIX_RE = re.compile(r"""^[fFrRbBuU]*['"]""") _SQL_KEYWORD_AT_START = re.compile( r"^\s*(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE|MERGE|" r"REPLACE|WITH|GRANT|REVOKE|EXEC|EXECUTE|KILL|SET|USE|SHOW|RENAME|COMMENT|CALL|EXPLAIN)\b", re.IGNORECASE, ) def extract_sql_literals(py_code: str) -> list[str]: """Return the SQL body of every string literal in `py_code` whose first keyword is a SQL verb. False positives (English sentences starting with 'select') are acceptable for safety; false negatives (f-strings where the body is built at runtime) are documented limitations. """ out = [] for m in _LITERAL_RE.finditer(py_code): body = m.group("body") if _SQL_KEYWORD_AT_START.match(body): out.append(body) return out def _first_keyword(stmt: Statement) -> str | None: for tok in stmt.flatten(): if tok.ttype is Keyword or tok.ttype is DML or tok.ttype is DDL: v = tok.normalized.upper() if v and v != "AS": # "AS" inside WITH ... AS ( ... ) is structural, not a verb return v return None def _check_cte_body(stmt: Statement) -> list[str]: """If `stmt` is a WITH ... statement, find the inner verb of the CTE and check that against the policy. Returns a list of forbidden kinds. """ # After WITH, the body has its own first DML/DDL keyword. sqlparse # flattens tokens in order, so scan past the WITH/CTE-name/AS tokens # and find the next DML/DDL. saw_with = False paren_depth = 0 for tok in stmt.flatten(): if tok.ttype is Keyword and tok.normalized.upper() == "WITH": saw_with = True continue if not saw_with: continue if tok.ttype in (Keyword.DML, Keyword.DDL) or tok.ttype is DML or tok.ttype is DDL: inner = tok.normalized.upper() if inner in FORBIDDEN_KINDS: return [f"WITH ... {inner}"] if inner in ALLOWED_TOP_LEVEL: return [] # Unknown verb; let the top-level check decide return [] return [] def find_forbidden_statements(sql_text: str) -> list[str]: """Parse `sql_text` and return the list of forbidden statement kinds. A statement is "forbidden" if its first verb is in FORBIDDEN_KINDS (DROP, DELETE, etc.) or if a WITH ... CTE is found. Returns [] when every statement is SELECT/INSERT/WITH-allowed. """ offenses: list[str] = [] parsed = sqlparse.parse(sql_text) for stmt in parsed: if not isinstance(stmt, Statement): continue # Skip empty statements (trailing semicolons, comments-only) if stmt.token_first(skip_cm=True) is None: continue first = _first_keyword(stmt) if first is None: continue if first in ALLOWED_TOP_LEVEL: if first == "WITH": offenses.extend(_check_cte_body(stmt)) continue if first in FORBIDDEN_KINDS: offenses.append(first) else: # Unknown verb (e.g. SHOW, SET, USE) — be conservative and # reject. Operators can extend the allow-list if needed. offenses.append(first) return offenses def validate_pyspark_code(py_code: str) -> list[str]: """Return a list of forbidden SQL statement kinds found in `py_code`. Empty list means the code is policy-compliant (or contains no SQL). For each SQL string literal extracted, every statement inside is checked. Multiple offenses in multiple literals are aggregated. """ offenses: list[str] = [] for body in extract_sql_literals(py_code): offenses.extend(find_forbidden_statements(body)) return offenses class SqlGuardError(ValueError): """Raised when py_code fails the SQL safety policy. -> HTTP 400 via the FastAPI ValueError handler in server.py. """