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:
@@ -0,0 +1,20 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/6/24
|
||||
@Author :tao.chen
|
||||
"""
|
||||
from common.logging import logger
|
||||
from spark_executor.core.job_writer import write_job_file
|
||||
|
||||
|
||||
def generate_job_file(code: str) -> dict[str, str]:
|
||||
"""Write a PySpark code string to disk; return its absolute path.
|
||||
|
||||
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/).
|
||||
"""
|
||||
logger.debug(f"generate_job_file enter code_bytes={len(code)}")
|
||||
path = write_job_file(code)
|
||||
logger.info(f"generate_job_file ok script_path={path}")
|
||||
return {"script_path": path}
|
||||
@@ -48,3 +48,14 @@ class GetJobLogsRequest(BaseModel):
|
||||
|
||||
class ConnectionNameRequest(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class GenerateJobFileRequest(BaseModel):
|
||||
code: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Full PySpark source code to write to disk. Will be passed verbatim "
|
||||
"to spark-submit after the agent calls prepare_submit_job on the "
|
||||
"returned path."
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user