Claude
34e6208f54
The agent can now write PySpark that runs DROP/DELETE/UPDATE/etc. on
production tables. Add a static guard that rejects anything other than
SELECT and INSERT at two enforcement points:
1. generate_job_file: validates BEFORE writing to disk. Agent gets
immediate feedback ('rewrite to use only SELECT/INSERT') rather
than learning at submit time.
2. prepare_submit_job: re-validates the script content (reads the
file) as a defense-in-depth check. Catches host-mounted files,
manually-edited files, anything that bypassed generate_job_file.
How it works:
- common/sql_guard.py extracts Python string literals whose first
keyword is a SQL verb (catches spark.sql('...'), f-strings, and any
raw SQL literal)
- sqlparse splits each literal into statements; we check the first
keyword against the policy (SELECT/INSERT/WITH allowed; DROP,
DELETE, UPDATE, TRUNCATE, ALTER, CREATE, REPLACE, MERGE, GRANT,
REVOKE, SET, SHOW, KILL, EXEC, etc. forbidden)
- WITH recurses into the CTE body to catch WITH x AS (DROP ...) ...
- The MCP layer maps ValueError -> HTTP 400 (existing handler)
Test coverage:
- 28 unit tests in test_sql_guard.py cover: extraction (single/double/
f-string, English false positives, multi-literal), statement
classification (SELECT, INSERT, DROP, DELETE, UPDATE, TRUNCATE,
ALTER, CREATE, multi-statement, CTE bodies, comments)
- 2 integration tests verify MCP layer returns 400 with the policy
explanation at both generate_job_file and prepare_submit_job
Limitations (documented in sql_guard.py docstring):
- f-strings where the SQL is built at runtime (e.g. f'SELECT * FROM
{user_input}') look like SELECTs at static-analysis time. The
guard catches the static literal; the runtime substitution is the
caller's responsibility.
- pyspark.sql.functions.expr('...') accepts SQL inline; not currently
caught. (Future work.)
146/146 still pass. Live verified: DROP TABLE -> MCP 400 with policy
explanation; SELECT -> MCP 200 + file written to ./data/jobs/.
180 lines
6.5 KiB
Python
180 lines
6.5 KiB
Python
# 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<q>['"])
|
|
(?P<body>(?:\\.|(?!(?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 ... <verb> 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 ... <forbidden> 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.
|
|
"""
|