The MCP tool was named `generate_job_file` from Stage 2 but it does
NOT generate PySpark code — the calling LLM writes the code in its own
context, and this tool only persists it to a file under
SPARK_EXECUTOR_JOBS_DIR so `spark-submit` can see it. The misleading
`generate_` prefix sent agents (and humans) looking for a code
generator that doesn't exist.
This commit folds three related polish changes into one (split later
with rebase -i if you want them as separate history):
1. The rename itself:
- `tools/generate.py` → `tools/write_job.py`
- `generate_job_file` → `write_job_file`
- `GenerateJobFileRequest` → `WriteJobFileRequest`
- `/generate_job_file` route → `/write_job_file`
- `operation_id="generate_job_file"` → `operation_id="write_job_file"`
The internal helper `core.job_writer.write_job_file` (which just
writes bytes to disk with no SQL guard) is imported with an
`_write_to_disk` alias to avoid the name collision with the
MCP-exposed function in the same module.
The description for the tool now explicitly states 'this tool
does NOT generate PySpark code. The calling LLM is expected to
have already written the code; this tool only persists it.'
2. Skill for LLM agents operating the service
(`docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md`,
449 lines). Covers the 16 tools, the two-step prepare/confirm
flow, the dual-ID contract (job_id vs application_id), the
PendingSubmission state machine, the Connection profile, the
job-file workflow, the error reference, common pitfalls, and a
full end-to-end word-count example.
3. Default `executor_memory` lowered 4G → 2G
(`_DEFAULTS_TO_CONFIRM` in `server.py`). Mirrors the matching
change in `test_mcp_routes.py` and the 5 unit tests that
reference the default. Aligns with the lighter workloads the
service is sized for in its current container profile.
Also tracked in git for the first time:
- `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`
(the original Stage 1/2/3 design plan, updated to use the new
tool name throughout).
Test rename:
- `tests/unit/test_generate_tool.py` → `test_write_job_tool.py`
- the new test file picks up an extra assertion that the SQL guard
rejects a `DROP TABLE` statement at write time.
243 tests pass (was 242; +1 new SQL-guard assertion). Zero regressions.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
|
|
The MCP-exposed `write_job_file` tool. Despite the historical name
|
|
`generate_job_file`, this tool does NOT generate PySpark code — the
|
|
calling LLM writes the code, and this tool persists it to disk so
|
|
`spark-submit` can see it. The rename (2026-06-29) tightens the naming
|
|
to match what the tool actually does, eliminating the "I thought this
|
|
would generate code for me" confusion.
|
|
|
|
It also runs the SQL safety policy (SELECT/INSERT only) on the code
|
|
BEFORE writing, so the agent gets immediate feedback if it slipped a
|
|
DROP / DELETE / UPDATE past its own code-generation step.
|
|
"""
|
|
from common.logging import logger
|
|
from common.sql_guard import validate_pyspark_code
|
|
# Internal helper already named write_job_file (it just writes bytes to
|
|
# disk; no SQL guard). Import with an alias to avoid the name collision
|
|
# with the MCP-exposed function below.
|
|
from spark_executor.core.job_writer import write_job_file as _write_to_disk
|
|
|
|
|
|
class SqlGuardViolation(ValueError):
|
|
"""Raised when write_job_file receives code that violates the SQL
|
|
safety policy. -> HTTP 400 via the FastAPI ValueError handler."""
|
|
pass
|
|
|
|
|
|
def write_job_file(code: str) -> dict[str, str]:
|
|
"""Persist a PySpark code string to disk and return its absolute path.
|
|
|
|
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/). Each call writes to a fresh timestamped file
|
|
under that directory, so repeated calls never overwrite one another.
|
|
"""
|
|
logger.debug(f"write_job_file enter code_bytes={len(code)}")
|
|
offenses = validate_pyspark_code(code)
|
|
if offenses:
|
|
logger.warning(
|
|
f"write_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."
|
|
)
|
|
path = _write_to_disk(code)
|
|
logger.info(f"write_job_file ok script_path={path}")
|
|
return {"script_path": path}
|