feat: register 12 MCP tools on FastAPI app
This commit is contained in:
+105
-2
@@ -1,12 +1,115 @@
|
|||||||
# coding=utf-8
|
# coding=utf-8
|
||||||
"""
|
"""
|
||||||
@Time :2026/6/24
|
@Time :2026/6/24
|
||||||
@Author :tao.chen
|
@Author :tao.chen
|
||||||
"""
|
"""
|
||||||
from fastapi import FastAPI
|
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 = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
|
||||||
|
|
||||||
|
|
||||||
@app.get("/health")
|
@app.get("/health")
|
||||||
def health_check():
|
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