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:
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user