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/
|
.pytest_cache/
|
||||||
.coverage
|
.coverage
|
||||||
htmlcov/
|
htmlcov/
|
||||||
|
|
||||||
|
# macOS
|
||||||
|
.DS_Store
|
||||||
|
|||||||
@@ -40,7 +40,7 @@
|
|||||||
| MCP 暴露 | fastapi-mcp |
|
| MCP 暴露 | fastapi-mcp |
|
||||||
| 数据模型 | Pydantic v2 |
|
| 数据模型 | Pydantic v2 |
|
||||||
| HTTP 客户端 | httpx、httpx-kerberos |
|
| HTTP 客户端 | httpx、httpx-kerberos |
|
||||||
| Spark 集成 | `spark-submit`、YARN REST API、`yarn logs` |
|
| Spark 集成 | `spark-submit`、YARN REST API (httpx 直连) |
|
||||||
| 日志 | loguru |
|
| 日志 | loguru |
|
||||||
| 进程管理 | uvicorn、gunicorn |
|
| 进程管理 | uvicorn、gunicorn |
|
||||||
| 包管理 | uv |
|
| 包管理 | uv |
|
||||||
@@ -104,6 +104,7 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
|
|||||||
| `list_connections` | 列出所有连接配置 |
|
| `list_connections` | 列出所有连接配置 |
|
||||||
| `get_connection` | 按名称读取连接配置 |
|
| `get_connection` | 按名称读取连接配置 |
|
||||||
| `delete_connection` | 删除连接配置 |
|
| `delete_connection` | 删除连接配置 |
|
||||||
|
| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) |
|
||||||
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
|
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
|
||||||
| `read_job_file` | 读取已存在的 PySpark 脚本内容 |
|
| `read_job_file` | 读取已存在的 PySpark 脚本内容 |
|
||||||
| `update_job_file` | 覆盖更新已存在的 PySpark 脚本 |
|
| `update_job_file` | 覆盖更新已存在的 PySpark 脚本 |
|
||||||
@@ -117,6 +118,11 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
|
|||||||
| `get_job_result` | 查询终态结果视图 |
|
| `get_job_result` | 查询终态结果视图 |
|
||||||
| `get_job_logs` | 拉取 YARN 聚合日志 |
|
| `get_job_logs` | 拉取 YARN 聚合日志 |
|
||||||
| `kill_job` | Kill YARN application |
|
| `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 工具
|
### Files MCP 工具
|
||||||
|
|
||||||
@@ -215,6 +221,30 @@ get_job_logs
|
|||||||
kill_job
|
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)
|
self._dump(records)
|
||||||
logger.info(f"connection saved name={conn.name} master={conn.master}")
|
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:
|
def delete(self, name: str) -> bool:
|
||||||
with self._lock:
|
with self._lock:
|
||||||
records = self._load()
|
records = self._load()
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ Resolution order for the output directory:
|
|||||||
"""
|
"""
|
||||||
import os
|
import os
|
||||||
import secrets
|
import secrets
|
||||||
from datetime import datetime
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
from common.config import settings
|
from common.config import settings
|
||||||
from common.logging import logger
|
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)
|
effective_dir = resolve_jobs_dir(jobs_dir)
|
||||||
os.makedirs(effective_dir, exist_ok=True)
|
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"
|
name = f"job_{stamp}_{secrets.token_hex(3)}.py"
|
||||||
path = os.path.join(effective_dir, name)
|
path = os.path.join(effective_dir, name)
|
||||||
abs_path = os.path.abspath(path)
|
abs_path = os.path.abspath(path)
|
||||||
|
|||||||
@@ -37,7 +37,7 @@ class PendingStore:
|
|||||||
return Path(self._data_dir) / self._dir_name
|
return Path(self._data_dir) / self._dir_name
|
||||||
|
|
||||||
def _date_path(self, created_at: datetime) -> Path:
|
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"
|
return self.dir_path / f"{created_at.date().isoformat()}.json"
|
||||||
|
|
||||||
def _legacy_path(self) -> Path:
|
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,
|
def _request(method: str, url: str, *, json_body: dict | None = None,
|
||||||
timeout: float = 30.0, verify: bool | str = True,
|
params: dict[str, str] | None = None, timeout: float = 30.0,
|
||||||
auth: httpx.Auth | None = None) -> httpx.Response:
|
verify: bool | str = True, auth: httpx.Auth | None = None) -> httpx.Response:
|
||||||
headers = {"Accept": "application/json"}
|
headers = {"Accept": "application/json"}
|
||||||
logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
|
logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
|
||||||
try:
|
try:
|
||||||
resp = httpx.request(
|
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:
|
except httpx.HTTPError as exc:
|
||||||
logger.error(f"YARN {method} {url} failed: {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]}")
|
logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
|
||||||
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
|
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
|
||||||
logger.info(f"YARN kill {application_id} -> ok")
|
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
|
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):
|
class Connection(BaseModel):
|
||||||
name: str
|
name: str
|
||||||
# Defaults to "yarn" because that's the literal string spark-submit wants
|
# 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_principal: str | None = None
|
||||||
auth_keytab: 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")
|
@field_validator("master")
|
||||||
@classmethod
|
@classmethod
|
||||||
def _check_master(cls, v: str) -> str:
|
def _check_master(cls, v: str) -> str:
|
||||||
|
|||||||
+267
-32
@@ -6,26 +6,41 @@
|
|||||||
from fastapi import FastAPI, Request
|
from fastapi import FastAPI, Request
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
|
from spark_executor.core.yarn_client import YarnError
|
||||||
from spark_executor.tools.connections import (
|
from spark_executor.tools.connections import (
|
||||||
delete_connection,
|
delete_connection,
|
||||||
get_connection,
|
get_connection,
|
||||||
list_connections,
|
list_connections,
|
||||||
save_connection,
|
save_connection,
|
||||||
|
update_connection,
|
||||||
)
|
)
|
||||||
from spark_executor.tools.write_job import write_job_file
|
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.job_file import read_job_file, update_job_file
|
||||||
from spark_executor.tools.kill import kill_job
|
from spark_executor.tools.kill import kill_job
|
||||||
from spark_executor.tools.logs import get_job_logs
|
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 (
|
from spark_executor.tools.requests import (
|
||||||
ConnectionNameRequest,
|
ConnectionNameRequest,
|
||||||
EmptyRequest,
|
EmptyRequest,
|
||||||
WriteJobFileRequest,
|
WriteJobFileRequest,
|
||||||
|
ExternalJobLogsRequest,
|
||||||
|
ExternalJobStatusRequest,
|
||||||
|
ExternalJobResultRequest,
|
||||||
|
ListApplicationsRequest,
|
||||||
|
FetchUrlRequest,
|
||||||
GetJobLogsRequest,
|
GetJobLogsRequest,
|
||||||
JobIdRequest,
|
JobIdRequest,
|
||||||
PendingIdRequest,
|
PendingIdRequest,
|
||||||
PrepareSubmitJobRequest,
|
PrepareSubmitJobRequest,
|
||||||
ReadJobFileRequest,
|
ReadJobFileRequest,
|
||||||
SaveConnectionRequest,
|
SaveConnectionRequest,
|
||||||
|
UpdateConnectionRequest,
|
||||||
UpdateJobFileRequest,
|
UpdateJobFileRequest,
|
||||||
UpdatePendingJobRequest,
|
UpdatePendingJobRequest,
|
||||||
)
|
)
|
||||||
@@ -69,6 +84,20 @@ async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONRespons
|
|||||||
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
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")
|
@app.get("/health")
|
||||||
def health_check():
|
def health_check():
|
||||||
return {"status": "ok"}
|
return {"status": "ok"}
|
||||||
@@ -117,7 +146,9 @@ def _prepare_submit_job(req: PrepareSubmitJobRequest):
|
|||||||
"Actually invoke spark-submit for the PendingSubmission identified "
|
"Actually invoke spark-submit for the PendingSubmission identified "
|
||||||
"by pending_id. Requires status=PENDING. On success, transitions the "
|
"by pending_id. Requires status=PENDING. On success, transitions the "
|
||||||
"pending entry to SUBMITTED and creates a Job record. On failure, "
|
"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):
|
def _confirm_submit_job(req: PendingIdRequest):
|
||||||
@@ -128,7 +159,13 @@ def _confirm_submit_job(req: PendingIdRequest):
|
|||||||
"/list_pending_jobs",
|
"/list_pending_jobs",
|
||||||
operation_id="list_pending_jobs",
|
operation_id="list_pending_jobs",
|
||||||
summary="List all pending submissions",
|
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()):
|
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
|
||||||
return list_pending_jobs()
|
return list_pending_jobs()
|
||||||
@@ -138,7 +175,13 @@ def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
|
|||||||
"/get_pending_job",
|
"/get_pending_job",
|
||||||
operation_id="get_pending_job",
|
operation_id="get_pending_job",
|
||||||
summary="Get a single pending submission",
|
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):
|
def _get_pending_job(req: PendingIdRequest):
|
||||||
return get_pending_job(req.pending_id)
|
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. "
|
"confirm_submit_job: the local job_id (12-char hex, e.g. "
|
||||||
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
||||||
"'application_17400000001_0001'). The lookup is by job_id first, "
|
"'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):
|
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. "
|
"confirm_submit_job: the local job_id (12-char hex, e.g. "
|
||||||
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
||||||
"'application_17400000001_0001'). The lookup is by job_id first, "
|
"'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):
|
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. "
|
"confirm_submit_job: the local job_id (12-char hex, e.g. "
|
||||||
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
|
||||||
"'application_17400000001_0001'). The lookup is by job_id first, "
|
"'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):
|
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. "
|
"PUT state=KILLED to YARN REST API for the job's application_id. "
|
||||||
"**job_id accepts BOTH identifiers** returned by "
|
"**job_id accepts BOTH identifiers** returned by "
|
||||||
"confirm_submit_job: the local job_id (12-char hex) and the YARN "
|
"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):
|
def _kill_job(req: JobIdRequest):
|
||||||
return kill_job(req.job_id)
|
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 ---
|
# --- Connection management tools ---
|
||||||
|
|
||||||
@app.post(
|
@app.post(
|
||||||
@@ -254,21 +398,39 @@ def _kill_job(req: JobIdRequest):
|
|||||||
operation_id="save_connection",
|
operation_id="save_connection",
|
||||||
summary="Save or update a named Spark connection",
|
summary="Save or update a named Spark connection",
|
||||||
description=(
|
description=(
|
||||||
"Upsert a Connection record (master URL, deploy mode, optional YARN RM URL, "
|
"Upsert a Connection record (master URL, deploy mode, optional YARN "
|
||||||
"spark_conf K/V) keyed by name. Used by prepare_submit_job via the "
|
"RM URL, spark_conf K/V) keyed by name. Referenced by "
|
||||||
"connection parameter."
|
"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):
|
def _save_connection(req: SaveConnectionRequest):
|
||||||
# exclude_none so we don't overwrite the function's default with explicit None
|
fields = req.model_dump(exclude_none=True)
|
||||||
return save_connection(**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(
|
@app.post(
|
||||||
"/list_connections",
|
"/list_connections",
|
||||||
operation_id="list_connections",
|
operation_id="list_connections",
|
||||||
summary="List all saved Spark 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()):
|
def _list_connections(_req: EmptyRequest = EmptyRequest()):
|
||||||
return list_connections()
|
return list_connections()
|
||||||
@@ -278,17 +440,58 @@ def _list_connections(_req: EmptyRequest = EmptyRequest()):
|
|||||||
"/get_connection",
|
"/get_connection",
|
||||||
operation_id="get_connection",
|
operation_id="get_connection",
|
||||||
summary="Get a single connection by name",
|
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):
|
def _get_connection(req: ConnectionNameRequest):
|
||||||
return get_connection(req.name)
|
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(
|
@app.post(
|
||||||
"/delete_connection",
|
"/delete_connection",
|
||||||
operation_id="delete_connection",
|
operation_id="delete_connection",
|
||||||
summary="Delete a saved 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):
|
def _delete_connection(req: ConnectionNameRequest):
|
||||||
return delete_connection(req.name)
|
return delete_connection(req.name)
|
||||||
@@ -301,16 +504,18 @@ def _delete_connection(req: ConnectionNameRequest):
|
|||||||
operation_id="write_job_file",
|
operation_id="write_job_file",
|
||||||
summary="Write LLM-authored PySpark code to disk",
|
summary="Write LLM-authored PySpark code to disk",
|
||||||
description=(
|
description=(
|
||||||
"Takes a PySpark code string the LLM has already composed in its "
|
"Persist PySpark code you've already written in your context to a "
|
||||||
"context and writes it to a timestamped file under "
|
"timestamped file under SPARK_EXECUTOR_JOBS_DIR (default "
|
||||||
"SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
|
"./data/jobs/). Returns the absolute path to pass as the "
|
||||||
"path for use as the script_path argument of prepare_submit_job — the "
|
"script_path argument of prepare_submit_job. The two-step pattern "
|
||||||
"two-step pattern means the LLM writes the file, the user can review "
|
"(write the file, then prepare) means the user can review the "
|
||||||
"it (via read_job_file), and only then is the job submitted.\n\n"
|
"file via read_job_file before anything runs.\n\n"
|
||||||
"Note: this tool does NOT generate PySpark code. The calling LLM is "
|
"Prerequisite: you should have already composed the PySpark code "
|
||||||
"expected to have already written the code; this tool only persists "
|
"in your own context before calling this tool — it only persists "
|
||||||
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
|
"code, it does not generate it. Code is run through the SQL safety "
|
||||||
"only) before being written — forbidden statements cause a 400."
|
"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):
|
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",
|
summary="Read the contents of an existing PySpark script",
|
||||||
description=(
|
description=(
|
||||||
"Returns the text content of an existing script file at the given "
|
"Returns the text content of an existing script file at the given "
|
||||||
"path. Caps reads at 1 MB. Typical use: after write_job_file "
|
"path. No size cap. Typical use: after write_job_file returns a "
|
||||||
"returns a path, call read_job_file on that path to inspect what "
|
"path, call read_job_file on that path to inspect what was actually "
|
||||||
"was actually written, before deciding to prepare_submit_job or "
|
"written, before deciding to prepare_submit_job or update_job_file."
|
||||||
"update_job_file."
|
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
def _read_job_file(req: ReadJobFileRequest):
|
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 "
|
"Replaces the entire content of an existing script file. Path must "
|
||||||
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
|
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
|
||||||
"to) — protects against overwriting host-mounted configs or other "
|
"to) — protects against overwriting host-mounted configs or other "
|
||||||
"non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
|
"non-script files. No size cap. Typical use: read_job_file, edit "
|
||||||
"edit the content (LLM or human), update_job_file, then "
|
"the content (LLM or human), update_job_file, then prepare_submit_job "
|
||||||
"prepare_submit_job with the same path."
|
"with the same path."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
def _update_job_file(req: UpdateJobFileRequest):
|
def _update_job_file(req: UpdateJobFileRequest):
|
||||||
return update_job_file(req.script_path, req.content)
|
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
|
@Author :tao.chen
|
||||||
"""
|
"""
|
||||||
from common.logging import logger
|
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
|
from spark_executor.models import Connection
|
||||||
|
|
||||||
|
|
||||||
|
_UNSET = object()
|
||||||
|
|
||||||
|
|
||||||
def save_connection(
|
def save_connection(
|
||||||
*,
|
*,
|
||||||
name: str,
|
name: str,
|
||||||
master: str,
|
master: str,
|
||||||
deploy_mode: str = "cluster",
|
deploy_mode: str = _UNSET, # type: ignore[assignment]
|
||||||
yarn_rm_url: str | None = None,
|
yarn_rm_url: str | None = _UNSET, # type: ignore[assignment]
|
||||||
spark_conf: dict[str, str] | None = None,
|
spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment]
|
||||||
ssl_verify: bool | None = None,
|
ssl_verify: bool | None = _UNSET, # type: ignore[assignment]
|
||||||
ssl_ca_bundle: str | None = None,
|
ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment]
|
||||||
auth_type: str = "none",
|
auth_type: str = _UNSET, # type: ignore[assignment]
|
||||||
auth_user: str | None = None,
|
auth_user: str | None = _UNSET, # type: ignore[assignment]
|
||||||
auth_password: str | None = None,
|
auth_password: str | None = _UNSET, # type: ignore[assignment]
|
||||||
auth_principal: str | None = None,
|
auth_principal: str | None = _UNSET, # type: ignore[assignment]
|
||||||
auth_keytab: str | None = None,
|
auth_keytab: str | None = _UNSET, # type: ignore[assignment]
|
||||||
) -> dict[str, str]:
|
url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment]
|
||||||
|
history_server_url: str | None = _UNSET, # type: ignore[assignment]
|
||||||
|
) -> dict[str, object]:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
|
f"save_connection enter name={name} master={master} "
|
||||||
f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}"
|
f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).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,
|
|
||||||
)
|
)
|
||||||
|
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)
|
store.save(conn)
|
||||||
return {"name": name, "status": "SAVED"}
|
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]]:
|
def list_connections() -> list[dict[str, object]]:
|
||||||
logger.debug("list_connections enter")
|
logger.debug("list_connections enter")
|
||||||
return [c.model_dump() for c in store.list_all()]
|
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
|
- update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR
|
||||||
(settings.jobs_dir) so the agent cannot overwrite host-mounted
|
(settings.jobs_dir) so the agent cannot overwrite host-mounted
|
||||||
configs or arbitrary files on the container FS.
|
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
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -24,9 +27,6 @@ from pathlib import Path
|
|||||||
from common.config import settings
|
from common.config import settings
|
||||||
from common.logging import logger
|
from common.logging import logger
|
||||||
|
|
||||||
MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB
|
|
||||||
|
|
||||||
|
|
||||||
class ScriptFileError(ValueError):
|
class ScriptFileError(ValueError):
|
||||||
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
|
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
|
||||||
handler in server.py.
|
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]:
|
def read_job_file(script_path: str) -> dict[str, object]:
|
||||||
"""Return the text content of an existing script file.
|
"""Return the text content of an existing script file.
|
||||||
|
|
||||||
Caps the read at 1 MB to keep MCP responses bounded; raises
|
No size cap. Raises ScriptFileError (-> 400) if the file is missing.
|
||||||
ScriptFileError (-> 400) if the file is missing or too large.
|
|
||||||
"""
|
"""
|
||||||
_check_readable(script_path)
|
_check_readable(script_path)
|
||||||
size = os.path.getsize(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}")
|
logger.debug(f"read_job_file enter script_path={script_path} size={size}")
|
||||||
with open(script_path, encoding="utf-8") as f:
|
with open(script_path, encoding="utf-8") as f:
|
||||||
content = f.read()
|
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]:
|
def update_job_file(script_path: str, content: str) -> dict[str, object]:
|
||||||
"""Overwrite an existing script file with new content.
|
"""Overwrite an existing script file with new content.
|
||||||
|
|
||||||
Restricted to paths under settings.jobs_dir. Caps writes at 1 MB.
|
Restricted to paths under settings.jobs_dir. No size cap.
|
||||||
Raises ScriptFileError (-> 400) if the path is missing, outside
|
Raises ScriptFileError (-> 400) if the path is missing or outside
|
||||||
the allowed dir, or the content is too large.
|
the allowed dir.
|
||||||
"""
|
"""
|
||||||
_check_writable(script_path)
|
_check_writable(script_path)
|
||||||
encoded_size = len(content.encode("utf-8"))
|
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(
|
logger.debug(
|
||||||
f"update_job_file enter script_path={script_path} "
|
f"update_job_file enter script_path={script_path} "
|
||||||
f"new_bytes={encoded_size}"
|
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}")
|
logger.debug(f"kill_job enter job_id={job_id}")
|
||||||
job = store.get_either(job_id)
|
job = store.get_either(job_id)
|
||||||
if job is None:
|
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)
|
conn = conn_store.get(job.connection)
|
||||||
if conn is None:
|
if conn is None:
|
||||||
raise KeyError(f"Connection not found: {job.connection}")
|
raise KeyError(f"Connection not found: {job.connection}")
|
||||||
|
|||||||
@@ -11,14 +11,31 @@ from spark_executor.tools.connections import store as conn_store
|
|||||||
store = JobStore()
|
store = JobStore()
|
||||||
|
|
||||||
|
|
||||||
def _unknown_job_error(uid: str) -> KeyError:
|
def _unknown_job_error(uid: str, external_tool_hint: str | None = None) -> Exception:
|
||||||
"""Standard "we tried both IDs and found nothing" message.
|
"""Build the right error for a not-found job.
|
||||||
|
|
||||||
The agent gets this from confirm_submit_job's response:
|
The agent gets this from confirm_submit_job's response:
|
||||||
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
|
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
|
||||||
and routinely confuses which to pass here. Spelling out that BOTH
|
and routinely confuses which to pass here. We differentiate two cases:
|
||||||
IDs were tried (and what they look like) saves a round trip.
|
|
||||||
|
- `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(
|
return KeyError(
|
||||||
f"No Job found for id={uid!r} (neither as job_id nor as "
|
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 "
|
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}")
|
logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
|
||||||
job = store.get_either(job_id)
|
job = store.get_either(job_id)
|
||||||
if job is None:
|
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)
|
conn = conn_store.get(job.connection)
|
||||||
if conn is None:
|
if conn is None:
|
||||||
raise KeyError(f"Connection not found: {job.connection}")
|
raise KeyError(f"Connection not found: {job.connection}")
|
||||||
|
|||||||
@@ -17,22 +17,143 @@ class EmptyRequest(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class SaveConnectionRequest(BaseModel):
|
class SaveConnectionRequest(BaseModel):
|
||||||
name: str
|
name: str = Field(
|
||||||
master: str
|
...,
|
||||||
deploy_mode: str = "cluster"
|
description=(
|
||||||
yarn_rm_url: str | None = None
|
"Unique connection name. Referenced by prepare_submit_job.connection "
|
||||||
spark_conf: dict[str, str] | None = None
|
"and get_external_*.connection_name. Saving with an existing name "
|
||||||
ssl_verify: bool | None = None
|
"overwrites that record."
|
||||||
ssl_ca_bundle: str | None = None
|
),
|
||||||
auth_type: str = "none"
|
)
|
||||||
auth_user: str | None = None
|
master: str = Field(
|
||||||
auth_password: str | None = None
|
...,
|
||||||
auth_principal: str | None = None
|
description=(
|
||||||
auth_keytab: str | None = None
|
"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):
|
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(
|
app_name: str = Field(
|
||||||
...,
|
...,
|
||||||
description="Human-readable application name for tracking the pending submission.",
|
description="Human-readable application name for tracking the pending submission.",
|
||||||
@@ -72,34 +193,115 @@ class PrepareSubmitJobRequest(BaseModel):
|
|||||||
|
|
||||||
|
|
||||||
class PendingIdRequest(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):
|
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(
|
script_path: str | None = Field(
|
||||||
default=None,
|
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):
|
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):
|
class GetJobLogsRequest(BaseModel):
|
||||||
job_id: str
|
job_id: str = Field(
|
||||||
tail_chars: int = 5000
|
...,
|
||||||
|
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):
|
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):
|
class WriteJobFileRequest(BaseModel):
|
||||||
@@ -142,3 +344,191 @@ class UpdateJobFileRequest(BaseModel):
|
|||||||
"Maximum 1 MB to keep the MCP response bounded."
|
"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}")
|
logger.debug(f"get_job_result enter job_id={job_id}")
|
||||||
job = store.get_either(job_id)
|
job = store.get_either(job_id)
|
||||||
if job is None:
|
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)
|
conn = conn_store.get(job.connection)
|
||||||
if conn is None:
|
if conn is None:
|
||||||
raise KeyError(f"Connection not found: {job.connection}")
|
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}")
|
logger.debug(f"get_job_status enter job_id={job_id}")
|
||||||
job = store.get_either(job_id)
|
job = store.get_either(job_id)
|
||||||
if job is None:
|
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)
|
conn = conn_store.get(job.connection)
|
||||||
if conn is None:
|
if conn is None:
|
||||||
raise KeyError(f"Connection not found: {job.connection}")
|
raise KeyError(f"Connection not found: {job.connection}")
|
||||||
|
|||||||
@@ -6,9 +6,8 @@
|
|||||||
import os
|
import os
|
||||||
import secrets
|
import secrets
|
||||||
import uuid
|
import uuid
|
||||||
from datetime import datetime
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
from common.config import settings
|
|
||||||
from common.logging import logger
|
from common.logging import logger
|
||||||
from common.sql_guard import validate_pyspark_code
|
from common.sql_guard import validate_pyspark_code
|
||||||
from spark_executor.core.connection_store import store as conn_store
|
from spark_executor.core.connection_store import store as conn_store
|
||||||
@@ -124,7 +123,7 @@ def prepare_submit_job(
|
|||||||
num_executors=num_executors,
|
num_executors=num_executors,
|
||||||
spark_conf=dict(conn.spark_conf),
|
spark_conf=dict(conn.spark_conf),
|
||||||
extra_args=dict(extra_args or {}),
|
extra_args=dict(extra_args or {}),
|
||||||
created_at=datetime.utcnow(),
|
created_at=datetime.now(timezone.utc),
|
||||||
status="PENDING",
|
status="PENDING",
|
||||||
)
|
)
|
||||||
pending_store.save(pending)
|
pending_store.save(pending)
|
||||||
@@ -248,7 +247,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
|||||||
application_id=application_id,
|
application_id=application_id,
|
||||||
script_path=pending.script_path,
|
script_path=pending.script_path,
|
||||||
queue=pending.queue,
|
queue=pending.queue,
|
||||||
submit_time=datetime.utcnow(),
|
submit_time=datetime.now(timezone.utc),
|
||||||
connection=pending.connection,
|
connection=pending.connection,
|
||||||
yarn_rm_url=pending.yarn_rm_url,
|
yarn_rm_url=pending.yarn_rm_url,
|
||||||
)
|
)
|
||||||
@@ -282,7 +281,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
|||||||
application_id=application_id,
|
application_id=application_id,
|
||||||
script_path=pending.script_path,
|
script_path=pending.script_path,
|
||||||
queue=pending.queue,
|
queue=pending.queue,
|
||||||
submit_time=datetime.utcnow(),
|
submit_time=datetime.now(timezone.utc),
|
||||||
connection=pending.connection,
|
connection=pending.connection,
|
||||||
yarn_rm_url=pending.yarn_rm_url,
|
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 import connection_store, pending_store
|
||||||
from spark_executor.core.connection_store import ConnectionStore
|
from spark_executor.core.connection_store import ConnectionStore
|
||||||
from spark_executor.core.pending_store import PendingStore
|
from spark_executor.core.pending_store import PendingStore
|
||||||
|
from spark_executor.models import Connection
|
||||||
from spark_executor.server import app
|
from spark_executor.server import app
|
||||||
from spark_executor.tools import connections, submit
|
from spark_executor.tools import connections, submit
|
||||||
|
|
||||||
@@ -64,6 +65,19 @@ def test_seventeen_tool_routes_registered():
|
|||||||
assert "/update_pending_job" in paths
|
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) ---
|
# --- operation_id: pin clean MCP tool names (no auto-generated suffixes) ---
|
||||||
#
|
#
|
||||||
# fastapi-mcp uses each route's OpenAPI `operationId` as the MCP tool name
|
# 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 r.status_code == 400
|
||||||
assert "only PENDING submissions can be updated" in r.json()["detail"]
|
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
|
# coding=utf-8
|
||||||
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -120,3 +121,113 @@ def test_delete_connection_returns_status(_fresh_store):
|
|||||||
def test_delete_connection_unknown_raises(_fresh_store):
|
def test_delete_connection_unknown_raises(_fresh_store):
|
||||||
with pytest.raises(KeyError):
|
with pytest.raises(KeyError):
|
||||||
connections.delete_connection("missing")
|
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")
|
update_job_file(str(tmp_path / "nope.py"), "x\n")
|
||||||
|
|
||||||
|
|
||||||
def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch):
|
def test_update_accepts_content_larger_than_former_1mb_cap(tmp_path: Path):
|
||||||
"""Cap at 1 MB so the MCP response stays bounded."""
|
"""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)
|
config.settings.jobs_dir = str(tmp_path)
|
||||||
p = tmp_path / "big.py"
|
p = tmp_path / "big.py"
|
||||||
p.write_text("# small\n")
|
p.write_text("# small\n")
|
||||||
# Synthesize 2 MB of content (don't actually write 2 MB to disk)
|
|
||||||
big = "x" * (2 * 1024 * 1024)
|
big = "x" * (2 * 1024 * 1024)
|
||||||
with pytest.raises(ScriptFileError, match="too large"):
|
out = update_job_file(str(p), big)
|
||||||
update_job_file(str(p), big)
|
assert out["bytes_written"] == len(big)
|
||||||
|
assert p.read_text() == 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)
|
|
||||||
|
|
||||||
|
|
||||||
# --- round-trip: write -> read -> update -> read ---
|
# --- round-trip: write -> read -> update -> read ---
|
||||||
|
|||||||
@@ -63,6 +63,14 @@ def test_kill_job_raises_for_unknown_job():
|
|||||||
kill.kill_job("missing")
|
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():
|
def test_kill_job_raises_when_connection_missing():
|
||||||
_fresh_stores()
|
_fresh_stores()
|
||||||
kill.store.put(
|
kill.store.put(
|
||||||
|
|||||||
@@ -68,6 +68,14 @@ def test_get_job_logs_raises_for_unknown_job():
|
|||||||
logs.get_job_logs("missing")
|
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):
|
def test_get_job_logs_raises_when_connection_missing(fresh_stores):
|
||||||
logs.store.put(
|
logs.store.put(
|
||||||
Job(
|
Job(
|
||||||
|
|||||||
@@ -136,6 +136,14 @@ def test_result_raises_keyerror_for_unknown_job():
|
|||||||
assert "application_id" in msg
|
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():
|
def test_result_raises_when_connection_missing():
|
||||||
_fresh_stores()
|
_fresh_stores()
|
||||||
result.store.put(
|
result.store.put(
|
||||||
|
|||||||
@@ -66,6 +66,14 @@ def test_get_job_status_raises_for_unknown_job():
|
|||||||
status.get_job_status("missing")
|
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():
|
def test_get_job_status_raises_when_connection_missing():
|
||||||
_fresh_stores()
|
_fresh_stores()
|
||||||
status.store.put(
|
status.store.put(
|
||||||
|
|||||||
Reference in New Issue
Block a user