Claude
994c2d67ab
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/.
62 lines
1.4 KiB
Python
62 lines
1.4 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
|
|
Pydantic request models for the FastAPI route layer. The underlying tool
|
|
functions in tools/*.py still take keyword arguments; these models exist only
|
|
so fastapi-mcp can call the routes via tools/call (which sends args as a
|
|
JSON body) without 422-ing on dict-typed parameters like spark_conf.
|
|
"""
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class EmptyRequest(BaseModel):
|
|
"""Used for tools that take no arguments (list_connections, list_pending_jobs)."""
|
|
pass
|
|
|
|
|
|
class SaveConnectionRequest(BaseModel):
|
|
name: str
|
|
master: str
|
|
deploy_mode: str = "cluster"
|
|
yarn_rm_url: str | None = None
|
|
spark_conf: dict[str, str] | None = None
|
|
|
|
|
|
class PrepareSubmitJobRequest(BaseModel):
|
|
connection: str
|
|
script_path: str
|
|
queue: str = "default"
|
|
executor_memory: str = "4G"
|
|
executor_cores: int = 2
|
|
num_executors: int = 2
|
|
|
|
|
|
class PendingIdRequest(BaseModel):
|
|
pending_id: str
|
|
|
|
|
|
class JobIdRequest(BaseModel):
|
|
job_id: str
|
|
|
|
|
|
class GetJobLogsRequest(BaseModel):
|
|
job_id: str
|
|
tail_chars: int = 5000
|
|
|
|
|
|
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."
|
|
),
|
|
)
|