mcp-server/spark_executor/server.py

# 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.generate import generate_job_file
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,
    GenerateJobFileRequest,
    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",
    summary="Prepare a Spark job submission (no spark-submit yet)",
    description=(
        "Snapshot the named Connection's master / deploy_mode / spark_conf / "
        "yarn_rm_url into a PendingSubmission record and persist it. "
        "Does NOT invoke spark-submit. Returns pending_id for use with "
        "confirm_submit_job (the user-second-confirmation step)."
    ),
)
def _prepare_submit_job(req: PrepareSubmitJobRequest):
    return prepare_submit_job(**req.model_dump())


@app.post(
    "/confirm_submit_job",
    summary="Confirm and submit a previously-prepared job",
    description=(
        "Actually invoke spark-submit for the PendingSubmission identified "
        "by pending_id. Requires status=PENDING. On success, transitions the "
        "pending entry to SUBMITTED and creates a Job record. On failure, "
        "marks the entry FAILED and re-raises."
    ),
)
def _confirm_submit_job(req: PendingIdRequest):
    return confirm_submit_job(pending_id=req.pending_id)


@app.post(
    "/list_pending_jobs",
    summary="List all pending submissions",
    description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).",
)
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
    return list_pending_jobs()


@app.post(
    "/get_pending_job",
    summary="Get a single pending submission",
    description="Return the PendingSubmission identified by pending_id, including its current status and outcome fields.",
)
def _get_pending_job(req: PendingIdRequest):
    return get_pending_job(req.pending_id)


@app.post(
    "/cancel_pending_job",
    summary="Cancel a pending submission",
    description=(
        "Flip a PENDING (or already-CANCELLED) PendingSubmission to CANCELLED. "
        "Refuses to cancel entries that are SUBMITTED or FAILED — those are "
        "terminal and must be killed via kill_job instead."
    ),
)
def _cancel_pending_job(req: PendingIdRequest):
    return cancel_pending_job(req.pending_id)


# --- Spark job tools ---

@app.post(
    "/get_job_status",
    summary="Query YARN for a job's current status",
    description=(
        "Return the YARN application state (RUNNING / SUCCEEDED / FAILED / "
        "KILLED / ACCEPTED / NEW / NEW_SAVING / SUBMITTED / etc.) plus the "
        "raw YARN REST response body."
    ),
)
def _get_job_status(req: JobIdRequest):
    return get_job_status(req.job_id)


@app.post(
    "/get_job_logs",
    summary="Fetch aggregated container logs for a job",
    description=(
        "Pull aggregated logs from the YARN ResourceManager. Returns the last "
        "tail_chars characters (default 5000). Requires yarn.log-aggregation-enable "
        "to be true on the target cluster."
    ),
)
def _get_job_logs(req: GetJobLogsRequest):
    return get_job_logs(req.job_id, tail_chars=req.tail_chars)


@app.post(
    "/kill_job",
    summary="Kill a running job",
    description="PUT state=KILLED to YARN REST API for the job's application_id.",
)
def _kill_job(req: JobIdRequest):
    return kill_job(req.job_id)


# --- Connection management tools ---

@app.post(
    "/save_connection",
    summary="Save or update a named Spark connection",
    description=(
        "Upsert a Connection record (master URL, deploy mode, optional YARN RM URL, "
        "spark_conf K/V) keyed by name. Used by prepare_submit_job via the "
        "connection parameter."
    ),
)
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",
    summary="List all saved Spark connections",
    description="Return every Connection in the registry (model_dump form).",
)
def _list_connections(_req: EmptyRequest = EmptyRequest()):
    return list_connections()


@app.post(
    "/get_connection",
    summary="Get a single connection by name",
    description="Return the Connection record, or 404 if not found.",
)
def _get_connection(req: ConnectionNameRequest):
    return get_connection(req.name)


@app.post(
    "/delete_connection",
    summary="Delete a saved connection",
    description="Remove a Connection by name. 404 if not found.",
)
def _delete_connection(req: ConnectionNameRequest):
    return delete_connection(req.name)


# --- LLM-driven PySpark generation (Stage 2) ---

@app.post(
    "/generate_job_file",
    summary="Write LLM-generated PySpark code to disk",
    description=(
        "Takes a PySpark code string and writes it to a timestamped file under "
        "SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
        "path for use as the script_path argument of prepare_submit_job — the "
        "two-step pattern means the LLM can produce code, the user can review "
        "the resulting file, and only then is the job submitted."
    ),
)
def _generate_job_file(req: GenerateJobFileRequest):
    return generate_job_file(req.code)