Compare commits
19
Commits
e285dc0f66
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
307da276bf | ||
|
|
a460b750a8 | ||
|
|
6109c6d11a | ||
|
|
3b698e1bdc | ||
|
|
a5e6d1ffe1 | ||
|
|
16fe011fa0 | ||
|
|
59cab3346e | ||
|
|
8ededc50a9 | ||
|
|
f43e5b6403 | ||
|
|
585a1f6e43 | ||
|
|
deafb5a26b | ||
|
|
7d4e512cb0 | ||
|
|
7fbad97a87 | ||
|
|
a5b9539663 | ||
|
|
6d7387b022 | ||
|
|
0291f36b01 | ||
|
|
627e70f697 | ||
|
|
f8536b63ad | ||
|
|
6cf68439a2 |
@@ -20,3 +20,6 @@ data/
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# macOS
|
||||
.DS_Store
|
||||
|
||||
@@ -40,7 +40,7 @@
|
||||
| MCP 暴露 | fastapi-mcp |
|
||||
| 数据模型 | Pydantic v2 |
|
||||
| HTTP 客户端 | httpx、httpx-kerberos |
|
||||
| Spark 集成 | `spark-submit`、YARN REST API、`yarn logs` |
|
||||
| Spark 集成 | `spark-submit`、YARN REST API (httpx 直连) |
|
||||
| 日志 | loguru |
|
||||
| 进程管理 | uvicorn、gunicorn |
|
||||
| 包管理 | uv |
|
||||
@@ -104,6 +104,7 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
|
||||
| `list_connections` | 列出所有连接配置 |
|
||||
| `get_connection` | 按名称读取连接配置 |
|
||||
| `delete_connection` | 删除连接配置 |
|
||||
| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) |
|
||||
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
|
||||
| `read_job_file` | 读取已存在的 PySpark 脚本内容 |
|
||||
| `update_job_file` | 覆盖更新已存在的 PySpark 脚本 |
|
||||
@@ -117,6 +118,11 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
|
||||
| `get_job_result` | 查询终态结果视图 |
|
||||
| `get_job_logs` | 拉取 YARN 聚合日志 |
|
||||
| `kill_job` | Kill YARN application |
|
||||
| `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态(按 `application_id` + `connection_name`) |
|
||||
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
|
||||
| `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 |
|
||||
| `list_applications` | 列出 YARN 上所有应用(按 `state` / `queue` / `limit` 过滤),绕过 JobStore |
|
||||
| `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.url_allowlist` glob allowlist 约束, 空则全拒) |
|
||||
|
||||
### Files MCP 工具
|
||||
|
||||
@@ -215,6 +221,30 @@ get_job_logs
|
||||
kill_job
|
||||
```
|
||||
|
||||
## RM / SHS API 端点
|
||||
|
||||
本服务**不再 shell-out `yarn` CLI**,所有集群交互都通过 `httpx` 直接打 REST(由 `spark_executor/core/yarn_client.py` 统一封装)。鉴权走 `Connection.auth_type`(`none` / `simple` / `basic` / `kerberos` + SPNEGO),SSL 走 `Connection.ssl_verify` / `ssl_ca_bundle`。
|
||||
|
||||
### YARN ResourceManager
|
||||
|
||||
| 方法 | 端点 | 用途 |
|
||||
| --- | --- | --- |
|
||||
| `GET` | `/ws/v1/cluster/apps?state=&queue=&limit=` | 列出 YARN 应用(支持 `state` / `queue` / `limit` 过滤),由 `list_applications` 触发 |
|
||||
| `GET` | `/ws/v1/cluster/apps/{appid}` | 取单个 app 的 state + amContainerLogs;日志降级路径会复用 |
|
||||
| `GET` | `/ws/v1/cluster/apps/{appid}/aggregated-logs` | 拉聚合后的 container 日志,404/501 时自动降级到 amContainerLogs 路径 |
|
||||
| `PUT` | `/ws/v1/cluster/apps/{appid}/state`,body `{"state":"KILLED"}` | 杀应用,由 `kill_job` 触发 |
|
||||
|
||||
> **日志降级路径**会**手动**跟随 3xx Location 跳到 NodeManager(最多 3 次)抓 AM driver 的 `/stdout`。原因:httpx 默认跨主机重定向会丢掉 `Authorization`,RM→NM 的 307 在很多集群会因此 401/403。driver 日志可以这样取,**executor 日志仍然依赖 `yarn.log-aggregation-enable=true`**,这个不走降级。
|
||||
|
||||
### Spark History Server
|
||||
|
||||
**当前没有任何内置调用。** `Connection.history_server_url` 字段只是原样存储,留作未来 SHS 专用工具消费(`models.py` 的 docstring 明确写明)。
|
||||
|
||||
要走 SHS API(典型路径 `/api/v1/applications/{id}/jobs`、`/api/v1/applications/{id}/stages` 等),**唯一**方式是用 `fetch_url` 工具:
|
||||
|
||||
1. 把 SHS 主机 glob 加到对应 `Connection.url_allowlist`(单标签主机如 `ccam*`,或多级域名如 `*.hadoop.internal`)。
|
||||
2. 调用 `fetch_url` 时传入完整 SHS URL,鉴权和 SSL 复用同一条 Connection 的配置。
|
||||
|
||||
## 配置项
|
||||
|
||||
### 应用配置
|
||||
|
||||
@@ -77,6 +77,29 @@ class ConnectionStore:
|
||||
self._dump(records)
|
||||
logger.info(f"connection saved name={conn.name} master={conn.master}")
|
||||
|
||||
def update(self, name: str, **fields) -> Connection:
|
||||
"""Apply `fields` to the existing Connection identified by `name` and persist.
|
||||
|
||||
PATCH semantics: only fields explicitly passed in `fields` are changed.
|
||||
Use Pydantic's `model_copy(update=fields)` to apply the patch.
|
||||
|
||||
Raises KeyError if no Connection with `name` exists.
|
||||
Raises pydantic.ValidationError if the patched Connection is invalid
|
||||
(e.g. `master='http://...'` fails the master validator).
|
||||
"""
|
||||
with self._lock:
|
||||
records = self._load()
|
||||
if name not in records:
|
||||
raise KeyError(f"Connection not found: {name}")
|
||||
existing = records[name]
|
||||
patched = Connection.model_validate(existing.model_copy(update=fields).model_dump())
|
||||
records[name] = patched
|
||||
self._dump(records)
|
||||
logger.info(
|
||||
f"connection updated name={name} fields={sorted(fields.keys())}"
|
||||
)
|
||||
return patched
|
||||
|
||||
def delete(self, name: str) -> bool:
|
||||
with self._lock:
|
||||
records = self._load()
|
||||
|
||||
@@ -14,7 +14,7 @@ Resolution order for the output directory:
|
||||
"""
|
||||
import os
|
||||
import secrets
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from common.config import settings
|
||||
from common.logging import logger
|
||||
@@ -63,7 +63,7 @@ def write_job_file(code: str, jobs_dir: str | None = None) -> str:
|
||||
|
||||
effective_dir = resolve_jobs_dir(jobs_dir)
|
||||
os.makedirs(effective_dir, exist_ok=True)
|
||||
stamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
|
||||
stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
|
||||
name = f"job_{stamp}_{secrets.token_hex(3)}.py"
|
||||
path = os.path.join(effective_dir, name)
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
@@ -37,7 +37,7 @@ class PendingStore:
|
||||
return Path(self._data_dir) / self._dir_name
|
||||
|
||||
def _date_path(self, created_at: datetime) -> Path:
|
||||
# created_at is datetime.utcnow(), so the shard date is a UTC date.
|
||||
# created_at is timezone-aware UTC, so the shard date is a UTC date.
|
||||
return self.dir_path / f"{created_at.date().isoformat()}.json"
|
||||
|
||||
def _legacy_path(self) -> Path:
|
||||
|
||||
@@ -121,13 +121,14 @@ def _base_url(yarn_rm_url: str | None) -> str:
|
||||
|
||||
|
||||
def _request(method: str, url: str, *, json_body: dict | None = None,
|
||||
timeout: float = 30.0, verify: bool | str = True,
|
||||
auth: httpx.Auth | None = None) -> httpx.Response:
|
||||
params: dict[str, str] | None = None, timeout: float = 30.0,
|
||||
verify: bool | str = True, auth: httpx.Auth | None = None) -> httpx.Response:
|
||||
headers = {"Accept": "application/json"}
|
||||
logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
|
||||
try:
|
||||
resp = httpx.request(
|
||||
method, url, json=json_body, headers=headers, timeout=timeout, verify=verify, auth=auth
|
||||
method, url, json=json_body, params=params, headers=headers,
|
||||
timeout=timeout, verify=verify, auth=auth
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
logger.error(f"YARN {method} {url} failed: {exc}")
|
||||
@@ -257,3 +258,51 @@ def kill_application(application_id: str, config: YarnClientConfig) -> None:
|
||||
logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
|
||||
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
|
||||
logger.info(f"YARN kill {application_id} -> ok")
|
||||
|
||||
|
||||
def list_applications(
|
||||
config: YarnClientConfig,
|
||||
*,
|
||||
state: str | None = None,
|
||||
queue: str | None = None,
|
||||
limit: int | None = None,
|
||||
) -> list[dict]:
|
||||
"""List YARN applications, optionally filtered.
|
||||
|
||||
YARN endpoint: GET /ws/v1/cluster/apps?state=...&queue=...&limit=...
|
||||
|
||||
Filters:
|
||||
- state: YARN application state. Common values:
|
||||
"NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING",
|
||||
"FINISHED", "FAILED", "KILLED".
|
||||
Note: "FINISHED" is the umbrella state covering SUCCEEDED/FAILED/KILLED.
|
||||
- queue: YARN queue name
|
||||
- limit: cap on number of returned apps (YARN has no pagination;
|
||||
callers that need a full enumeration should make multiple
|
||||
calls with state=... filters or accept the cap)
|
||||
|
||||
Returns a list of YARN app dicts (each with id, name, user, queue,
|
||||
state, finalStatus, applicationType, startedTime, finishedTime,
|
||||
trackingUrl, progress, etc). Empty list if no apps match.
|
||||
|
||||
Raises YarnError on transport / 4xx / 5xx.
|
||||
"""
|
||||
params: dict[str, str] = {}
|
||||
if state is not None:
|
||||
params["state"] = state
|
||||
if queue is not None:
|
||||
params["queue"] = queue
|
||||
if limit is not None:
|
||||
params["limit"] = str(limit)
|
||||
|
||||
url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps"
|
||||
resp = _request("GET", url, params=params,
|
||||
verify=config.verify_for_httpx(),
|
||||
auth=config.auth_for_httpx())
|
||||
if resp.status_code >= 400:
|
||||
raise YarnError(
|
||||
f"YARN list applications failed: {resp.status_code} {resp.text[:200]}"
|
||||
)
|
||||
data = resp.json()
|
||||
apps_container = data.get("apps") or {}
|
||||
return apps_container.get("app", []) or []
|
||||
|
||||
@@ -40,6 +40,35 @@ class SubmitResult(BaseModel):
|
||||
tracking_url: str | None = None
|
||||
|
||||
|
||||
class FetchUrlResult(BaseModel):
|
||||
url: str
|
||||
status_code: int
|
||||
content_type: str
|
||||
body: str
|
||||
|
||||
|
||||
class ApplicationSummary(BaseModel):
|
||||
"""A YARN application summary from /ws/v1/cluster/apps.
|
||||
|
||||
Field names are mapped from the YARN JSON keys to clearer
|
||||
snake_case names by the tool function. Unused YARN fields
|
||||
(memorySeconds, vcoreSeconds, preemptedResource*, etc.) are
|
||||
not exposed — the LLM doesn't need them.
|
||||
"""
|
||||
application_id: str
|
||||
name: str
|
||||
user: str
|
||||
queue: str
|
||||
state: str
|
||||
final_status: str | None = None
|
||||
application_type: str | None = None
|
||||
application_tags: str = ""
|
||||
started_time: int = 0
|
||||
finished_time: int = 0
|
||||
tracking_url: str | None = None
|
||||
progress: float | None = None
|
||||
|
||||
|
||||
class Connection(BaseModel):
|
||||
name: str
|
||||
# Defaults to "yarn" because that's the literal string spark-submit wants
|
||||
@@ -61,6 +90,35 @@ class Connection(BaseModel):
|
||||
auth_principal: str | None = None
|
||||
auth_keytab: str | None = None
|
||||
|
||||
url_allowlist: list[str] = Field(
|
||||
default_factory=list,
|
||||
description=(
|
||||
"List of fnmatch glob patterns for hosts the fetch_url tool may access. "
|
||||
"The list is mandatory-opt-in: an empty list (the default) denies all "
|
||||
"hosts, so you must populate it before fetch_url can access any URL. "
|
||||
"Set or change via save_connection (pass url_allowlist on create) or "
|
||||
"update_connection (PATCH the field on an existing connection). "
|
||||
"Useful for clusters whose hostnames do NOT share a common suffix — "
|
||||
"e.g. single-label hosts like 'ccam1'-'ccam99' (configure ['ccam*']) "
|
||||
"or HDFS namenode on a different subdomain ('*.hadoop.internal'). "
|
||||
"Patterns are matched against the URL host only (no port, no path). "
|
||||
"fnmatch rules apply: '*' does NOT match '.', so 'ccam*' matches "
|
||||
"'ccam50' but not 'ccam50.evil.com'."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional Spark History Server (SHS) base URL, e.g. "
|
||||
"'http://history.prod.internal:18080'. Stored as part of the "
|
||||
"connection for reference and to be consumed by future SHS-"
|
||||
"specific tools. Currently **not consumed by any tool** — to "
|
||||
"fetch SHS endpoints today, use fetch_url with the SHS host "
|
||||
"added to url_allowlist. The SHS host and the YARN RM host "
|
||||
"are usually different, so this is independent of yarn_rm_url."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("master")
|
||||
@classmethod
|
||||
def _check_master(cls, v: str) -> str:
|
||||
|
||||
+267
-32
@@ -6,26 +6,41 @@
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
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,
|
||||
list_applications,
|
||||
)
|
||||
from spark_executor.tools.fetch_url import fetch_url
|
||||
from spark_executor.tools.requests import (
|
||||
ConnectionNameRequest,
|
||||
EmptyRequest,
|
||||
WriteJobFileRequest,
|
||||
ExternalJobLogsRequest,
|
||||
ExternalJobStatusRequest,
|
||||
ExternalJobResultRequest,
|
||||
ListApplicationsRequest,
|
||||
FetchUrlRequest,
|
||||
GetJobLogsRequest,
|
||||
JobIdRequest,
|
||||
PendingIdRequest,
|
||||
PrepareSubmitJobRequest,
|
||||
ReadJobFileRequest,
|
||||
SaveConnectionRequest,
|
||||
UpdateConnectionRequest,
|
||||
UpdateJobFileRequest,
|
||||
UpdatePendingJobRequest,
|
||||
)
|
||||
@@ -69,6 +84,20 @@ async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONRespons
|
||||
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
||||
|
||||
|
||||
@app.exception_handler(YarnError)
|
||||
async def _yarnerror_handler(_request: Request, exc: YarnError) -> JSONResponse:
|
||||
"""YARN RM unreachable or returned an error.
|
||||
|
||||
Surfaces as HTTP 502 Bad Gateway (we are a gateway to YARN). The
|
||||
detail includes whatever yarn_client put in the YarnError message
|
||||
(host:port unreachable, HTTP status from YARN, parse failure,
|
||||
etc). The LLM can act on this to distinguish "YARN is down" from
|
||||
"the request was bad" (the latter would be a 400 from
|
||||
_valueerror_handler instead).
|
||||
"""
|
||||
return JSONResponse(status_code=502, content={"detail": str(exc)})
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"status": "ok"}
|
||||
@@ -117,7 +146,9 @@ def _prepare_submit_job(req: PrepareSubmitJobRequest):
|
||||
"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."
|
||||
"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):
|
||||
@@ -128,7 +159,13 @@ def _confirm_submit_job(req: PendingIdRequest):
|
||||
"/list_pending_jobs",
|
||||
operation_id="list_pending_jobs",
|
||||
summary="List all pending submissions",
|
||||
description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).",
|
||||
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()
|
||||
@@ -138,7 +175,13 @@ def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
|
||||
"/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.",
|
||||
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)
|
||||
@@ -186,7 +229,13 @@ def _cancel_pending_job(req: PendingIdRequest):
|
||||
"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."
|
||||
"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):
|
||||
@@ -206,7 +255,13 @@ def _get_job_status(req: JobIdRequest):
|
||||
"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."
|
||||
"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):
|
||||
@@ -225,7 +280,13 @@ def _get_job_result(req: JobIdRequest):
|
||||
"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."
|
||||
"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):
|
||||
@@ -240,13 +301,96 @@ def _get_job_logs(req: GetJobLogsRequest):
|
||||
"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."
|
||||
"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)
|
||||
|
||||
|
||||
@app.post(
|
||||
"/list_applications",
|
||||
operation_id="list_applications",
|
||||
summary="List YARN applications on a cluster, optionally filtered",
|
||||
description=(
|
||||
"Query YARN's /ws/v1/cluster/apps endpoint through the named "
|
||||
"Connection, returning a list of ApplicationSummary records. "
|
||||
"Bypasses the local JobStore — useful for enumerating apps that "
|
||||
"were not submitted through this service.\n\n"
|
||||
"**Filters:** state (YARN state, e.g. 'RUNNING', 'FINISHED', "
|
||||
"'FAILED'), queue (YARN queue name), limit (default 100, max ~10000). "
|
||||
"YARN has no offset-based pagination, so for large clusters combine "
|
||||
"state/queue filters to scope the result. The `FINISHED` state "
|
||||
"covers SUCCEEDED/FAILED/KILLED.\n\n"
|
||||
"Returns an empty list if no apps match. The Connection's auth_type "
|
||||
"/ auth_user / auth_password / ssl_verify / ssl_ca_bundle are reused "
|
||||
"for the request."
|
||||
),
|
||||
)
|
||||
def _list_applications(req: ListApplicationsRequest):
|
||||
return list_applications(
|
||||
req.connection_name,
|
||||
state=req.state,
|
||||
queue=req.queue,
|
||||
limit=req.limit,
|
||||
)
|
||||
|
||||
|
||||
# --- Connection management tools ---
|
||||
|
||||
@app.post(
|
||||
@@ -254,21 +398,39 @@ def _kill_job(req: JobIdRequest):
|
||||
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. Used by prepare_submit_job via the "
|
||||
"connection parameter."
|
||||
"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). When the name already exists, "
|
||||
"only the provided fields are changed (PATCH semantics); omitted "
|
||||
"fields keep their previous values."
|
||||
),
|
||||
)
|
||||
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))
|
||||
fields = req.model_dump(exclude_none=True)
|
||||
name = fields.pop("name")
|
||||
try:
|
||||
get_connection(name)
|
||||
except KeyError:
|
||||
# exclude_none so we don't overwrite the function's default with explicit None
|
||||
return save_connection(name=name, **fields)
|
||||
return update_connection(name=name, **fields)
|
||||
|
||||
|
||||
@app.post(
|
||||
"/list_connections",
|
||||
operation_id="list_connections",
|
||||
summary="List all saved Spark connections",
|
||||
description="Return every Connection in the registry (model_dump form).",
|
||||
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()
|
||||
@@ -278,17 +440,58 @@ def _list_connections(_req: EmptyRequest = EmptyRequest()):
|
||||
"/get_connection",
|
||||
operation_id="get_connection",
|
||||
summary="Get a single connection by name",
|
||||
description="Return the Connection record, or 404 if not found.",
|
||||
description=(
|
||||
"Return the full Connection record, or 404 if not found. **Call this "
|
||||
"whenever you need any cluster-level config** — common lookups: "
|
||||
"yarn_rm_url (YARN RM endpoint), history_server_url (Spark History "
|
||||
"Server), url_allowlist (which hosts fetch_url may access), "
|
||||
"auth_type / auth_user / ssl_verify (for any YARN REST or HTTP call), "
|
||||
"master / deploy_mode / spark_conf (for prepare_submit_job).\n\n"
|
||||
"The response is a full Pydantic model dump — all fields including "
|
||||
"secrets (auth_password, auth_keytab). Treat it as sensitive. "
|
||||
"Useful to verify a connection was saved correctly, or to discover "
|
||||
"the right endpoint to call before invoking get_external_*, "
|
||||
"fetch_url, or list_applications."
|
||||
),
|
||||
)
|
||||
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.",
|
||||
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)
|
||||
@@ -301,16 +504,18 @@ def _delete_connection(req: ConnectionNameRequest):
|
||||
operation_id="write_job_file",
|
||||
summary="Write LLM-authored PySpark code to disk",
|
||||
description=(
|
||||
"Takes a PySpark code string the LLM has already composed in its "
|
||||
"context 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 writes the file, the user can review "
|
||||
"it (via read_job_file), and only then is the job submitted.\n\n"
|
||||
"Note: this tool does NOT generate PySpark code. The calling LLM is "
|
||||
"expected to have already written the code; this tool only persists "
|
||||
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
|
||||
"only) before being written — forbidden statements cause a 400."
|
||||
"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):
|
||||
@@ -323,10 +528,9 @@ def _write_job_file(req: WriteJobFileRequest):
|
||||
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."
|
||||
"path. No size cap. 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):
|
||||
@@ -341,10 +545,41 @@ def _read_job_file(req: ReadJobFileRequest):
|
||||
"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."
|
||||
"non-script files. No size cap. 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 match one of the fnmatch "
|
||||
"glob patterns in the named Connection's url_allowlist. An empty or "
|
||||
"omitted allowlist denies every host. There are no scheme or IP-literal "
|
||||
"guardrails — the allowlist is the only gate — so keep it tight. The "
|
||||
"Connection's saved auth is reused, so the agent does not need cluster "
|
||||
"credentials.\n\n"
|
||||
"**Limits:** 30s timeout, redirects followed.\n\n"
|
||||
"**Errors:** when the host is unreachable (connect refused, DNS "
|
||||
"failure, TLS handshake error, timeout, etc.) this tool returns "
|
||||
"HTTP 400 with the exception class and message in the response "
|
||||
"detail — e.g. `fetch_url could not reach ...: ConnectError: "
|
||||
"Connection refused`. Upstream HTTP 4xx/5xx responses that DID "
|
||||
"come back are returned in the result body with their status code "
|
||||
"preserved (not translated to an error) so you can see what the "
|
||||
"server actually said."
|
||||
),
|
||||
)
|
||||
def _fetch_url(req: FetchUrlRequest):
|
||||
return fetch_url(req.url, req.connection_name)
|
||||
|
||||
@@ -4,47 +4,98 @@
|
||||
@Author :tao.chen
|
||||
"""
|
||||
from common.logging import logger
|
||||
from spark_executor.core.connection_store import ConnectionStore, store
|
||||
from spark_executor.core.connection_store import store
|
||||
from spark_executor.models import Connection
|
||||
|
||||
|
||||
_UNSET = object()
|
||||
|
||||
|
||||
def save_connection(
|
||||
*,
|
||||
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,
|
||||
) -> dict[str, str]:
|
||||
deploy_mode: str = _UNSET, # type: ignore[assignment]
|
||||
yarn_rm_url: str | None = _UNSET, # type: ignore[assignment]
|
||||
spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment]
|
||||
ssl_verify: bool | None = _UNSET, # type: ignore[assignment]
|
||||
ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment]
|
||||
auth_type: str = _UNSET, # type: ignore[assignment]
|
||||
auth_user: str | None = _UNSET, # type: ignore[assignment]
|
||||
auth_password: str | None = _UNSET, # type: ignore[assignment]
|
||||
auth_principal: str | None = _UNSET, # type: ignore[assignment]
|
||||
auth_keytab: str | None = _UNSET, # type: ignore[assignment]
|
||||
url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment]
|
||||
history_server_url: str | None = _UNSET, # type: ignore[assignment]
|
||||
) -> dict[str, object]:
|
||||
logger.debug(
|
||||
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
|
||||
f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}"
|
||||
)
|
||||
conn = Connection(
|
||||
name=name,
|
||||
master=master,
|
||||
deploy_mode=deploy_mode,
|
||||
yarn_rm_url=yarn_rm_url,
|
||||
spark_conf=spark_conf or {},
|
||||
ssl_verify=ssl_verify,
|
||||
ssl_ca_bundle=ssl_ca_bundle,
|
||||
auth_type=auth_type,
|
||||
auth_user=auth_user,
|
||||
auth_password=auth_password,
|
||||
auth_principal=auth_principal,
|
||||
auth_keytab=auth_keytab,
|
||||
f"save_connection enter name={name} master={master} "
|
||||
f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}"
|
||||
)
|
||||
existing = store.get(name)
|
||||
if existing is not None:
|
||||
fields: dict[str, object] = {"master": master}
|
||||
if deploy_mode is not _UNSET:
|
||||
fields["deploy_mode"] = deploy_mode
|
||||
if yarn_rm_url is not _UNSET:
|
||||
fields["yarn_rm_url"] = yarn_rm_url
|
||||
if spark_conf is not _UNSET:
|
||||
fields["spark_conf"] = spark_conf or {}
|
||||
if ssl_verify is not _UNSET:
|
||||
fields["ssl_verify"] = ssl_verify
|
||||
if ssl_ca_bundle is not _UNSET:
|
||||
fields["ssl_ca_bundle"] = ssl_ca_bundle
|
||||
if auth_type is not _UNSET:
|
||||
fields["auth_type"] = auth_type
|
||||
if auth_user is not _UNSET:
|
||||
fields["auth_user"] = auth_user
|
||||
if auth_password is not _UNSET:
|
||||
fields["auth_password"] = auth_password
|
||||
if auth_principal is not _UNSET:
|
||||
fields["auth_principal"] = auth_principal
|
||||
if auth_keytab is not _UNSET:
|
||||
fields["auth_keytab"] = auth_keytab
|
||||
if url_allowlist is not _UNSET:
|
||||
fields["url_allowlist"] = url_allowlist or []
|
||||
if history_server_url is not _UNSET:
|
||||
fields["history_server_url"] = history_server_url
|
||||
return update_connection(name, **fields)
|
||||
new_fields: dict[str, object] = {
|
||||
"name": name,
|
||||
"master": master,
|
||||
"deploy_mode": deploy_mode if deploy_mode is not _UNSET else "cluster",
|
||||
"yarn_rm_url": yarn_rm_url if yarn_rm_url is not _UNSET else None,
|
||||
"spark_conf": (spark_conf if spark_conf is not _UNSET else None) or {},
|
||||
"ssl_verify": ssl_verify if ssl_verify is not _UNSET else None,
|
||||
"ssl_ca_bundle": ssl_ca_bundle if ssl_ca_bundle is not _UNSET else None,
|
||||
"auth_type": auth_type if auth_type is not _UNSET else "none",
|
||||
"auth_user": auth_user if auth_user is not _UNSET else None,
|
||||
"auth_password": auth_password if auth_password is not _UNSET else None,
|
||||
"auth_principal": auth_principal if auth_principal is not _UNSET else None,
|
||||
"auth_keytab": auth_keytab if auth_keytab is not _UNSET else None,
|
||||
"url_allowlist": (url_allowlist if url_allowlist is not _UNSET else None) or [],
|
||||
"history_server_url": history_server_url if history_server_url is not _UNSET else None,
|
||||
}
|
||||
conn = Connection(**new_fields)
|
||||
store.save(conn)
|
||||
return {"name": name, "status": "SAVED"}
|
||||
|
||||
|
||||
def update_connection(name: str, **fields) -> dict[str, object]:
|
||||
"""Update an existing Connection's mutable fields.
|
||||
|
||||
`name` is the identifier (immutable). PATCH semantics: only the fields
|
||||
you pass are changed. To clear an optional field (e.g. `yarn_rm_url`),
|
||||
use `delete_connection(name=...)` followed by `save_connection(...)`.
|
||||
|
||||
Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf,
|
||||
ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password,
|
||||
auth_principal, auth_keytab, url_allowlist, history_server_url.
|
||||
"""
|
||||
logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}")
|
||||
return store.update(name, **fields).model_dump()
|
||||
|
||||
|
||||
def list_connections() -> list[dict[str, object]]:
|
||||
logger.debug("list_connections enter")
|
||||
return [c.model_dump() for c in store.list_all()]
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/7/8
|
||||
@Author :tao.chen
|
||||
|
||||
Tools for inspecting YARN applications NOT submitted through this service.
|
||||
These bypass the local JobStore and require the caller to supply both the
|
||||
YARN application_id and the name of a saved Connection.
|
||||
"""
|
||||
import json
|
||||
|
||||
from common.logging import logger
|
||||
from spark_executor.core.yarn_client import (
|
||||
YarnClientConfig,
|
||||
get_application_logs,
|
||||
get_application_status,
|
||||
list_applications as list_applications_yarn, # alias to avoid collision
|
||||
)
|
||||
from spark_executor.models import ApplicationSummary, JobResult, JobStatus
|
||||
from spark_executor.tools.connections import store as conn_store
|
||||
|
||||
|
||||
def get_external_job_logs(application_id: str, connection_name: str, tail_chars: int = 5000) -> str:
|
||||
"""Fetch aggregated container logs for a YARN application by ID."""
|
||||
logger.debug(f"get_external_job_logs enter application_id={application_id} connection_name={connection_name} tail_chars={tail_chars}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
full = get_application_logs(application_id, config)
|
||||
tailed = full[-tail_chars:] if len(full) > tail_chars else full
|
||||
logger.info(
|
||||
f"get_external_job_logs ok application_id={application_id} connection_name={connection_name} "
|
||||
f"full_chars={len(full)} returned_chars={len(tailed)}"
|
||||
)
|
||||
return tailed
|
||||
|
||||
|
||||
def get_external_job_status(application_id: str, connection_name: str) -> JobStatus:
|
||||
"""Query YARN for an external application's current status."""
|
||||
logger.debug(f"get_external_job_status enter application_id={application_id} connection_name={connection_name}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
state, raw = get_application_status(application_id, config)
|
||||
logger.info(f"get_external_job_status ok application_id={application_id} connection_name={connection_name} state={state}")
|
||||
return JobStatus(application_id=application_id, state=state, raw=raw)
|
||||
|
||||
|
||||
def get_external_job_result(application_id: str, connection_name: str) -> JobResult:
|
||||
"""Query YARN for an external application's terminal result view."""
|
||||
logger.debug(f"get_external_job_result enter application_id={application_id} connection_name={connection_name}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
state, raw = get_application_status(application_id, config)
|
||||
app = json.loads(raw).get("app", {})
|
||||
result = JobResult(
|
||||
application_id=application_id,
|
||||
state=state,
|
||||
final_status=app.get("finalStatus"),
|
||||
diagnostics=app.get("diagnostics"),
|
||||
tracking_url=app.get("trackingUrl"),
|
||||
started_time=app.get("startedTime"),
|
||||
finished_time=app.get("finishedTime"),
|
||||
)
|
||||
logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}")
|
||||
return result
|
||||
|
||||
|
||||
def list_applications(
|
||||
connection_name: str,
|
||||
state: str | None = None,
|
||||
queue: str | None = None,
|
||||
limit: int = 100,
|
||||
) -> list[ApplicationSummary]:
|
||||
"""List YARN applications on the named cluster, optionally filtered.
|
||||
|
||||
Bypasses the local JobStore (this is for apps not submitted through
|
||||
this service). The YARN ResourceManager REST endpoint
|
||||
/ws/v1/cluster/apps is queried through the connection's auth/SSL
|
||||
config.
|
||||
|
||||
Defaults: limit=100 (YARN has no offset-based pagination, so large
|
||||
clusters should use state/queue filters to scope the result).
|
||||
"""
|
||||
logger.debug(
|
||||
f"list_applications enter connection_name={connection_name} "
|
||||
f"state={state} queue={queue} limit={limit}"
|
||||
)
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
raw_apps = list_applications_yarn(
|
||||
config, state=state, queue=queue, limit=limit
|
||||
)
|
||||
summaries = [
|
||||
ApplicationSummary(
|
||||
application_id=app.get("id", ""),
|
||||
name=app.get("name", ""),
|
||||
user=app.get("user", ""),
|
||||
queue=app.get("queue", ""),
|
||||
state=app.get("state", ""),
|
||||
final_status=app.get("finalStatus"),
|
||||
application_type=app.get("applicationType"),
|
||||
application_tags=app.get("applicationTags", ""),
|
||||
started_time=app.get("startedTime", 0),
|
||||
finished_time=app.get("finishedTime", 0),
|
||||
tracking_url=app.get("trackingUrl"),
|
||||
progress=app.get("progress"),
|
||||
)
|
||||
for app in raw_apps
|
||||
]
|
||||
logger.info(
|
||||
f"list_applications ok connection_name={connection_name} count={len(summaries)}"
|
||||
)
|
||||
return summaries
|
||||
@@ -0,0 +1,136 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/7/9
|
||||
@Author :tao.chen
|
||||
|
||||
Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the
|
||||
cluster's network when it cannot reach those hosts directly. Security:
|
||||
the host is checked against an explicit fnmatch glob allowlist configured
|
||||
on the connection (`url_allowlist`). Empty or missing allowlist means no
|
||||
URL access; the allowlist is the only gate — scheme and IP-literal checks
|
||||
are intentionally NOT performed. Reuses the connection's saved auth/SSL
|
||||
config so the agent doesn't need cluster credentials.
|
||||
"""
|
||||
import fnmatch
|
||||
from urllib.parse import urlparse, urljoin
|
||||
|
||||
import httpx
|
||||
|
||||
from common.logging import logger
|
||||
from spark_executor.core.yarn_client import YarnClientConfig
|
||||
from spark_executor.models import FetchUrlResult
|
||||
from spark_executor.tools.connections import store as conn_store
|
||||
|
||||
_REQUEST_TIMEOUT_SECONDS = 30
|
||||
|
||||
_MAX_REDIRECTS = 3
|
||||
_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
|
||||
|
||||
|
||||
def _host_matches_any_glob(host: str, allowlist: list[str]) -> bool:
|
||||
"""True if `host` matches any of the fnmatch glob patterns.
|
||||
|
||||
fnmatch is case-sensitive on Linux (our deployment target). `*` in a
|
||||
pattern does NOT cross '.' boundaries, so 'ccam*' matches 'ccam50' but
|
||||
NOT 'ccam50.evil.com' (which is what we want — domain boundary matters
|
||||
for SSRF).
|
||||
"""
|
||||
host_labels = host.split(".")
|
||||
for p in allowlist:
|
||||
pat_labels = p.split(".")
|
||||
if len(pat_labels) != len(host_labels):
|
||||
continue
|
||||
if all(fnmatch.fnmatchcase(h, pat) for h, pat in zip(host_labels, pat_labels)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _validate_url_host(url: str, allowlist: list[str] | None) -> None:
|
||||
"""Reject the URL unless its host matches a glob in `allowlist`.
|
||||
|
||||
The only check. The url_allowlist is the single source of truth for
|
||||
what fetch_url is allowed to access — no scheme, IP-literal, or
|
||||
"must be the same as yarn_rm_url" guardrails. The caller is
|
||||
responsible for writing a tight allowlist.
|
||||
|
||||
The only structural check: the URL must have a host (otherwise
|
||||
the glob match has nothing to test). Anything else is the
|
||||
allowlist's job.
|
||||
"""
|
||||
host = urlparse(url).hostname
|
||||
if not host:
|
||||
raise ValueError(f"URL has no host: {url!r}")
|
||||
if _host_matches_any_glob(host, allowlist or []):
|
||||
return
|
||||
raise ValueError(
|
||||
f"URL host {host!r} is not in Connection.url_allowlist {allowlist!r}. "
|
||||
f"Add the host pattern to url_allowlist (or use a broader glob) "
|
||||
f"and try again."
|
||||
)
|
||||
|
||||
|
||||
def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
|
||||
"""Proxy an HTTP GET to url using the auth/SSL settings of connection_name."""
|
||||
logger.debug(f"fetch_url enter url={url} connection_name={connection_name}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
|
||||
_validate_url_host(url, conn.url_allowlist)
|
||||
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
auth = config.auth_for_httpx()
|
||||
verify = config.verify_for_httpx()
|
||||
|
||||
for hop in range(_MAX_REDIRECTS + 1):
|
||||
try:
|
||||
resp = httpx.get(
|
||||
url,
|
||||
auth=auth,
|
||||
verify=verify,
|
||||
timeout=_REQUEST_TIMEOUT_SECONDS,
|
||||
follow_redirects=False,
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
# Network-level failure (no response received). Surface the
|
||||
# exception class + message so the LLM can act on it.
|
||||
# ValueError -> 400 via the existing handler in server.py.
|
||||
# The cause chain (`from exc`) preserves the original
|
||||
# exception for loguru.
|
||||
raise ValueError(
|
||||
f"fetch_url could not reach {url!r}: "
|
||||
f"{type(exc).__name__}: {exc}. "
|
||||
f"Check that the URL is reachable from the MCP service, "
|
||||
f"the host is in Connection.url_allowlist, and the "
|
||||
f"connection's auth/SSL settings are correct."
|
||||
) from exc
|
||||
if resp.status_code not in _REDIRECT_STATUSES:
|
||||
logger.info(
|
||||
f"fetch_url ok url={url} connection_name={connection_name} "
|
||||
f"status_code={resp.status_code} "
|
||||
f"content_type={resp.headers.get('content-type', '')} "
|
||||
f"body_chars={len(resp.text)}"
|
||||
)
|
||||
|
||||
return FetchUrlResult(
|
||||
url=url,
|
||||
status_code=resp.status_code,
|
||||
content_type=resp.headers.get("content-type", ""),
|
||||
body=resp.text,
|
||||
)
|
||||
|
||||
location = resp.headers.get("Location") or resp.headers.get("location")
|
||||
if not location:
|
||||
raise ValueError(
|
||||
f"HTTP GET returned {resp.status_code} with no Location header at {url}"
|
||||
)
|
||||
next_url = urljoin(url, location)
|
||||
logger.debug(
|
||||
f"fetch_url redirect url={url} status={resp.status_code} to={next_url}"
|
||||
)
|
||||
_validate_url_host(next_url, conn.url_allowlist)
|
||||
url = next_url
|
||||
|
||||
raise ValueError(
|
||||
f"fetch_url exceeded {_MAX_REDIRECTS} redirects (last url: {url})"
|
||||
)
|
||||
@@ -16,7 +16,10 @@ Safety:
|
||||
- update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR
|
||||
(settings.jobs_dir) so the agent cannot overwrite host-mounted
|
||||
configs or arbitrary files on the container FS.
|
||||
- 1 MB cap on both read and write payloads to keep MCP responses bounded.
|
||||
- No size cap on read or write payloads — the caller controls the
|
||||
file size. MCP response size is bounded by httpx and the JSON
|
||||
transport; for very large files the LLM should split via
|
||||
write_job_file and read in chunks via read_job_file + update.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -24,9 +27,6 @@ from pathlib import Path
|
||||
from common.config import settings
|
||||
from common.logging import logger
|
||||
|
||||
MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB
|
||||
|
||||
|
||||
class ScriptFileError(ValueError):
|
||||
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
|
||||
handler in server.py.
|
||||
@@ -68,16 +68,10 @@ def _check_writable(script_path: str) -> None:
|
||||
def read_job_file(script_path: str) -> dict[str, object]:
|
||||
"""Return the text content of an existing script file.
|
||||
|
||||
Caps the read at 1 MB to keep MCP responses bounded; raises
|
||||
ScriptFileError (-> 400) if the file is missing or too large.
|
||||
No size cap. Raises ScriptFileError (-> 400) if the file is missing.
|
||||
"""
|
||||
_check_readable(script_path)
|
||||
size = os.path.getsize(script_path)
|
||||
if size > MAX_FILE_BYTES:
|
||||
raise ScriptFileError(
|
||||
f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} "
|
||||
f"byte cap). Edit it via a host volume mount instead."
|
||||
)
|
||||
logger.debug(f"read_job_file enter script_path={script_path} size={size}")
|
||||
with open(script_path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
@@ -88,17 +82,12 @@ def read_job_file(script_path: str) -> dict[str, object]:
|
||||
def update_job_file(script_path: str, content: str) -> dict[str, object]:
|
||||
"""Overwrite an existing script file with new content.
|
||||
|
||||
Restricted to paths under settings.jobs_dir. Caps writes at 1 MB.
|
||||
Raises ScriptFileError (-> 400) if the path is missing, outside
|
||||
the allowed dir, or the content is too large.
|
||||
Restricted to paths under settings.jobs_dir. No size cap.
|
||||
Raises ScriptFileError (-> 400) if the path is missing or outside
|
||||
the allowed dir.
|
||||
"""
|
||||
_check_writable(script_path)
|
||||
encoded_size = len(content.encode("utf-8"))
|
||||
if encoded_size > MAX_FILE_BYTES:
|
||||
raise ScriptFileError(
|
||||
f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} "
|
||||
f"byte cap). Split the script into multiple files."
|
||||
)
|
||||
logger.debug(
|
||||
f"update_job_file enter script_path={script_path} "
|
||||
f"new_bytes={encoded_size}"
|
||||
|
||||
@@ -21,7 +21,13 @@ def kill_job(job_id: str) -> dict[str, str]:
|
||||
logger.debug(f"kill_job enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint=(
|
||||
"the YARN CLI (`yarn application -kill <app_id>`) or the YARN "
|
||||
"UI directly — this service has no external kill tool"
|
||||
),
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -11,14 +11,31 @@ from spark_executor.tools.connections import store as conn_store
|
||||
store = JobStore()
|
||||
|
||||
|
||||
def _unknown_job_error(uid: str) -> KeyError:
|
||||
"""Standard "we tried both IDs and found nothing" message.
|
||||
def _unknown_job_error(uid: str, external_tool_hint: str | None = None) -> Exception:
|
||||
"""Build the right error for a not-found job.
|
||||
|
||||
The agent gets this from confirm_submit_job's response:
|
||||
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
|
||||
and routinely confuses which to pass here. Spelling out that BOTH
|
||||
IDs were tried (and what they look like) saves a round trip.
|
||||
and routinely confuses which to pass here. We differentiate two cases:
|
||||
|
||||
- `uid` looks like a YARN application_id (starts with 'application_')
|
||||
AND the caller passed an `external_tool_hint`: the YARN app likely
|
||||
exists, this tool just can't serve it because it was not submitted
|
||||
through this service. Raise ValueError (-> 400 via the FastAPI
|
||||
handler) pointing the agent at the right external tool.
|
||||
- Otherwise: no local record of either form of id. Raise KeyError
|
||||
(-> 404). Spelling out what both IDs look like saves a round trip.
|
||||
"""
|
||||
if uid.startswith("application_") and external_tool_hint:
|
||||
return ValueError(
|
||||
f"job_id={uid!r} looks like a YARN application_id (starts with "
|
||||
f"'application_'), but this tool only works for jobs submitted "
|
||||
f"through this MCP service (no local JobStore record). For YARN "
|
||||
f"applications not submitted here, use {external_tool_hint} "
|
||||
f"instead. (If you actually submitted this job through this "
|
||||
f"service, pass the local job_id — it is a 12-char hex like "
|
||||
f"'a1b2c3d4e5f6'.)"
|
||||
)
|
||||
return KeyError(
|
||||
f"No Job found for id={uid!r} (neither as job_id nor as "
|
||||
f"application_id). Pass the job_id from confirm_submit_job's "
|
||||
@@ -37,7 +54,10 @@ def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
|
||||
logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_logs(application_id, connection_name, tail_chars)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -17,22 +17,143 @@ class EmptyRequest(BaseModel):
|
||||
|
||||
|
||||
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
|
||||
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."
|
||||
),
|
||||
)
|
||||
|
||||
url_allowlist: list[str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional list of fnmatch glob patterns for hosts the fetch_url tool "
|
||||
"may access. See Connection.url_allowlist for full semantics. "
|
||||
"Example for single-label host clusters: ['ccam*'] allows any host "
|
||||
"starting with 'ccam' (ccam1, ccam2, ..., ccam99). If omitted, None, "
|
||||
"or empty, the saved connection will have url_allowlist=[] (the "
|
||||
"default), meaning fetch_url will reject every URL until the list is "
|
||||
"populated via update_connection."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional Spark History Server (SHS) base URL, e.g. "
|
||||
"'http://history.prod.internal:18080'. Stored for reference. "
|
||||
"Currently not consumed by any tool — see Connection.history_server_url."
|
||||
),
|
||||
)
|
||||
|
||||
class PrepareSubmitJobRequest(BaseModel):
|
||||
connection: str
|
||||
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.",
|
||||
@@ -72,34 +193,115 @@ class PrepareSubmitJobRequest(BaseModel):
|
||||
|
||||
|
||||
class PendingIdRequest(BaseModel):
|
||||
pending_id: str
|
||||
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
|
||||
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 SQL guard.",
|
||||
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."
|
||||
),
|
||||
)
|
||||
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
|
||||
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
|
||||
tail_chars: int = 5000
|
||||
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
|
||||
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):
|
||||
@@ -142,3 +344,191 @@ class UpdateJobFileRequest(BaseModel):
|
||||
"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.",
|
||||
)
|
||||
|
||||
class FetchUrlRequest(BaseModel):
|
||||
url: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Absolute URL to fetch. The only access control is the named "
|
||||
"Connection's url_allowlist: the URL host must match one of the "
|
||||
"fnmatch glob patterns in that list. An empty or omitted allowlist "
|
||||
"denies every host. IP literals and non-HTTP schemes are allowed "
|
||||
"if and only if they are matched by the allowlist. The Connection's "
|
||||
"saved auth is reused for the outbound request — the agent does not "
|
||||
"need cluster credentials."
|
||||
),
|
||||
)
|
||||
connection_name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection (see list_connections). The "
|
||||
"Connection's url_allowlist is the host gate for this tool — "
|
||||
"**not yarn_rm_url** (yarn_rm_url is only used by the "
|
||||
"get_external_* tools to locate the YARN RM). The "
|
||||
"Connection's auth_type / auth_user / auth_password / "
|
||||
"ssl_verify / ssl_ca_bundle are reused for the request."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UpdateConnectionRequest(BaseModel):
|
||||
name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of the existing Connection to update (immutable identifier). "
|
||||
"If you want to rename, use delete_connection + save_connection."
|
||||
),
|
||||
)
|
||||
master: str | None = Field(
|
||||
default=None,
|
||||
description="New Spark master URL. Omit to keep current.",
|
||||
)
|
||||
deploy_mode: str | None = Field(
|
||||
default=None,
|
||||
description="New Spark deploy mode ('cluster' or 'client'). Omit to keep current.",
|
||||
)
|
||||
yarn_rm_url: str | None = Field(
|
||||
default=None,
|
||||
description="New YARN ResourceManager REST base URL. Omit to keep current.",
|
||||
)
|
||||
spark_conf: dict[str, str] | None = Field(
|
||||
default=None,
|
||||
description="Replacement Spark conf dict (not merged with existing). Omit to keep current.",
|
||||
)
|
||||
ssl_verify: bool | None = Field(
|
||||
default=None,
|
||||
description="New SSL verify setting. Omit to keep current.",
|
||||
)
|
||||
ssl_ca_bundle: str | None = Field(
|
||||
default=None,
|
||||
description="New SSL CA bundle path. Omit to keep current.",
|
||||
)
|
||||
auth_type: str | None = Field(
|
||||
default=None,
|
||||
description="New auth mode. Omit to keep current.",
|
||||
)
|
||||
auth_user: str | None = Field(
|
||||
default=None,
|
||||
description="New auth username. Omit to keep current.",
|
||||
)
|
||||
auth_password: str | None = Field(
|
||||
default=None,
|
||||
description="New auth password. Omit to keep current.",
|
||||
)
|
||||
auth_principal: str | None = Field(
|
||||
default=None,
|
||||
description="New Kerberos principal. Omit to keep current.",
|
||||
)
|
||||
auth_keytab: str | None = Field(
|
||||
default=None,
|
||||
description="New Kerberos keytab path. Omit to keep current.",
|
||||
)
|
||||
url_allowlist: list[str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Replacement url_allowlist list (not merged). Omit to keep current. "
|
||||
"Pass an empty list to deny all hosts (the default for new connections). "
|
||||
"Example: ['ccam*'] allows ccam1-ccam99."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"New Spark History Server base URL. Omit to keep current. "
|
||||
"Pass None to clear (use delete_connection + save_connection "
|
||||
"if you need explicit clear semantics; same caveat as other "
|
||||
"Optional fields in this tool)."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ListApplicationsRequest(BaseModel):
|
||||
connection_name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection (see list_connections) pointing at "
|
||||
"the YARN cluster to query."
|
||||
),
|
||||
)
|
||||
state: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional YARN application state filter. One of: 'NEW', "
|
||||
"'NEW_SAVING', 'SUBMITTED', 'ACCEPTED', 'RUNNING', 'FINISHED', "
|
||||
"'FAILED', 'KILLED'. 'FINISHED' is the umbrella state covering "
|
||||
"SUCCEEDED/FAILED/KILLED. None = no state filter (returns all "
|
||||
"states up to `limit`)."
|
||||
),
|
||||
)
|
||||
queue: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional YARN queue name filter (e.g. 'default', 'prod'). "
|
||||
"None = no queue filter."
|
||||
),
|
||||
)
|
||||
limit: int = Field(
|
||||
default=100,
|
||||
ge=1,
|
||||
le=10000,
|
||||
description=(
|
||||
"Maximum number of applications to return. YARN has no "
|
||||
"offset-based pagination, so for large clusters use state/queue "
|
||||
"filters to scope the result. Max 10000 in practice (YARN's own "
|
||||
"limit on the limit param)."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -24,7 +24,10 @@ def get_job_result(job_id: str) -> JobResult:
|
||||
logger.debug(f"get_job_result enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_result(application_id, connection_name)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -23,7 +23,10 @@ def get_job_status(job_id: str) -> JobStatus:
|
||||
logger.debug(f"get_job_status enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_status(application_id, connection_name)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -6,9 +6,8 @@
|
||||
import os
|
||||
import secrets
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from common.config import settings
|
||||
from common.logging import logger
|
||||
from common.sql_guard import validate_pyspark_code
|
||||
from spark_executor.core.connection_store import store as conn_store
|
||||
@@ -124,7 +123,7 @@ def prepare_submit_job(
|
||||
num_executors=num_executors,
|
||||
spark_conf=dict(conn.spark_conf),
|
||||
extra_args=dict(extra_args or {}),
|
||||
created_at=datetime.utcnow(),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
status="PENDING",
|
||||
)
|
||||
pending_store.save(pending)
|
||||
@@ -248,7 +247,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
||||
application_id=application_id,
|
||||
script_path=pending.script_path,
|
||||
queue=pending.queue,
|
||||
submit_time=datetime.utcnow(),
|
||||
submit_time=datetime.now(timezone.utc),
|
||||
connection=pending.connection,
|
||||
yarn_rm_url=pending.yarn_rm_url,
|
||||
)
|
||||
@@ -282,7 +281,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
||||
application_id=application_id,
|
||||
script_path=pending.script_path,
|
||||
queue=pending.queue,
|
||||
submit_time=datetime.utcnow(),
|
||||
submit_time=datetime.now(timezone.utc),
|
||||
connection=pending.connection,
|
||||
yarn_rm_url=pending.yarn_rm_url,
|
||||
)
|
||||
|
||||
@@ -8,6 +8,7 @@ from fastapi.testclient import TestClient
|
||||
from spark_executor.core import connection_store, pending_store
|
||||
from spark_executor.core.connection_store import ConnectionStore
|
||||
from spark_executor.core.pending_store import PendingStore
|
||||
from spark_executor.models import Connection
|
||||
from spark_executor.server import app
|
||||
from spark_executor.tools import connections, submit
|
||||
|
||||
@@ -64,6 +65,19 @@ def test_seventeen_tool_routes_registered():
|
||||
assert "/update_pending_job" in paths
|
||||
|
||||
|
||||
def test_twenty_three_tool_routes_registered():
|
||||
paths = {r.path for r in app.routes}
|
||||
for path in (
|
||||
"/get_external_job_logs",
|
||||
"/get_external_job_status",
|
||||
"/get_external_job_result",
|
||||
"/list_applications",
|
||||
"/fetch_url",
|
||||
"/update_connection",
|
||||
):
|
||||
assert path in paths, f"missing MCP tool route: {path}"
|
||||
|
||||
|
||||
# --- operation_id: pin clean MCP tool names (no auto-generated suffixes) ---
|
||||
#
|
||||
# fastapi-mcp uses each route's OpenAPI `operationId` as the MCP tool name
|
||||
@@ -531,3 +545,147 @@ def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch):
|
||||
)
|
||||
assert r.status_code == 400
|
||||
assert "only PENDING submissions can be updated" in r.json()["detail"]
|
||||
|
||||
|
||||
# --- YarnError -> 502 handler ---
|
||||
|
||||
|
||||
def test_yarn_error_returns_502_with_detail(tmp_path, monkeypatch):
|
||||
"""When a YARN REST call raises YarnError (host unreachable, YARN
|
||||
returned 4xx/5xx, parse failure), the route must surface HTTP 502
|
||||
with the YarnError message in the detail — not 500 'Internal
|
||||
Server Error' with no info. Mirrors the same fix for fetch_url.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.status.get_application_status",
|
||||
side_effect=YarnError("YARN connection failed: ConnectError: Connection refused"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_job_status",
|
||||
json={"job_id": "a1b2c3d4e5f6"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
detail = r.json()["detail"]
|
||||
assert "YARN connection failed" in detail
|
||||
assert "Connection refused" in detail
|
||||
|
||||
|
||||
def test_external_job_status_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_status uses get_application_status under the
|
||||
hood — same handler, same 502, same detail."""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
side_effect=YarnError("YARN application 'application_xxx' not found"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_status",
|
||||
json={"application_id": "application_1740000000001_0001", "connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "not found" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_external_job_result_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_result also uses get_application_status (same
|
||||
underlying YARN endpoint), and exercises the same YarnError path."""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
side_effect=YarnError("YARN connection failed: ConnectError: Connection refused"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_result",
|
||||
json={"application_id": "application_1740000000001_0001", "connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN connection failed" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_external_job_logs_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_logs uses get_application_logs — different
|
||||
function, but raises YarnError the same way. Handler must catch
|
||||
it the same way."""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_logs",
|
||||
side_effect=YarnError("YARN GET logs returned HTTP 500: cluster overloaded"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_logs",
|
||||
json={
|
||||
"application_id": "application_1740000000001_0001",
|
||||
"connection_name": "prod",
|
||||
"tail_chars": 5000,
|
||||
},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN GET logs returned HTTP 500" in r.json()["detail"]
|
||||
assert "cluster overloaded" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_list_applications_returns_502_on_yarn_error(tmp_path):
|
||||
"""list_applications uses list_applications_yarn — yet another
|
||||
YARN-touching tool. Same YarnError -> 502 contract."""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
side_effect=YarnError("YARN list applications failed: 503 Service Unavailable"),
|
||||
):
|
||||
r = c.post(
|
||||
"/list_applications",
|
||||
json={"connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN list applications failed" in r.json()["detail"]
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# coding=utf-8
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@@ -120,3 +121,113 @@ def test_delete_connection_returns_status(_fresh_store):
|
||||
def test_delete_connection_unknown_raises(_fresh_store):
|
||||
with pytest.raises(KeyError):
|
||||
connections.delete_connection("missing")
|
||||
|
||||
|
||||
def test_save_connection_preserves_url_allowlist_on_existing_record(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn", url_allowlist=["ccam*"])
|
||||
out = connections.save_connection(name="prod", master="spark://new:7077")
|
||||
assert out["master"] == "spark://new:7077"
|
||||
assert out["url_allowlist"] == ["ccam*"]
|
||||
|
||||
|
||||
def test_update_connection_changes_specified_field(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
out = connections.update_connection(name="prod", master="spark://new:7077")
|
||||
assert out["master"] == "spark://new:7077"
|
||||
assert _fresh_store.get("prod").master == "spark://new:7077"
|
||||
|
||||
|
||||
def test_update_connection_keeps_omitted_fields(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
deploy_mode="client",
|
||||
yarn_rm_url="http://rm:8088",
|
||||
)
|
||||
out = connections.update_connection(name="prod", deploy_mode="cluster")
|
||||
assert out["deploy_mode"] == "cluster"
|
||||
assert out["master"] == "yarn"
|
||||
assert out["yarn_rm_url"] == "http://rm:8088"
|
||||
|
||||
|
||||
def test_update_connection_with_no_fields_is_noop(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn", deploy_mode="client")
|
||||
out = connections.update_connection(name="prod")
|
||||
assert out["master"] == "yarn"
|
||||
assert out["deploy_mode"] == "client"
|
||||
|
||||
|
||||
def test_update_connection_changes_url_allowlist(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
out = connections.update_connection(name="prod", url_allowlist=["ccam*"])
|
||||
assert out["url_allowlist"] == ["ccam*"]
|
||||
|
||||
|
||||
def test_update_connection_replaces_not_merges_dict(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn", spark_conf={"a": "1"})
|
||||
out = connections.update_connection(name="prod", spark_conf={"b": "2"})
|
||||
assert out["spark_conf"] == {"b": "2"}
|
||||
|
||||
|
||||
def test_update_connection_replaces_not_merges_url_allowlist(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
out = connections.update_connection(name="prod", url_allowlist=["nm*"])
|
||||
assert out["url_allowlist"] == ["nm*"]
|
||||
|
||||
|
||||
def test_update_connection_raises_for_unknown_name(_fresh_store):
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
connections.update_connection(name="missing", master="yarn")
|
||||
|
||||
|
||||
def test_update_connection_validates_patched_connection(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
with pytest.raises(ValueError, match="master must be"):
|
||||
connections.update_connection(name="prod", master="http://bad")
|
||||
|
||||
|
||||
def test_update_connection_persists_to_disk(_fresh_store, tmp_path):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
connections.update_connection(name="prod", master="spark://new:7077")
|
||||
raw = json.loads((tmp_path / "connections.json").read_text())
|
||||
assert raw["prod"]["master"] == "spark://new:7077"
|
||||
|
||||
|
||||
# --- history_server_url ---
|
||||
|
||||
def test_save_connection_with_history_server_url(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
history_server_url="http://history.prod.internal:18080",
|
||||
)
|
||||
assert _fresh_store.get("prod").history_server_url == "http://history.prod.internal:18080"
|
||||
|
||||
|
||||
def test_save_connection_history_server_url_defaults_to_none(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
assert _fresh_store.get("prod").history_server_url is None
|
||||
|
||||
|
||||
def test_update_connection_changes_history_server_url(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
out = connections.update_connection(
|
||||
name="prod",
|
||||
history_server_url="http://history2.prod.internal:18080",
|
||||
)
|
||||
assert out["history_server_url"] == "http://history2.prod.internal:18080"
|
||||
assert _fresh_store.get("prod").history_server_url == "http://history2.prod.internal:18080"
|
||||
|
||||
|
||||
def test_update_connection_keeps_history_server_url_when_omitted(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
history_server_url="http://history.prod.internal:18080",
|
||||
)
|
||||
out = connections.update_connection(name="prod", master="spark://new:7077")
|
||||
assert out["history_server_url"] == "http://history.prod.internal:18080"
|
||||
|
||||
@@ -0,0 +1,305 @@
|
||||
# coding=utf-8
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from spark_executor.core import connection_store
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
from spark_executor.models import ApplicationSummary, Connection
|
||||
from spark_executor.tools import connections, external_jobs
|
||||
from spark_executor.tools.requests import ListApplicationsRequest
|
||||
|
||||
|
||||
def _fresh_stores():
|
||||
"""Reset connection store singletons for a single test."""
|
||||
store = connection_store.ConnectionStore()
|
||||
connection_store.store = store
|
||||
connections.store = store
|
||||
external_jobs.conn_store = store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_stores(tmp_path, monkeypatch):
|
||||
"""Reset connection store singletons to an isolated tmp_path."""
|
||||
monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
||||
_fresh_stores()
|
||||
|
||||
|
||||
def test_get_external_job_logs_returns_tailed():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
long_log = "LOG" * 3000
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_logs",
|
||||
return_value=long_log,
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="prod", tail_chars=100
|
||||
)
|
||||
assert out == long_log[-100:]
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_logs_returns_full_when_short():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
short_log = "short log"
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_logs",
|
||||
return_value=short_log,
|
||||
):
|
||||
out = external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="prod", tail_chars=5000
|
||||
)
|
||||
assert out == short_log
|
||||
|
||||
|
||||
def test_get_external_job_logs_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
|
||||
|
||||
def test_get_external_job_status_returns_state():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
raw = json.dumps({"app": {"state": "RUNNING"}})
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
return_value=("RUNNING", raw),
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_status(
|
||||
application_id="application_1", connection_name="prod"
|
||||
)
|
||||
assert out.application_id == "application_1"
|
||||
assert out.state == "RUNNING"
|
||||
assert out.raw == raw
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_status_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_status(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
|
||||
|
||||
def test_get_external_job_result_parses_app_fields():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
raw = json.dumps(
|
||||
{
|
||||
"app": {
|
||||
"state": "FINISHED",
|
||||
"finalStatus": "SUCCEEDED",
|
||||
"diagnostics": "",
|
||||
"trackingUrl": "http://rm:8088/proxy/application_1",
|
||||
"startedTime": 100,
|
||||
"finishedTime": 200,
|
||||
}
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
return_value=("FINISHED", raw),
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_result(
|
||||
application_id="application_1", connection_name="prod"
|
||||
)
|
||||
assert out.application_id == "application_1"
|
||||
assert out.state == "FINISHED"
|
||||
assert out.final_status == "SUCCEEDED"
|
||||
assert out.diagnostics == ""
|
||||
assert out.tracking_url == "http://rm:8088/proxy/application_1"
|
||||
assert out.started_time == 100
|
||||
assert out.finished_time == 200
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_result_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_result(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
|
||||
|
||||
def test_list_applications_returns_summaries(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[
|
||||
{
|
||||
"id": "application_1",
|
||||
"name": "app-one",
|
||||
"user": "alice",
|
||||
"queue": "default",
|
||||
"state": "RUNNING",
|
||||
"finalStatus": "UNDEFINED",
|
||||
"applicationType": "SPARK",
|
||||
"applicationTags": "tag1",
|
||||
"startedTime": 1000,
|
||||
"finishedTime": 0,
|
||||
"trackingUrl": "http://rm:8088/proxy/application_1",
|
||||
"progress": 75.0,
|
||||
},
|
||||
{
|
||||
"id": "application_2",
|
||||
"name": "app-two",
|
||||
"user": "bob",
|
||||
"queue": "research",
|
||||
"state": "FINISHED",
|
||||
"finalStatus": "SUCCEEDED",
|
||||
"applicationType": "SPARK",
|
||||
"applicationTags": "",
|
||||
"startedTime": 2000,
|
||||
"finishedTime": 3000,
|
||||
"trackingUrl": "http://rm:8088/proxy/application_2",
|
||||
"progress": 100.0,
|
||||
},
|
||||
],
|
||||
) as m:
|
||||
out = external_jobs.list_applications("prod")
|
||||
assert len(out) == 2
|
||||
assert all(isinstance(item, ApplicationSummary) for item in out)
|
||||
assert out[0].application_id == "application_1"
|
||||
assert out[1].application_id == "application_2"
|
||||
args = m.call_args.args
|
||||
assert args[0].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_list_applications_passes_state_filter(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[],
|
||||
) as m:
|
||||
external_jobs.list_applications("prod", state="RUNNING")
|
||||
assert m.call_args.kwargs == {"state": "RUNNING", "queue": None, "limit": 100}
|
||||
|
||||
|
||||
def test_list_applications_passes_queue_filter(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[],
|
||||
) as m:
|
||||
external_jobs.list_applications("prod", queue="research")
|
||||
assert m.call_args.kwargs == {"state": None, "queue": "research", "limit": 100}
|
||||
|
||||
|
||||
def test_list_applications_passes_limit_filter(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[],
|
||||
) as m:
|
||||
external_jobs.list_applications("prod", limit=50)
|
||||
assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 50}
|
||||
|
||||
|
||||
def test_list_applications_default_limit_is_100(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[],
|
||||
) as m:
|
||||
external_jobs.list_applications("prod")
|
||||
assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 100}
|
||||
assert ListApplicationsRequest(connection_name="prod").limit == 100
|
||||
|
||||
|
||||
def test_list_applications_returns_empty_list_when_no_apps(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[],
|
||||
):
|
||||
out = external_jobs.list_applications("prod")
|
||||
assert out == []
|
||||
|
||||
|
||||
def test_list_applications_raises_for_missing_connection(fresh_stores):
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.list_applications("missing")
|
||||
|
||||
|
||||
def test_list_applications_maps_yarn_json_to_summary(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
yarn_app = {
|
||||
"id": "application_42",
|
||||
"name": "mapped-app",
|
||||
"user": "carol",
|
||||
"queue": "prod",
|
||||
"state": "ACCEPTED",
|
||||
"finalStatus": "UNDEFINED",
|
||||
"applicationType": "MAPREDUCE",
|
||||
"applicationTags": "batch",
|
||||
"startedTime": 12345,
|
||||
"finishedTime": 0,
|
||||
"trackingUrl": "http://rm:8088/proxy/application_42",
|
||||
"progress": 12.5,
|
||||
}
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
return_value=[yarn_app],
|
||||
):
|
||||
out = external_jobs.list_applications("prod")
|
||||
assert len(out) == 1
|
||||
summary = out[0]
|
||||
assert summary.application_id == "application_42"
|
||||
assert summary.name == "mapped-app"
|
||||
assert summary.user == "carol"
|
||||
assert summary.queue == "prod"
|
||||
assert summary.state == "ACCEPTED"
|
||||
assert summary.final_status == "UNDEFINED"
|
||||
assert summary.application_type == "MAPREDUCE"
|
||||
assert summary.application_tags == "batch"
|
||||
assert summary.started_time == 12345
|
||||
assert summary.finished_time == 0
|
||||
assert summary.tracking_url == "http://rm:8088/proxy/application_42"
|
||||
assert summary.progress == 12.5
|
||||
|
||||
|
||||
def test_list_applications_raises_on_404(fresh_stores):
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.list_applications_yarn",
|
||||
side_effect=YarnError("YARN list applications failed: 404"),
|
||||
):
|
||||
with pytest.raises(YarnError, match="YARN list applications failed"):
|
||||
external_jobs.list_applications("prod")
|
||||
@@ -0,0 +1,402 @@
|
||||
# coding=utf-8
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from spark_executor.core import connection_store
|
||||
from spark_executor.core.connection_store import ConnectionStore
|
||||
from spark_executor.models import Connection
|
||||
from spark_executor.tools import connections, fetch_url
|
||||
from spark_executor.tools.requests import ListApplicationsRequest
|
||||
|
||||
|
||||
def _fresh_stores(tmp_path, monkeypatch):
|
||||
"""Reset connection store singletons for a single test."""
|
||||
monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
||||
store = ConnectionStore()
|
||||
monkeypatch.setattr(connection_store, "store", store)
|
||||
connections.store = store
|
||||
fetch_url.conn_store = store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_stores(tmp_path: Path, monkeypatch):
|
||||
"""Reset connection store singletons for a single test."""
|
||||
_fresh_stores(tmp_path, monkeypatch)
|
||||
yield
|
||||
|
||||
|
||||
def test_fetch_url_returns_body_and_status(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"})
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod")
|
||||
assert out.url == "http://nm01.prod.internal:8042/node"
|
||||
assert out.status_code == 200
|
||||
assert out.content_type == "text/html"
|
||||
assert out.body == "hello"
|
||||
assert "truncated" not in fetch_url.FetchUrlResult.model_fields
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_raises_for_missing_connection(fresh_stores):
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
fetch_url.fetch_url("http://rm.prod.internal:8088/", "missing")
|
||||
|
||||
|
||||
def test_fetch_url_passes_auth_from_connection(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="auth",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
auth_type="basic",
|
||||
auth_user="u",
|
||||
auth_password="p",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="ok")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth")
|
||||
auth = m.call_args.kwargs["auth"]
|
||||
assert isinstance(auth, httpx.BasicAuth)
|
||||
import base64
|
||||
|
||||
creds = base64.b64decode(auth._auth_header.split()[1]).decode()
|
||||
assert creds == "u:p"
|
||||
|
||||
|
||||
def test_fetch_url_passes_ssl_verify_from_connection(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="insecure",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
ssl_verify=False,
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="ok")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure")
|
||||
assert m.call_args.kwargs["verify"] is False
|
||||
|
||||
|
||||
def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="ccam",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam")
|
||||
assert out.status_code == 200
|
||||
assert out.body == "hello"
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["ccam*", "*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url(
|
||||
"http://history.prod.internal:18080/api/v1/info", "prod"
|
||||
)
|
||||
assert out.status_code == 200
|
||||
assert out.body == "hello"
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_glob_does_not_match_unrelated_host(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="ccam",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://evil.com/foo", "ccam")
|
||||
|
||||
|
||||
def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="ccam",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://ccam50.evil.com/", "ccam")
|
||||
|
||||
|
||||
def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("http://ccam50:8088/", "prod")
|
||||
assert out.status_code == 200
|
||||
assert out.body == "hello"
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_accepts_ip_literal_when_in_url_allowlist(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
url_allowlist=["*.*.*.*"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod")
|
||||
assert out.status_code == 200
|
||||
assert out.body == "hello"
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_accepts_https_when_in_url_allowlist(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
url_allowlist=["ccam*.example.com"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(200, text="hello")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get", return_value=resp
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod")
|
||||
assert out.status_code == 200
|
||||
assert out.body == "hello"
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_rejects_when_url_has_no_host(fresh_stores):
|
||||
with pytest.raises(ValueError, match="URL has no host"):
|
||||
fetch_url._validate_url_host("", ["*"])
|
||||
with pytest.raises(ValueError, match="URL has no host"):
|
||||
fetch_url._validate_url_host("not-a-url", ["*"])
|
||||
|
||||
|
||||
def test_fetch_url_rejects_when_url_allowlist_empty(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(name="prod", master="yarn", url_allowlist=[])
|
||||
)
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://anything.com/", "prod")
|
||||
|
||||
|
||||
def test_fetch_url_error_message_mentions_url_allowlist(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
)
|
||||
)
|
||||
with pytest.raises(ValueError, match="url_allowlist"):
|
||||
fetch_url.fetch_url("http://evil.com/foo", "prod")
|
||||
|
||||
|
||||
def test_fetch_url_omitted_url_allowlist_defaults_to_empty_and_rejects(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal",
|
||||
)
|
||||
)
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://nm.prod.internal/", "prod")
|
||||
|
||||
|
||||
def test_fetch_url_rejects_redirect_to_disallowed_host(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="ccam",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
redirect = httpx.Response(302, headers={"Location": "http://evil.com/"})
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
return_value=redirect,
|
||||
) as m:
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://ccam50/foo", "ccam")
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_fetch_url_follows_redirect_to_allowed_host(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.internal"],
|
||||
)
|
||||
)
|
||||
redirect = httpx.Response(
|
||||
302, headers={"Location": "http://other.internal/"}
|
||||
)
|
||||
final = httpx.Response(200, text="ok")
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
side_effect=[redirect, final],
|
||||
) as m:
|
||||
out = fetch_url.fetch_url("http://foo.internal/", "prod")
|
||||
assert out.status_code == 200
|
||||
assert out.body == "ok"
|
||||
assert m.call_count == 2
|
||||
assert m.call_args_list[1].args[0] == "http://other.internal/"
|
||||
|
||||
|
||||
def test_fetch_url_rejects_redirect_to_ip_literal_not_in_allowlist(fresh_stores):
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="ccam",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
url_allowlist=["ccam*"],
|
||||
)
|
||||
)
|
||||
redirect = httpx.Response(302, headers={"Location": "http://10.0.0.1/"})
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
return_value=redirect,
|
||||
) as m:
|
||||
with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
|
||||
fetch_url.fetch_url("http://ccam50/foo", "ccam")
|
||||
assert m.call_count == 1
|
||||
|
||||
|
||||
def test_list_applications_request_limit_bounds():
|
||||
assert ListApplicationsRequest(connection_name="prod", limit=10000).limit == 10000
|
||||
with pytest.raises(ValidationError):
|
||||
ListApplicationsRequest(connection_name="prod", limit=0)
|
||||
with pytest.raises(ValidationError):
|
||||
ListApplicationsRequest(connection_name="prod", limit=10001)
|
||||
|
||||
|
||||
# --- Network error handling (httpx.HTTPError -> 400 with detail) ---
|
||||
|
||||
|
||||
def test_fetch_url_raises_400_with_detail_on_connect_error(fresh_stores):
|
||||
"""When httpx.get raises ConnectError (host unreachable / port closed),
|
||||
fetch_url must raise ValueError (-> 400) with the exception class
|
||||
and message in the detail. Previously this bubbled up as a bare
|
||||
500 Internal Server Error with no info."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
side_effect=httpx.ConnectError("Connection refused"),
|
||||
):
|
||||
with pytest.raises(ValueError) as ei:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
msg = str(ei.value)
|
||||
assert "ConnectError" in msg
|
||||
assert "Connection refused" in msg
|
||||
assert "nm01.prod.internal" in msg
|
||||
# Cause chain preserved for loguru
|
||||
assert isinstance(ei.value.__cause__, httpx.ConnectError)
|
||||
|
||||
|
||||
def test_fetch_url_raises_400_with_detail_on_timeout(fresh_stores):
|
||||
"""When httpx.get raises TimeoutException (request exceeded 30s),
|
||||
fetch_url must raise ValueError with the exception details surfaced."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
side_effect=httpx.ReadTimeout("Timed out reading"),
|
||||
):
|
||||
with pytest.raises(ValueError) as ei:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
msg = str(ei.value)
|
||||
assert "ReadTimeout" in msg
|
||||
assert "Timed out reading" in msg
|
||||
assert isinstance(ei.value.__cause__, httpx.ReadTimeout)
|
||||
|
||||
|
||||
def test_fetch_url_returns_body_for_4xx_5xx_upstream(fresh_stores):
|
||||
"""Upstream HTTP errors (4xx/5xx responses that DID come back) are
|
||||
NOT translated to ValueError — the FetchUrlResult carries the status
|
||||
code and body so the LLM can see what the server actually said. This
|
||||
is the intentional contrast with the no-response-at-all case."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(503, text="Service Unavailable - try again later")
|
||||
with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp):
|
||||
out = fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
assert out.status_code == 503
|
||||
assert "Service Unavailable" in out.body
|
||||
@@ -118,25 +118,17 @@ def test_update_rejects_missing_file(tmp_path: Path):
|
||||
update_job_file(str(tmp_path / "nope.py"), "x\n")
|
||||
|
||||
|
||||
def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch):
|
||||
"""Cap at 1 MB so the MCP response stays bounded."""
|
||||
def test_update_accepts_content_larger_than_former_1mb_cap(tmp_path: Path):
|
||||
"""The 1 MB cap was removed (commit a5e6d1f and following). Content
|
||||
larger than the old cap must now be accepted; size is bounded by
|
||||
the MCP transport, not the tool."""
|
||||
config.settings.jobs_dir = str(tmp_path)
|
||||
p = tmp_path / "big.py"
|
||||
p.write_text("# small\n")
|
||||
# Synthesize 2 MB of content (don't actually write 2 MB to disk)
|
||||
big = "x" * (2 * 1024 * 1024)
|
||||
with pytest.raises(ScriptFileError, match="too large"):
|
||||
update_job_file(str(p), big)
|
||||
|
||||
|
||||
def test_update_rejects_1mb_plus_1_byte(tmp_path: Path):
|
||||
"""Exactly at the boundary: 1 MB + 1 byte must be rejected."""
|
||||
config.settings.jobs_dir = str(tmp_path)
|
||||
p = tmp_path / "x.py"
|
||||
p.write_text("# t\n")
|
||||
just_over = "x" * (1024 * 1024 + 1)
|
||||
with pytest.raises(ScriptFileError, match="too large"):
|
||||
update_job_file(str(p), just_over)
|
||||
out = update_job_file(str(p), big)
|
||||
assert out["bytes_written"] == len(big)
|
||||
assert p.read_text() == big
|
||||
|
||||
|
||||
# --- round-trip: write -> read -> update -> read ---
|
||||
|
||||
@@ -63,6 +63,14 @@ def test_kill_job_raises_for_unknown_job():
|
||||
kill.kill_job("missing")
|
||||
|
||||
|
||||
def test_kill_job_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) — kill_job has no external
|
||||
equivalent, so the error points to the YARN CLI / UI."""
|
||||
with pytest.raises(ValueError, match="YARN CLI"):
|
||||
kill.kill_job("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_kill_job_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
kill.store.put(
|
||||
|
||||
@@ -68,6 +68,14 @@ def test_get_job_logs_raises_for_unknown_job():
|
||||
logs.get_job_logs("missing")
|
||||
|
||||
|
||||
def test_get_job_logs_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_logs"):
|
||||
logs.get_job_logs("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_get_job_logs_raises_when_connection_missing(fresh_stores):
|
||||
logs.store.put(
|
||||
Job(
|
||||
|
||||
@@ -136,6 +136,14 @@ def test_result_raises_keyerror_for_unknown_job():
|
||||
assert "application_id" in msg
|
||||
|
||||
|
||||
def test_result_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_result"):
|
||||
result.get_job_result("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_result_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
result.store.put(
|
||||
|
||||
@@ -66,6 +66,14 @@ def test_get_job_status_raises_for_unknown_job():
|
||||
status.get_job_status("missing")
|
||||
|
||||
|
||||
def test_get_job_status_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_status"):
|
||||
status.get_job_status("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_get_job_status_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
status.store.put(
|
||||
|
||||
Reference in New Issue
Block a user