mcp-server/spark_executor/tools/requests.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen

Pydantic request models for the FastAPI route layer. The underlying tool
functions in tools/*.py still take keyword arguments; these models exist only
so fastapi-mcp can call the routes via tools/call (which sends args as a
JSON body) without 422-ing on dict-typed parameters like spark_conf.
"""
from pydantic import BaseModel, Field


class EmptyRequest(BaseModel):
    """Used for tools that take no arguments (list_connections, list_pending_jobs)."""
    pass


class SaveConnectionRequest(BaseModel):
    name: str
    master: str
    deploy_mode: str = "cluster"
    yarn_rm_url: str | None = None
    spark_conf: dict[str, str] | None = None
    ssl_verify: bool | None = None
    ssl_ca_bundle: str | None = None
    auth_type: str = "none"
    auth_user: str | None = None
    auth_password: str | None = None
    auth_principal: str | None = None
    auth_keytab: str | None = None


class PrepareSubmitJobRequest(BaseModel):
    connection: str
    app_name: str = Field(
        ...,
        description="Human-readable application name for tracking the pending submission.",
    )
    script_path: str = Field(
        ...,
        description=(
            "Absolute path to the PySpark script inside the container's "
            "filesystem. Must point at an existing regular file. For "
            "LLM-generated code, call write_job_file(code=...) first "
            "and pass the returned script_path here. For pre-existing "
            "files on the host, mount them via a docker volume and pass "
            "the in-container path. Returns 400 with a remediation hint "
            "if the path is missing or not a file."
        ),
    )
    queue: str = Field(
        default="default",
        description="YARN queue to submit to. Must be explicitly confirmed by the caller.",
    )
    executor_memory: str = Field(
        default="2G",
        description="Executor memory, e.g. '2G'. Must be explicitly confirmed by the caller.",
    )
    executor_cores: int = Field(
        default=2,
        description="Number of cores per executor. Must be explicitly confirmed by the caller.",
    )
    num_executors: int = Field(
        default=2,
        description="Total number of executors. Must be explicitly confirmed by the caller.",
    )
    extra_args: dict[str, str] | None = Field(
        default=None,
        description="Additional spark-submit flags (e.g. jars, py-files) confirmed at prepare time.",
    )


class PendingIdRequest(BaseModel):
    pending_id: str


class UpdatePendingJobRequest(BaseModel):
    pending_id: str
    script_path: str | None = Field(
        default=None,
        description="Optional new absolute path to the PySpark script. If provided, the file must exist and pass SQL guard.",
    )
    queue: str | None = None
    executor_memory: str | None = None
    executor_cores: int | None = None
    num_executors: int | None = None
    app_name: str | None = None
    extra_args: dict[str, str] | None = None


class JobIdRequest(BaseModel):
    job_id: str


class GetJobLogsRequest(BaseModel):
    job_id: str
    tail_chars: int = 5000


class ConnectionNameRequest(BaseModel):
    name: str


class WriteJobFileRequest(BaseModel):
    code: str = Field(
        ...,
        description=(
            "Full PySpark source code to write to disk. The LLM is expected "
            "to have already written this code in its own context; this tool "
            "persists it to a file under SPARK_EXECUTOR_JOBS_DIR so "
            "spark-submit can see it. The returned script_path is what you "
            "pass to prepare_submit_job."
        ),
    )


class ReadJobFileRequest(BaseModel):
    script_path: str = Field(
        ...,
        description=(
            "Absolute path to a PySpark script inside the container's "
            "filesystem. Must point at an existing regular file."
        ),
    )


class UpdateJobFileRequest(BaseModel):
    script_path: str = Field(
        ...,
        description=(
            "Absolute path to an existing PySpark script inside the "
            "container's filesystem. Must be under SPARK_EXECUTOR_JOBS_DIR "
            "(the same dir write_job_file writes to) — protects against "
            "overwriting host-mounted configs or other critical files."
        ),
    )
    content: str = Field(
        ...,
        description=(
            "New file content (replaces the file in full; no merge/diff). "
            "Maximum 1 MB to keep the MCP response bounded."
        ),
    )