feat: register 12 MCP tools on FastAPI app
This commit is contained in:
+105
-2
@@ -1,12 +1,115 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/6/24
|
||||
@Author :tao.chen
|
||||
@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.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"}
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.
|
||||
|
||||
# --- 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,
|
||||
)
|
||||
|
||||
|
||||
@app.post("/confirm_submit_job")
|
||||
def _confirm_submit_job(pending_id: str):
|
||||
return confirm_submit_job(pending_id=pending_id)
|
||||
|
||||
|
||||
@app.post("/list_pending_jobs")
|
||||
def _list_pending_jobs():
|
||||
return list_pending_jobs()
|
||||
|
||||
|
||||
@app.post("/get_pending_job")
|
||||
def _get_pending_job(pending_id: str):
|
||||
return get_pending_job(pending_id)
|
||||
|
||||
|
||||
@app.post("/cancel_pending_job")
|
||||
def _cancel_pending_job(pending_id: str):
|
||||
return cancel_pending_job(pending_id)
|
||||
|
||||
|
||||
# --- Spark job tools ---
|
||||
|
||||
@app.post("/get_job_status")
|
||||
def _get_job_status(job_id: str):
|
||||
return get_job_status(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)
|
||||
|
||||
|
||||
@app.post("/kill_job")
|
||||
def _kill_job(job_id: str):
|
||||
return kill_job(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,
|
||||
)
|
||||
|
||||
|
||||
@app.post("/list_connections")
|
||||
def _list_connections():
|
||||
return list_connections()
|
||||
|
||||
|
||||
@app.post("/get_connection")
|
||||
def _get_connection(name: str):
|
||||
return get_connection(name)
|
||||
|
||||
|
||||
@app.post("/delete_connection")
|
||||
def _delete_connection(name: str):
|
||||
return delete_connection(name)
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/6/24
|
||||
@Author :tao.chen
|
||||
"""
|
||||
@@ -0,0 +1,33 @@
|
||||
# coding=utf-8
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from spark_executor.server import app
|
||||
|
||||
|
||||
def test_health_still_present():
|
||||
c = TestClient(app)
|
||||
r = c.get("/health")
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"status": "ok"}
|
||||
|
||||
|
||||
def test_twelve_tool_routes_registered():
|
||||
paths = {r.path for r in app.routes}
|
||||
for path in (
|
||||
# pending-submission flow (5)
|
||||
"/prepare_submit_job",
|
||||
"/confirm_submit_job",
|
||||
"/list_pending_jobs",
|
||||
"/get_pending_job",
|
||||
"/cancel_pending_job",
|
||||
# job lifecycle (3)
|
||||
"/get_job_status",
|
||||
"/get_job_logs",
|
||||
"/kill_job",
|
||||
# connection management (4)
|
||||
"/save_connection",
|
||||
"/list_connections",
|
||||
"/get_connection",
|
||||
"/delete_connection",
|
||||
):
|
||||
assert path in paths, f"missing MCP tool route: {path}"
|
||||
Reference in New Issue
Block a user