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/.
This commit is contained in:
Claude
2026-06-25 12:52:16 +08:00
parent 15e9a75bc7
commit 34e6208f54
7 changed files with 477 additions and 1 deletions
+35
View File
@@ -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"})