diff --git a/spark_executor/server.py b/spark_executor/server.py index 578c733..16ad11d 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -3,7 +3,8 @@ @Time :2026/6/24 @Author :tao.chen """ -from fastapi import FastAPI +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse from spark_executor.tools.connections import ( delete_connection, @@ -34,6 +35,23 @@ from spark_executor.tools.submit import ( app = FastAPI(title="Spark Executor", 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"} diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 4b85cf8..c77a975 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -110,3 +110,69 @@ def test_list_connections_works_with_empty_body(): r = c.post("/list_connections", json={}) assert r.status_code == 200 assert r.json() == [] + + +# --- Exception handlers: KeyError -> 404, ValueError -> 400 --- + +def test_unknown_job_id_returns_404(): + c = TestClient(app) + r = c.post("/get_job_status", json={"job_id": "missing"}) + assert r.status_code == 404 + assert "missing" in r.json()["detail"] + + +def test_unknown_pending_id_returns_404(): + c = TestClient(app) + r = c.post("/get_pending_job", json={"pending_id": "p_nope"}) + assert r.status_code == 404 + assert "p_nope" in r.json()["detail"] + + +def test_unknown_connection_name_returns_404(): + c = TestClient(app) + r = c.post("/get_connection", json={"name": "nope"}) + assert r.status_code == 404 + assert "nope" in r.json()["detail"] + + +def test_unknown_connection_in_prepare_returns_404(): + c = TestClient(app) + r = c.post( + "/prepare_submit_job", + json={"connection": "nope", "script_path": "/tmp/x.py"}, + ) + assert r.status_code == 404 + assert "nope" in r.json()["detail"] + + +def test_confirm_non_pending_returns_400(): + """A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400.""" + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + prep = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/x.py"}, + ).json() + pid = prep["pending_id"] + c.post("/cancel_pending_job", json={"pending_id": pid}) + r = c.post("/confirm_submit_job", json={"pending_id": pid}) + assert r.status_code == 400 + assert "CANCELLED" in r.json()["detail"] + + +def test_cancel_already_submitted_returns_400(): + """Cancelling a SUBMITTED pending should refuse with ValueError -> 400.""" + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + prep = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/x.py"}, + ).json() + pid = prep["pending_id"] + # Simulate a SUBMITTED state by mutating the pending directly + p = pending_store.store.get(pid) + p.status = "SUBMITTED" + pending_store.store.save(p) + r = c.post("/cancel_pending_job", json={"pending_id": pid}) + assert r.status_code == 400 + assert "SUBMITTED" in r.json()["detail"]