feat: SQL safety policy (SELECT/INSERT only) at submit time

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/.

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<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.
+    """

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",
 ]
 

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}

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(

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"})

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) == []

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"