130 lines
3.6 KiB
Python
130 lines
3.6 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
"""
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
|
|
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 MCP", version="0.0.1", description="Spark Executor MCP Server")
|
|
|
|
|
|
# --- Exception handlers: translate tool-layer errors into proper HTTP statuses ---
|
|
#
|
|
# Tool functions raise KeyError for "unknown id" (job_id, pending_id, connection
|
|
# name) and ValueError for invalid state transitions (e.g. confirming a
|
|
# CANCELLED pending). Without these handlers FastAPI would map them to a bare
|
|
# 500 "Internal Server Error" which is useless to MCP clients.
|
|
|
|
@app.exception_handler(KeyError)
|
|
async def _keyerror_handler(_request: Request, exc: KeyError) -> JSONResponse:
|
|
return JSONResponse(status_code=404, content={"detail": str(exc)})
|
|
|
|
|
|
@app.exception_handler(ValueError)
|
|
async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONResponse:
|
|
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
|
|
|
|
|
@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)
|