Claude
35bff11edf
Closes the review-and-edit loop for LLM-generated PySpark code:
generate_job_file(code=...) -> {script_path}
read_job_file(script_path=...) -> {content, path, size}
update_job_file(path, content) -> {path, bytes_written}
prepare_submit_job(path) -> {pending_id, ...}
Or, in a single edit cycle:
1. generate (LLM writes initial draft)
2. read (LLM or human inspects)
3. update (overwrite with edited version)
4. prepare (submit for two-step confirmation)
Safety:
- read_job_file has no path restriction (read-only; useful for
inspecting any file the agent can see: scripts, logs/, hadoop-conf/)
- update_job_file is sandboxed to settings.jobs_dir (the same dir
generate_job_file writes to). Rejects paths outside that tree,
including ../-traversal attempts. This protects host-mounted
configs (/etc/passwd, hadoop-conf/*) from being overwritten by
the agent.
- 1 MB cap on both reads and writes so MCP responses stay bounded.
Pydantic body models (ReadJobFileRequest, UpdateJobFileRequest) follow
the Stage 1 pattern so tools/call roundtrips long code strings without
the FastAPI query-length 422.
Tests (13 new):
- 9 unit tests: read success/missing/empty/dir, update success/outside/
relative-escape/missing/oversize/1mb+1, full edit cycle round-trip
- 4 integration tests: read via MCP, missing file 400, write+readback
via MCP, outside-jobs_dir rejection via MCP
163/146 still pass. Live verified end-to-end: generate -> read v1
-> update -> read v2; update /etc/passwd correctly 400'd with
'script_path must be under ... data/jobs/'.
102 lines
2.8 KiB
Python
102 lines
2.8 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 = Field(
|
|
...,
|
|
description=(
|
|
"Absolute path to the PySpark script inside the container's "
|
|
"filesystem. Must point at an existing regular file. For "
|
|
"LLM-generated code, call generate_job_file(code=...) first "
|
|
"and pass the returned script_path here. For pre-existing "
|
|
"files on the host, mount them via a docker volume and pass "
|
|
"the in-container path. Returns 400 with a remediation hint "
|
|
"if the path is missing or not a file."
|
|
),
|
|
)
|
|
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."
|
|
),
|
|
)
|
|
|
|
|
|
class ReadJobFileRequest(BaseModel):
|
|
script_path: str = Field(
|
|
...,
|
|
description=(
|
|
"Absolute path to a PySpark script inside the container's "
|
|
"filesystem. Must point at an existing regular file."
|
|
),
|
|
)
|
|
|
|
|
|
class UpdateJobFileRequest(BaseModel):
|
|
script_path: str = Field(
|
|
...,
|
|
description=(
|
|
"Absolute path to an existing PySpark script inside the "
|
|
"container's filesystem. Must be under SPARK_EXECUTOR_JOBS_DIR "
|
|
"(the same dir generate_job_file writes to) — protects against "
|
|
"overwriting host-mounted configs or other critical files."
|
|
),
|
|
)
|
|
content: str = Field(
|
|
...,
|
|
description=(
|
|
"New file content (replaces the file in full; no merge/diff). "
|
|
"Maximum 1 MB to keep the MCP response bounded."
|
|
),
|
|
)
|