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,
    update_connection,
)
from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_job_file
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.external_jobs import (
    get_external_job_logs,
    get_external_job_status,
    get_external_job_result,
)
from spark_executor.tools.fetch_url import fetch_url
from spark_executor.tools.requests import (
    ConnectionNameRequest,
    EmptyRequest,
    WriteJobFileRequest,
    ExternalJobLogsRequest,
    ExternalJobStatusRequest,
    ExternalJobResultRequest,
    FetchUrlRequest,
    GetJobLogsRequest,
    JobIdRequest,
    PendingIdRequest,
    PrepareSubmitJobRequest,
    ReadJobFileRequest,
    SaveConnectionRequest,
    UpdateConnectionRequest,
    UpdateJobFileRequest,
    UpdatePendingJobRequest,
)
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,
    update_pending_job,
)

from spark_executor.tools.result import get_job_result

app = FastAPI(title="Spark Executor MCP", version="0.0.1", description="Spark Executor MCP Server")


_DEFAULTS_TO_CONFIRM = {
    "queue": "default",
    "executor_memory": "2G",
    "executor_cores": 2,
    "num_executors": 2,
}


# --- 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",
    operation_id="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).\n\n"
        "REQUIRED PATTERN for LLM-generated code: call write_job_file(code=...) "
        "first, then pass the returned script_path here. Direct submission with a "
        "synthetic path (one that only exists in the agent's context) will be "
        "rejected with HTTP 400 — the script must exist inside the container's "
        "filesystem. For pre-existing files, mount the host directory into the "
        "container and pass the in-container path."
    ),
)
def _prepare_submit_job(req: PrepareSubmitJobRequest):
    omitted = [f for f in _DEFAULTS_TO_CONFIRM if f not in req.model_fields_set]
    if omitted:
        details = ", ".join(f"{f}={_DEFAULTS_TO_CONFIRM[f]!r}" for f in omitted)
        raise ValueError(
            f"Please confirm default values: {details}. "
            f"Resubmit with these fields explicitly set."
        )
    return prepare_submit_job(**req.model_dump())


@app.post(
    "/confirm_submit_job",
    operation_id="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. A FAILED pending can be "
        "re-confirmed — it resets to PENDING for a single fresh attempt — so "
        "transient failures (e.g. YARN RM was down) are recoverable."
    ),
)
def _confirm_submit_job(req: PendingIdRequest):
    return confirm_submit_job(pending_id=req.pending_id)


@app.post(
    "/list_pending_jobs",
    operation_id="list_pending_jobs",
    summary="List all pending submissions",
    description=(
        "Return every PendingSubmission in any status (PENDING, SUBMITTED, "
        "CANCELLED, FAILED). Call this before prepare_submit_job to check if a "
        "submission with the same parameters is already in flight, or after a "
        "batch of confirm_submit_job calls to inspect the lifecycle of recent "
        "submissions."
    ),
)
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
    return list_pending_jobs()


@app.post(
    "/get_pending_job",
    operation_id="get_pending_job",
    summary="Get a single pending submission",
    description=(
        "Return the PendingSubmission identified by pending_id, including its "
        "current status and outcome fields. Use this to inspect a pending "
        "submission between prepare_submit_job and confirm_submit_job (e.g. "
        "to confirm the snapshotted connection), or to read the error field "
        "of a FAILED submission before re-confirming."
    ),
)
def _get_pending_job(req: PendingIdRequest):
    return get_pending_job(req.pending_id)


@app.post(
    "/update_pending_job",
    operation_id="update_pending_job",
    summary="Update an unsubmitted pending submission",
    description=(
        "Modify parameters of a PENDING submission before confirm_submit_job. "
        "Only the provided fields are changed. If script_path is changed, the "
        "new file must exist and pass the SQL guard."
    ),
)
def _update_pending_job(req: UpdatePendingJobRequest):
    return update_pending_job(**req.model_dump(exclude_none=True))


@app.post(
    "/cancel_pending_job",
    operation_id="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",
    operation_id="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.\n\n"
        "**job_id accepts BOTH identifiers** returned by "
        "confirm_submit_job: the local job_id (12-char hex, e.g. "
        "'a1b2c3d4e5f6') and the YARN application_id (e.g. "
        "'application_17400000001_0001'). The lookup is by job_id first, "
        "then by application_id. **If you pass a YARN application_id and "
        "the app is NOT in the local JobStore, this tool returns HTTP 400** "
        "(not 404) with a hint message naming the right external tool. "
        "**For YARN applications NOT submitted through this service** "
        "(no local JobStore record), use "
        "`get_external_job_status(application_id, connection_name)` "
        "directly — it bypasses the local registry and queries YARN."
    ),
)
def _get_job_status(req: JobIdRequest):
    return get_job_status(req.job_id)


@app.post(
    "/get_job_result",
    operation_id="get_job_result",
    summary="Query YARN for a job's terminal result view",
    description=(
        "Return a terminal-oriented view of a Spark job: final_status, "
        "diagnostics, tracking_url, started_time, and finished_time. "
        "This is distinct from get_job_status, which is for polling the "
        "running YARN state and returns the raw YARN response.\n\n"
        "**job_id accepts BOTH identifiers** returned by "
        "confirm_submit_job: the local job_id (12-char hex, e.g. "
        "'a1b2c3d4e5f6') and the YARN application_id (e.g. "
        "'application_17400000001_0001'). The lookup is by job_id first, "
        "then by application_id. **If you pass a YARN application_id and "
        "the app is NOT in the local JobStore, this tool returns HTTP 400** "
        "(not 404) with a hint message naming the right external tool. "
        "**For YARN applications NOT submitted through this service** "
        "(no local JobStore record), use "
        "`get_external_job_result(application_id, connection_name)` "
        "directly — it bypasses the local registry and queries YARN."
    ),
)
def _get_job_result(req: JobIdRequest):
    return get_job_result(req.job_id)


@app.post(
    "/get_job_logs",
    operation_id="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.\n\n"
        "**job_id accepts BOTH identifiers** returned by "
        "confirm_submit_job: the local job_id (12-char hex, e.g. "
        "'a1b2c3d4e5f6') and the YARN application_id (e.g. "
        "'application_17400000001_0001'). The lookup is by job_id first, "
        "then by application_id. **If you pass a YARN application_id and "
        "the app is NOT in the local JobStore, this tool returns HTTP 400** "
        "(not 404) with a hint message naming the right external tool. "
        "**For YARN applications NOT submitted through this service** "
        "(no local JobStore record), use "
        "`get_external_job_logs(application_id, connection_name, tail_chars)` "
        "directly — it bypasses the local registry and queries YARN."
    ),
)
def _get_job_logs(req: GetJobLogsRequest):
    return get_job_logs(req.job_id, tail_chars=req.tail_chars)


@app.post(
    "/kill_job",
    operation_id="kill_job",
    summary="Kill a running job",
    description=(
        "PUT state=KILLED to YARN REST API for the job's application_id. "
        "**job_id accepts BOTH identifiers** returned by "
        "confirm_submit_job: the local job_id (12-char hex) and the YARN "
        "application_id. The lookup is by job_id first, then by application_id. "
        "**If you pass a YARN application_id and the app is NOT in the "
        "local JobStore, this tool returns HTTP 400** (not 404) with a "
        "hint pointing to the YARN CLI / UI. **This tool only works for "
        "jobs submitted through this service**; there is no external "
        "equivalent. For YARN applications you did not submit here, use "
        "the YARN CLI / UI directly to kill them."
    ),
)
def _kill_job(req: JobIdRequest):
    return kill_job(req.job_id)


# --- External YARN job tools (bypass JobStore) ---

@app.post(
    "/get_external_job_logs",
    operation_id="get_external_job_logs",
    summary="Query YARN logs for an application not submitted through this service",
    description=(
        "Fetch aggregated container logs for a YARN application using its "
        "application_id and a saved Connection. This bypasses the local JobStore, "
        "so it works for jobs submitted outside this MCP service. "
        "application_id format is 'application_<14-digit-timestamp>_<sequence>'. "
        "For jobs submitted via this service, use get_job_logs(job_id=...) instead."
    ),
)
def _get_external_job_logs(req: ExternalJobLogsRequest):
    return get_external_job_logs(req.application_id, req.connection_name, req.tail_chars)


@app.post(
    "/get_external_job_status",
    operation_id="get_external_job_status",
    summary="Query YARN status for an application not submitted through this service",
    description=(
        "Return the YARN application state and raw REST response for an "
        "application using its application_id and a saved Connection. "
        "This bypasses the local JobStore, so it works for jobs submitted "
        "outside this MCP service. For jobs submitted via this service, "
        "use get_job_status(job_id=...) instead."
    ),
)
def _get_external_job_status(req: ExternalJobStatusRequest):
    return get_external_job_status(req.application_id, req.connection_name)


@app.post(
    "/get_external_job_result",
    operation_id="get_external_job_result",
    summary="Query YARN terminal result for an application not submitted through this service",
    description=(
        "Return a terminal-oriented view (final_status, diagnostics, tracking_url, "
        "started_time, finished_time) for a YARN application using its application_id "
        "and a saved Connection. This bypasses the local JobStore. "
        "For jobs submitted via this service, use get_job_result(job_id=...) instead."
    ),
)
def _get_external_job_result(req: ExternalJobResultRequest):
    return get_external_job_result(req.application_id, req.connection_name)


# --- Connection management tools ---

@app.post(
    "/save_connection",
    operation_id="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. Referenced by "
        "prepare_submit_job via the connection parameter, and by the 3 "
        "get_external_* tools via connection_name. **The Connection's "
        "yarn_rm_url is required for get_external_* to work** — spark-submit "
        "can discover the RM for submissions, but direct YARN REST queries "
        "need an explicit URL. Saving with an existing name overwrites the "
        "record in place (no version history)."
    ),
)
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",
    operation_id="list_connections",
    summary="List all saved Spark connections",
    description=(
        "Return every Connection in the registry (model_dump form). Call "
        "this before save_connection to see existing names (saving with an "
        "existing name overwrites), or after save_connection to verify the "
        "record you just stored."
    ),
)
def _list_connections(_req: EmptyRequest = EmptyRequest()):
    return list_connections()


@app.post(
    "/get_connection",
    operation_id="get_connection",
    summary="Get a single connection by name",
    description=(
        "Return the Connection record, or 404 if not found. Useful to "
        "verify a connection was saved correctly, or to inspect the "
        "resolved yarn_rm_url / auth_type before submitting a job or "
        "calling one of the get_external_* tools."
    ),
)
def _get_connection(req: ConnectionNameRequest):
    return get_connection(req.name)


@app.post(
    "/update_connection",
    operation_id="update_connection",
    summary="Update an existing connection's fields",
    description=(
        "Apply a partial update (PATCH) to an existing Connection record. "
        "Only the fields you provide are changed; the rest are kept as-is. "
        "The `name` is the immutable identifier (use delete_connection + "
        "save_connection to rename).\n\n"
        "To CLEAR an optional field (e.g. remove `yarn_rm_url`), use "
        "delete_connection followed by save_connection with the field omitted. "
        "This tool cannot clear fields — only replace them.\n\n"
        "Returns the full updated Connection record. 404 if no Connection "
        "with the given name exists."
    ),
)
def _update_connection(req: UpdateConnectionRequest):
    fields = req.model_dump(exclude_none=True)
    fields.pop("name", None)  # name is the identity, not a field to patch
    return update_connection(name=req.name, **fields)


@app.post(
    "/delete_connection",
    operation_id="delete_connection",
    summary="Delete a saved connection",
    description=(
        "Remove a Connection by name. 404 if not found. Deleting a "
        "Connection does NOT affect any pending submission or running job "
        "that already references it (the connection details are snapshotted "
        "at prepare_submit_job time, and YARN holds the live submission "
        "state). New prepare_submit_job calls will fail until you re-save "
        "the connection with the same name."
    ),
)
def _delete_connection(req: ConnectionNameRequest):
    return delete_connection(req.name)


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

@app.post(
    "/write_job_file",
    operation_id="write_job_file",
    summary="Write LLM-authored PySpark code to disk",
    description=(
        "Persist PySpark code you've already written in your context to a "
        "timestamped file under SPARK_EXECUTOR_JOBS_DIR (default "
        "./data/jobs/). Returns the absolute path to pass as the "
        "script_path argument of prepare_submit_job. The two-step pattern "
        "(write the file, then prepare) means the user can review the "
        "file via read_job_file before anything runs.\n\n"
        "Prerequisite: you should have already composed the PySpark code "
        "in your own context before calling this tool — it only persists "
        "code, it does not generate it. Code is run through the SQL safety "
        "policy (SELECT/INSERT only) before being written; forbidden "
        "statements cause a 400 with details about which line broke the "
        "policy."
    ),
)
def _write_job_file(req: WriteJobFileRequest):
    return write_job_file(req.code)


@app.post(
    "/read_job_file",
    operation_id="read_job_file",
    summary="Read the contents of an existing PySpark script",
    description=(
        "Returns the text content of an existing script file at the given "
        "path. Caps reads at 1 MB. Typical use: after write_job_file "
        "returns a path, call read_job_file on that path to inspect what "
        "was actually written, before deciding to prepare_submit_job or "
        "update_job_file."
    ),
)
def _read_job_file(req: ReadJobFileRequest):
    return read_job_file(req.script_path)


@app.post(
    "/update_job_file",
    operation_id="update_job_file",
    summary="Overwrite an existing PySpark script with new content",
    description=(
        "Replaces the entire content of an existing script file. Path must "
        "be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
        "to) — protects against overwriting host-mounted configs or other "
        "non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
        "edit the content (LLM or human), update_job_file, then "
        "prepare_submit_job with the same path."
    ),
)
def _update_job_file(req: UpdateJobFileRequest):
    return update_job_file(req.script_path, req.content)

# --- HTTP fetch proxy (host allowlist via Connection.yarn_rm_url) ---

@app.post(
    "/fetch_url",
    operation_id="fetch_url",
    summary="Fetch a URL on the cluster's network and return the body",
    description=(
        "Proxy an HTTP GET to a URL on the cluster's network, returning the "
        "response body. Useful when the agent is on a different network from "
        "the cluster and cannot reach YARN tracking pages, Spark History "
        "Server, or NodeManager web UIs directly.\n\n"
        "**Security constraints:** the URL host must share at least 2 labels "
        "of suffix with the named Connection's yarn_rm_url host (e.g. if "
        "yarn_rm_url is 'rm.prod.internal:8088', you may fetch "
        "'http://nm01.prod.internal:8042/...' but NOT 'http://evil.com/...'). "
        "IP literals (10.0.0.1, ::1) and non-http(s) schemes (file://, "
        "gopher://) are rejected. The Connection's saved auth is reused, so "
        "the agent does not need cluster credentials.\n\n"
        "**Limits:** response body capped at 1 MB (truncated=true if larger), "
        "30s timeout, redirects followed."
    ),
)
def _fetch_url(req: FetchUrlRequest):
    return fetch_url(req.url, req.connection_name)