# 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>_'. " "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)