fix: switch FastAPI routes to Pydantic body models for MCP tools/call

The plan's route signatures used query parameters, which worked for direct
HTTP callers and unit tests, but fastapi-mcp's HTTP transport passes
tools/call arguments as a JSON body. dict-typed parameters like
spark_conf arrived as a string and the route returned 422.

Refactor each route to take a single Pydantic body model (saved in
spark_executor/tools/requests.py). Underlying tool functions unchanged.

Integration tests in tests/integration/test_mcp_routes.py now exercise
the full body-based roundtrip (save_connection with spark_conf, prepare
→ list → get, list_connections with empty body).
This commit is contained in:
Claude
2026-06-24 14:56:31 +08:00
parent cc7eafaa16
commit 6b9f278294
3 changed files with 164 additions and 39 deletions
+35 -39
View File
@@ -13,6 +13,15 @@ from spark_executor.tools.connections import (
)
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.requests import (
ConnectionNameRequest,
EmptyRequest,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
PrepareSubmitJobRequest,
SaveConnectionRequest,
)
from spark_executor.tools.status import get_job_status
from spark_executor.tools.submit import (
cancel_pending_job,
@@ -31,85 +40,72 @@ def health_check():
# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.
# Each route takes a single Pydantic body model so tools/call (which sends args
# as JSON body) works for every tool, including those with dict-typed params
# like spark_conf.
# --- Pending submission flow (two-step submit) ---
@app.post("/prepare_submit_job")
def _prepare_submit_job(connection: str, script_path: str, queue: str = "default",
executor_memory: str = "4G", executor_cores: int = 2,
num_executors: int = 2):
return prepare_submit_job(
connection=connection,
script_path=script_path,
queue=queue,
executor_memory=executor_memory,
executor_cores=executor_cores,
num_executors=num_executors,
)
def _prepare_submit_job(req: PrepareSubmitJobRequest):
return prepare_submit_job(**req.model_dump())
@app.post("/confirm_submit_job")
def _confirm_submit_job(pending_id: str):
return confirm_submit_job(pending_id=pending_id)
def _confirm_submit_job(req: PendingIdRequest):
return confirm_submit_job(pending_id=req.pending_id)
@app.post("/list_pending_jobs")
def _list_pending_jobs():
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
return list_pending_jobs()
@app.post("/get_pending_job")
def _get_pending_job(pending_id: str):
return get_pending_job(pending_id)
def _get_pending_job(req: PendingIdRequest):
return get_pending_job(req.pending_id)
@app.post("/cancel_pending_job")
def _cancel_pending_job(pending_id: str):
return cancel_pending_job(pending_id)
def _cancel_pending_job(req: PendingIdRequest):
return cancel_pending_job(req.pending_id)
# --- Spark job tools ---
@app.post("/get_job_status")
def _get_job_status(job_id: str):
return get_job_status(job_id)
def _get_job_status(req: JobIdRequest):
return get_job_status(req.job_id)
@app.post("/get_job_logs")
def _get_job_logs(job_id: str, tail_chars: int = 5000):
return get_job_logs(job_id, tail_chars=tail_chars)
def _get_job_logs(req: GetJobLogsRequest):
return get_job_logs(req.job_id, tail_chars=req.tail_chars)
@app.post("/kill_job")
def _kill_job(job_id: str):
return kill_job(job_id)
def _kill_job(req: JobIdRequest):
return kill_job(req.job_id)
# --- Connection management tools ---
@app.post("/save_connection")
def _save_connection(name: str, master: str, deploy_mode: str = "cluster",
yarn_rm_url: str | None = None,
spark_conf: dict[str, str] | None = None):
return save_connection(
name=name,
master=master,
deploy_mode=deploy_mode,
yarn_rm_url=yarn_rm_url,
spark_conf=spark_conf,
)
def _save_connection(req: SaveConnectionRequest):
# exclude_none so we don't overwrite the function's default with explicit None
return save_connection(**req.model_dump(exclude_none=True))
@app.post("/list_connections")
def _list_connections():
def _list_connections(_req: EmptyRequest = EmptyRequest()):
return list_connections()
@app.post("/get_connection")
def _get_connection(name: str):
return get_connection(name)
def _get_connection(req: ConnectionNameRequest):
return get_connection(req.name)
@app.post("/delete_connection")
def _delete_connection(name: str):
return delete_connection(name)
def _delete_connection(req: ConnectionNameRequest):
return delete_connection(req.name)