Files
mcp-server/spark_executor/tools/write_job.py
T
ClaudeandClaude Fable 5 523e6a9c76 refactor(mcp): rename generate_job_file to write_job_file to match what it does
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>
2026-06-29 19:13:29 +08:00

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}