feat: add generate_job_file MCP tool (Stage 2 Task 22)

Exposes a single new MCP tool: generate_job_file(code) -> {script_path}.

Wiring follows the Stage 1 conventions (Pydantic body model for the
route, loguru DEBUG/INFO logging, summary+description for the startup
log, MCP arg-pattern via the body model so tools/call roundtrips long
code strings without 422-ing on FastAPI query length limits).

Flow:
  1. LLM calls generate_job_file(code=...)          -> {script_path}
  2. LLM calls prepare_submit_job(connection=...,
                                  script_path=...) -> {pending_id}
  3. User reviews the file + pending record
  4. LLM calls confirm_submit_job(pending_id=...)  -> spark-submit runs

The output directory is SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/,
gitignored, persists across container restarts via the existing
./data volume mount in docker-compose.yml).

3 new tests:
  - Unit: env-var override, default fallback in tmp cwd
  - Integration: end-to-end body call with a > FastAPI-query-limit code
    string (the canary test that would have caught the Stage 1
    query-params-422 bug)

Live MCP smoke verified: tools/list shows 14 tools (13 from Stage 1 +
the new one), tools/call generate_job_file returns the absolute path
under ./data/jobs/.
This commit is contained in:
Claude
2026-06-25 10:19:59 +08:00
parent 9340061b30
commit 994c2d67ab
5 changed files with 97 additions and 1 deletions
+20 -1
View File
@@ -1,4 +1,5 @@
# coding=utf-8
import os
from pathlib import Path
import pytest
@@ -32,7 +33,7 @@ def test_health_still_present():
assert r.json() == {"status": "ok"}
def test_twelve_tool_routes_registered():
def test_thirteen_tool_routes_registered():
paths = {r.path for r in app.routes}
for path in (
# pending-submission flow (5)
@@ -50,6 +51,8 @@ def test_twelve_tool_routes_registered():
"/list_connections",
"/get_connection",
"/delete_connection",
# LLM-driven PySpark generation (1) — Stage 2
"/generate_job_file",
):
assert path in paths, f"missing MCP tool route: {path}"
@@ -112,6 +115,22 @@ def test_list_connections_works_with_empty_body():
assert r.json() == []
# --- Stage 2: generate_job_file ---
def test_generate_job_file_accepts_code_string_in_body(tmp_path, monkeypatch):
"""End-to-end: a Pydantic body model lets tools/call pass a code string
that FastAPI would reject if it were a query parameter (length limits)."""
monkeypatch.chdir(tmp_path)
c = TestClient(app)
code = "print('from MCP integration test')\n" * 100 # > FastAPI query limit
r = c.post("/generate_job_file", json={"code": code})
assert r.status_code == 200, r.text
p = r.json()["script_path"]
assert os.path.isfile(p)
with open(p) as f:
assert f.read() == code
# --- Exception handlers: KeyError -> 404, ValueError -> 400 ---
def test_unknown_job_id_returns_404():
+27
View File
@@ -0,0 +1,27 @@
# coding=utf-8
import os
from pathlib import Path
from spark_executor.tools import generate
def test_generate_writes_code_and_returns_path(monkeypatch, tmp_path: Path):
monkeypatch.chdir(tmp_path)
out = generate.generate_job_file(
"from pyspark.sql import SparkSession\n"
"spark = SparkSession.builder.getOrCreate()\n"
)
assert "script_path" in out
p = out["script_path"]
assert os.path.isabs(p)
assert p.startswith(str(tmp_path / "data" / "jobs"))
assert p.endswith(".py")
with open(p) as f:
assert "SparkSession.builder.getOrCreate()" in f.read()
def test_generate_uses_env_var_when_set(monkeypatch, tmp_path: Path):
monkeypatch.setenv("SPARK_EXECUTOR_JOBS_DIR", str(tmp_path / "custom"))
out = generate.generate_job_file("x = 1\n")
assert out["script_path"].startswith(str(tmp_path / "custom"))
assert os.path.isfile(out["script_path"])