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 = Field(
        ...,
        description=(
            "Unique connection name. Referenced by prepare_submit_job.connection "
            "and get_external_*.connection_name. Saving with an existing name "
            "overwrites that record."
        ),
    )
    master: str = Field(
        ...,
        description=(
            "Spark master URL. **Required** at the MCP layer (even though "
            "the underlying Connection model has a default of 'yarn'). "
            "Other valid values: 'yarn' (default for YARN), 'spark://host:port' "
            "(Standalone), 'k8s://...', 'mesos://...', 'local[N]' or 'local[*]'. "
            "The validator rejects 'yarn-cluster' and bare 'http://...' URLs — "
            "those are common typos. The literal 'yarn' (not 'yarn-cluster') is "
            "what spark-submit wants for --master."
        ),
    )
    deploy_mode: str = Field(
        default="cluster",
        description=(
            "Spark deploy mode. 'cluster' (default, driver runs in YARN) or "
            "'client' (driver runs where spark-submit is invoked). Most YARN "
            "production submissions use 'cluster'."
        ),
    )
    yarn_rm_url: str | None = Field(
        default=None,
        description=(
            "YARN ResourceManager REST base URL, e.g. 'http://rm-host:8088'. "
            "**Required** for the 3 get_external_* tools to query YARN "
            "directly. Optional for prepare_submit_job — spark-submit "
            "discovers the RM via the cluster config when this is unset."
        ),
    )
    spark_conf: dict[str, str] | None = Field(
        default=None,
        description=(
            "Dict of Spark conf key→value pairs, passed as --conf flags to "
            "spark-submit. Example: {'spark.executor.memory': '4g', "
            "'spark.sql.shuffle.partitions': '200'}. None or empty means "
            "no extra --conf flags."
        ),
    )
    ssl_verify: bool | None = Field(
        default=None,
        description=(
            "Whether to verify the YARN RM TLS certificate. None (default) "
            "falls back to the global setting; explicit True/False overrides "
            "the global default for this connection. Set False only for "
            "self-signed dev clusters."
        ),
    )
    ssl_ca_bundle: str | None = Field(
        default=None,
        description=(
            "Absolute path to a CA bundle file for YARN RM TLS verification. "
            "Only relevant when the RM uses a private CA. Ignored when "
            "ssl_verify=False."
        ),
    )
    auth_type: str = Field(
        default="none",
        description=(
            "Authentication mode for YARN REST calls. One of: 'none' "
            "(default, no auth header), 'simple' (pseudo-auth, "
            "auth_user required), 'basic' (HTTP Basic, auth_user + "
            "auth_password required), 'kerberos' (SPNEGO via the system "
            "ticket cache — run kinit beforehand)."
        ),
    )
    auth_user: str | None = Field(
        default=None,
        description=(
            "Username for auth_type='simple' or 'basic'. Ignored when "
            "auth_type='none' or 'kerberos'."
        ),
    )
    auth_password: str | None = Field(
        default=None,
        description=(
            "Password for auth_type='basic'. Ignored otherwise. Sent on "
            "every YARN REST request — store with care."
        ),
    )
    auth_principal: str | None = Field(
        default=None,
        description=(
            "Kerberos principal (e.g. 'user@REALM'). Display/audit only; "
            "the actual SPNEGO handshake uses the system ticket cache. "
            "Run `kinit <principal>` on the host before invoking the tools."
        ),
    )
    auth_keytab: str | None = Field(
        default=None,
        description=(
            "Absolute path to a Kerberos keytab file. Optional convenience "
            "for 'kinit -kt' workflows. The service does NOT auto-initialize "
            "from the keytab — you must `kinit -kt <auth_keytab> <auth_principal>` "
            "yourself before calling the tools."
        ),
    )


class PrepareSubmitJobRequest(BaseModel):
    connection: str = Field(
        ...,
        description=(
            "Name of a saved Connection (call save_connection first, or "
            "list_connections to see available names). The Connection's "
            "master / deploy_mode / spark_conf / yarn_rm_url are snapshotted "
            "into the pending submission at prepare time, so editing the "
            "Connection afterwards does NOT retarget this pending job."
        ),
    )
    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 = Field(
        ...,
        description=(
            "ID of a pending submission. Format: 'p_' + 12 hex chars, "
            "e.g. 'p_a1b2c3d4e5f6'. Returned by prepare_submit_job; "
            "visible via list_pending_jobs."
        ),
    )


class UpdatePendingJobRequest(BaseModel):
    pending_id: str = Field(
        ...,
        description=(
            "ID of the pending submission to modify. Format: 'p_' + 12 hex "
            "chars, e.g. 'p_a1b2c3d4e5f6'. Only PENDING submissions can "
            "be updated — once SUBMITTED, CANCELLED, or FAILED, the "
            "pending is terminal."
        ),
    )
    script_path: str | None = Field(
        default=None,
        description=(
            "Optional new absolute path to the PySpark script. If provided, "
            "the file must exist and pass the SQL guard. Omit to keep the "
            "current script_path."
        ),
    )
    queue: str | None = Field(
        default=None,
        description=(
            "New YARN queue name. Omit to keep the current value (PATCH "
            "semantics: only fields you provide are changed)."
        ),
    )
    executor_memory: str | None = Field(
        default=None,
        description=(
            "New executor memory, e.g. '4G'. Omit to keep the current value."
        ),
    )
    executor_cores: int | None = Field(
        default=None,
        description="New cores per executor. Omit to keep the current value.",
    )
    num_executors: int | None = Field(
        default=None,
        description="New total executor count. Omit to keep the current value.",
    )
    app_name: str | None = Field(
        default=None,
        description=(
            "New human-readable application name (visible in YARN UI). "
            "Omit to keep the current value."
        ),
    )
    extra_args: dict[str, str] | None = Field(
        default=None,
        description=(
            "Replacement dict of additional spark-submit flags (e.g. "
            "{'jars': '/path/to.jar'}). Unlike the scalar fields, providing "
            "this REPLACES the entire dict — it is not deep-merged. Omit to "
            "keep the current value."
        ),
    )


class JobIdRequest(BaseModel):
    job_id: str = Field(
        ...,
        description=(
            "Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') or "
            "the YARN application_id (e.g. 'application_17400000001_0001') of "
            "a job submitted through this service. The lookup tries job_id "
            "first, then application_id. For YARN applications not submitted "
            "here, use the get_external_* tools instead."
        ),
    )


class GetJobLogsRequest(BaseModel):
    job_id: str = Field(
        ...,
        description=(
            "Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') "
            "or the YARN application_id (e.g. 'application_17400000001_0001') "
            "of a job submitted through this service. For YARN applications "
            "not submitted here, use get_external_job_logs instead."
        ),
    )
    tail_chars: int = Field(
        default=5000,
        description=(
            "Return only the last N characters of the aggregated container "
            "logs. Default 5000. Use a larger value if the head of the log "
            "(stack traces, driver errors) is being truncated."
        ),
    )


class ConnectionNameRequest(BaseModel):
    name: str = Field(
        ...,
        description=(
            "Name of a saved Connection. Use list_connections to see "
            "available names. Saving with this name updates an existing "
            "record (see save_connection)."
        ),
    )


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."
        ),
    )


class ExternalJobLogsRequest(BaseModel):
    application_id: str = Field(
        ...,
        description=(
            "YARN application_id of the external job, format "
            "'application_<14-digit-timestamp>_<sequence>' (e.g. "
            "'application_1740000000001_0001'). This tool is for jobs "
            "NOT submitted through this MCP service — for those, use "
            "get_job_logs(job_id=...) instead."
        ),
    )
    connection_name: str = Field(
        ...,
        description=(
            "Name of a saved Connection (see list_connections) pointing "
            "at the YARN cluster where the application ran."
        ),
    )
    tail_chars: int = Field(
        default=5000,
        description="Return only the last N characters of the aggregated container logs.",
    )


class ExternalJobStatusRequest(BaseModel):
    application_id: str = Field(
        ...,
        description=(
            "YARN application_id of the external job, format "
            "'application_<14-digit-timestamp>_<sequence>'. Use "
            "get_job_status(job_id=...) for jobs submitted through this service."
        ),
    )
    connection_name: str = Field(
        ...,
        description="Name of a saved Connection pointing at the YARN cluster.",
    )


class ExternalJobResultRequest(BaseModel):
    application_id: str = Field(
        ...,
        description=(
            "YARN application_id of the external job, format "
            "'application_<14-digit-timestamp>_<sequence>'. Use "
            "get_job_result(job_id=...) for jobs submitted through this service."
        ),
    )
    connection_name: str = Field(
        ...,
        description="Name of a saved Connection pointing at the YARN cluster.",
    )