Claude
6b9f278294
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).
112 lines
2.8 KiB
Python
112 lines
2.8 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
"""
|
|
from fastapi import FastAPI
|
|
|
|
from spark_executor.tools.connections import (
|
|
delete_connection,
|
|
get_connection,
|
|
list_connections,
|
|
save_connection,
|
|
)
|
|
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,
|
|
confirm_submit_job,
|
|
get_pending_job,
|
|
list_pending_jobs,
|
|
prepare_submit_job,
|
|
)
|
|
|
|
app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
|
|
|
|
|
|
@app.get("/health")
|
|
def health_check():
|
|
return {"status": "ok"}
|
|
|
|
|
|
# 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(req: PrepareSubmitJobRequest):
|
|
return prepare_submit_job(**req.model_dump())
|
|
|
|
|
|
@app.post("/confirm_submit_job")
|
|
def _confirm_submit_job(req: PendingIdRequest):
|
|
return confirm_submit_job(pending_id=req.pending_id)
|
|
|
|
|
|
@app.post("/list_pending_jobs")
|
|
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
|
|
return list_pending_jobs()
|
|
|
|
|
|
@app.post("/get_pending_job")
|
|
def _get_pending_job(req: PendingIdRequest):
|
|
return get_pending_job(req.pending_id)
|
|
|
|
|
|
@app.post("/cancel_pending_job")
|
|
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(req: JobIdRequest):
|
|
return get_job_status(req.job_id)
|
|
|
|
|
|
@app.post("/get_job_logs")
|
|
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(req: JobIdRequest):
|
|
return kill_job(req.job_id)
|
|
|
|
|
|
# --- Connection management tools ---
|
|
|
|
@app.post("/save_connection")
|
|
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(_req: EmptyRequest = EmptyRequest()):
|
|
return list_connections()
|
|
|
|
|
|
@app.post("/get_connection")
|
|
def _get_connection(req: ConnectionNameRequest):
|
|
return get_connection(req.name)
|
|
|
|
|
|
@app.post("/delete_connection")
|
|
def _delete_connection(req: ConnectionNameRequest):
|
|
return delete_connection(req.name)
|