diff --git a/common/sql_guard.py b/common/sql_guard.py new file mode 100644 index 0000000..2e178d8 --- /dev/null +++ b/common/sql_guard.py @@ -0,0 +1,179 @@ +# 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. + """ diff --git a/pyproject.toml b/pyproject.toml index e2e2b92..d160c4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "loguru>=0.7.3", "pydantic>=2.13.4", "pydantic-core>=2.46.4", + "sqlparse>=0.5.0", "uvicorn>=0.49.0", ] diff --git a/spark_executor/tools/generate.py b/spark_executor/tools/generate.py index d671822..6ce431c 100644 --- a/spark_executor/tools/generate.py +++ b/spark_executor/tools/generate.py @@ -4,17 +4,40 @@ @Author :tao.chen """ from common.logging import logger +from common.sql_guard import validate_pyspark_code from spark_executor.core.job_writer import write_job_file +class SqlGuardViolation(ValueError): + """Raised when generate_job_file receives code that violates the SQL + safety policy. -> HTTP 400 via the FastAPI ValueError handler. + """ + pass + + def generate_job_file(code: str) -> dict[str, str]: """Write a PySpark code string to disk; return its absolute path. - Use the returned path as the `script_path` argument of prepare_submit_job. + Validates the code against the SQL safety policy (SELECT/INSERT only) + BEFORE writing, so the agent gets immediate feedback rather than + learning at prepare_submit_job time. Use the returned path as the + `script_path` argument of prepare_submit_job. + The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var (default: ./data/jobs/). """ logger.debug(f"generate_job_file enter code_bytes={len(code)}") + offenses = validate_pyspark_code(code) + if offenses: + logger.warning( + f"generate_job_file rejected: SQL policy violation(s): {offenses}" + ) + raise SqlGuardViolation( + f"PySpark code violates SQL safety policy. " + f"Only SELECT and INSERT statements are allowed. " + f"Found forbidden statement(s): {offenses}. " + f"Rewrite the code to use only SELECT/INSERT (or DataFrame DSL)." + ) path = write_job_file(code) logger.info(f"generate_job_file ok script_path={path}") return {"script_path": path} diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 793c0b5..e7ba0f0 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -9,6 +9,7 @@ import uuid from datetime import datetime from common.logging import logger +from common.sql_guard import validate_pyspark_code from spark_executor.core.connection_store import store as conn_store from spark_executor.core.job_store import JobStore from spark_executor.core.log_parser import parse_spark_submit_output @@ -42,6 +43,19 @@ def _check_script_path(script_path: str) -> None: raise _script_path_error(script_path) +def _sql_guard_error(offenses: list[str]) -> ValueError: + """Instructive 400 error when the script's SQL fails the safety policy.""" + return ValueError( + f"Script SQL violates the safety policy. " + f"Only SELECT and INSERT statements are allowed. " + f"Found forbidden statement(s): {offenses}. " + f"Edit the script and try again. (If the code was generated via " + f"generate_job_file, the generator should have caught this — check " + f"for f-string-based SQL injection where the static analysis can't " + f"see the runtime value.)" + ) + + def _new_pending_id() -> str: return "p_" + secrets.token_hex(6) @@ -64,6 +78,7 @@ def prepare_submit_job( # Order of checks matters for the error the agent sees: # 1. Unknown connection -> 404 (KeyError -> 404 in server.py) # 2. Missing script file -> 400 (ValueError -> 400) + # 3. SQL policy violation -> 400 (ValueError -> 400) # Connection is checked first because it is a more fundamental problem # (the agent is asking about a cluster that doesn't exist), and the # agent shouldn't have to fix the script path only to learn the @@ -78,6 +93,18 @@ def prepare_submit_job( # a 400 + clear remediation than to let spark-submit fail later with # an opaque FileNotFoundError -> 500. _check_script_path(script_path) + # SQL safety: re-validate the script even though generate_job_file + # already guards its own output. Catches files written by other means + # (host volume mounts, manual edits). + with open(script_path, encoding="utf-8") as f: + script_body = f.read() + offenses = validate_pyspark_code(script_body) + if offenses: + logger.warning( + f"prepare_submit_job rejected: SQL policy violation(s) in " + f"{script_path}: {offenses}" + ) + raise _sql_guard_error(offenses) pending_id = _new_pending_id() pending = PendingSubmission( diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index c94bf2d..9c0995e 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -109,6 +109,41 @@ def test_prepare_rejects_nonexistent_script_path_with_400(): assert "generate_job_file" in detail +def test_prepare_rejects_sql_policy_violation_with_400(tmp_path): + """SQL guard: even if the file exists, prepare_submit_job must reject + code containing DROP/DELETE/etc. (defense in depth — generate_job_file + also enforces this, but a host-mounted file might bypass that).""" + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + bad_script = tmp_path / "evil.py" + bad_script.write_text('spark.sql("DROP TABLE users")\n') + r = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": str(bad_script)}, + ) + assert r.status_code == 400 + detail = r.json()["detail"] + assert "DROP" in detail + assert "SELECT" in detail and "INSERT" in detail # the policy is explained + + +def test_generate_rejects_sql_policy_violation_with_400(tmp_path): + """Same policy at generate_job_file — agent gets immediate feedback + before the file is even written to disk.""" + from common import config + config.settings.jobs_dir = str(tmp_path / "jobs") + c = TestClient(app) + r = c.post( + "/generate_job_file", + json={"code": 'spark.sql("DELETE FROM events")\n'}, + ) + assert r.status_code == 400 + detail = r.json()["detail"] + assert "DELETE" in detail + # And the file should NOT have been written + assert list(tmp_path.glob("*.py")) == [] + + def test_list_and_get_pending_job_roundtrip_via_body(tmp_path): c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) diff --git a/tests/unit/test_sql_guard.py b/tests/unit/test_sql_guard.py new file mode 100644 index 0000000..64f42b0 --- /dev/null +++ b/tests/unit/test_sql_guard.py @@ -0,0 +1,200 @@ +# coding=utf-8 +import pytest + +from common.sql_guard import ( + SqlGuardError, + extract_sql_literals, + find_forbidden_statements, + validate_pyspark_code, +) + + +# --- extract_sql_literals --- + +def test_extract_picks_up_spark_sql_double_quoted(): + code = 'spark.sql("SELECT * FROM users")' + out = extract_sql_literals(code) + assert out == ["SELECT * FROM users"] + + +def test_extract_picks_up_spark_sql_single_quoted(): + code = "spark.sql('SELECT 1')" + out = extract_sql_literals(code) + assert out == ["SELECT 1"] + + +def test_extract_ignores_english_sentence_starting_with_select(): + """False positive tolerated for safety: 'select the best option' would + be picked up. That's fine — it's a non-SQL literal that gets + classified, but won't contain any forbidden keyword.""" + code = 'log("select the best option")' + out = extract_sql_literals(code) + assert out == ["select the best option"] + # And the validator returns no offenses for it (no SQL verbs in body) + assert find_forbidden_statements(out[0]) == [] + + +def test_extract_ignores_non_sql_strings(): + code = 'msg = "hello world"\nlog(msg)' + assert extract_sql_literals(code) == [] + + +def test_extract_handles_fstring_prefix(): + code = 'spark.sql(f"SELECT * FROM {table}")' + out = extract_sql_literals(code) + assert out == ["SELECT * FROM {table}"] + + +def test_extract_collects_multiple_literals(): + code = ''' +df1 = spark.sql("SELECT * FROM t1") +df2 = spark.sql("SELECT * FROM t2") +''' + out = extract_sql_literals(code) + assert out == ["SELECT * FROM t1", "SELECT * FROM t2"] + + +# --- find_forbidden_statements --- + +def test_select_only_is_allowed(): + assert find_forbidden_statements("SELECT * FROM users") == [] + + +def test_select_with_where_is_allowed(): + assert find_forbidden_statements( + "SELECT id, name FROM users WHERE age > 18 ORDER BY id" + ) == [] + + +def test_insert_is_allowed(): + assert find_forbidden_statements( + "INSERT INTO events SELECT * FROM raw_events" + ) == [] + + +def test_drop_table_is_forbidden(): + assert find_forbidden_statements("DROP TABLE users") == ["DROP"] + + +def test_drop_database_is_forbidden(): + assert find_forbidden_statements("DROP DATABASE prod") == ["DROP"] + + +def test_delete_is_forbidden(): + assert find_forbidden_statements("DELETE FROM users WHERE id = 1") == ["DELETE"] + + +def test_update_is_forbidden(): + assert find_forbidden_statements( + "UPDATE users SET name = 'x' WHERE id = 1" + ) == ["UPDATE"] + + +def test_truncate_is_forbidden(): + assert find_forbidden_statements("TRUNCATE TABLE events") == ["TRUNCATE"] + + +def test_alter_is_forbidden(): + assert find_forbidden_statements("ALTER TABLE users ADD COLUMN x INT") == ["ALTER"] + + +def test_create_table_is_forbidden(): + assert find_forbidden_statements( + "CREATE TABLE foo (id INT, name STRING)" + ) == ["CREATE"] + + +def test_with_cte_select_is_allowed(): + sql = "WITH active AS (SELECT * FROM users WHERE active) SELECT * FROM active" + assert find_forbidden_statements(sql) == [] + + +def test_with_cte_drop_is_forbidden(): + """With CTE body containing a DROP. sqlparse may split this into two + statements or keep it as one — either way, the DROP must be flagged.""" + sql = "WITH x AS (DROP TABLE y) SELECT * FROM x" + out = find_forbidden_statements(sql) + # The DROP must be reported (label may be 'DROP' or 'WITH ... DROP' + # depending on how sqlparse tokenizes; both prove it's caught). + assert any("DROP" in o for o in out) + assert any(o in {"DROP", "WITH ... DROP"} for o in out) + + +def test_multi_statement_selects_all_allowed(): + sql = "SELECT 1; SELECT 2; INSERT INTO t VALUES (1)" + assert find_forbidden_statements(sql) == [] + + +def test_multi_statement_one_bad_is_caught(): + sql = "SELECT 1; DROP TABLE users; SELECT 2" + assert find_forbidden_statements(sql) == ["DROP"] + + +def test_comments_and_whitespace_dont_confuse_parser(): + sql = """ + -- this is a comment + /* multi-line + comment */ + SELECT * FROM users + """ + assert find_forbidden_statements(sql) == [] + + +# --- validate_pyspark_code (end-to-end through the Python string extractor) --- + +def test_validate_clean_pyspark_code(): + code = ''' +from pyspark.sql import SparkSession +spark = SparkSession.builder.getOrCreate() +df = spark.sql("SELECT * FROM users") +df.show() +''' + assert validate_pyspark_code(code) == [] + + +def test_validate_rejects_drop_in_pyspark_code(): + code = 'spark.sql("DROP TABLE users")' + out = validate_pyspark_code(code) + assert "DROP" in out + + +def test_validate_rejects_delete_with_where(): + code = 'spark.sql("DELETE FROM events WHERE id = 1")' + out = validate_pyspark_code(code) + assert "DELETE" in out + + +def test_validate_rejects_update(): + code = 'spark.sql("UPDATE users SET x = 1")' + out = validate_pyspark_code(code) + assert "UPDATE" in out + + +def test_validate_collects_multiple_offenses(): + code = ''' +spark.sql("DROP TABLE a") +spark.sql("DELETE FROM b") +''' + out = validate_pyspark_code(code) + assert "DROP" in out + assert "DELETE" in out + + +def test_validate_passes_when_no_sql_present(): + code = ''' +from pyspark.sql import SparkSession +spark = SparkSession.builder.getOrCreate() +df = spark.range(0, 100) +df.show() +''' + assert validate_pyspark_code(code) == [] + + +def test_validate_passes_with_dataframe_dsl(): + """DataFrame operations (filter, select, groupBy, agg) are not raw SQL.""" + code = ''' +df = spark.read.parquet("/data/foo") +filtered = df.filter(df.age > 18).select("id", "name").groupBy("name").count() +filtered.show() +''' + assert validate_pyspark_code(code) == [] diff --git a/uv.lock b/uv.lock index 9ad6562..7a857a1 100644 --- a/uv.lock +++ b/uv.lock @@ -849,6 +849,7 @@ dependencies = [ { name = "loguru" }, { name = "pydantic" }, { name = "pydantic-core" }, + { name = "sqlparse" }, { name = "uvicorn" }, ] @@ -867,6 +868,7 @@ requires-dist = [ { name = "loguru", specifier = ">=0.7.3" }, { name = "pydantic", specifier = ">=2.13.4" }, { name = "pydantic-core", specifier = ">=2.46.4" }, + { name = "sqlparse", specifier = ">=0.5.0" }, { name = "uvicorn", specifier = ">=0.49.0" }, ] @@ -876,6 +878,15 @@ dev = [ { name = "pytest-mock", specifier = ">=3.14" }, ] +[[package]] +name = "sqlparse" +version = "0.5.5" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" }, +] + [[package]] name = "sse-starlette" version = "3.4.5"