Merge pull request 'Feat/fetch url tool' (#5) from feat/fetch-url-tool into main
Reviewed-on: https://gitea-production-a772.up.railway.app/taochen/mcp-server/pulls/5
This commit is contained in:
@@ -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 脚本 |
|
||||||
@@ -120,6 +121,8 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
|
|||||||
| `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态(按 `application_id` + `connection_name`) |
|
| `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态(按 `application_id` + `connection_name`) |
|
||||||
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
|
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
|
||||||
| `get_external_job_logs` | 拉取外部 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 工具
|
||||||
|
|
||||||
|
|||||||
@@ -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,21 @@ 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. "
|
||||||
|
"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'."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
@field_validator("master")
|
@field_validator("master")
|
||||||
@classmethod
|
@classmethod
|
||||||
def _check_master(cls, v: str) -> str:
|
def _check_master(cls, v: str) -> str:
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from spark_executor.tools.connections import (
|
|||||||
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
|
||||||
@@ -20,7 +21,9 @@ from spark_executor.tools.external_jobs import (
|
|||||||
get_external_job_logs,
|
get_external_job_logs,
|
||||||
get_external_job_status,
|
get_external_job_status,
|
||||||
get_external_job_result,
|
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,
|
||||||
@@ -28,12 +31,15 @@ from spark_executor.tools.requests import (
|
|||||||
ExternalJobLogsRequest,
|
ExternalJobLogsRequest,
|
||||||
ExternalJobStatusRequest,
|
ExternalJobStatusRequest,
|
||||||
ExternalJobResultRequest,
|
ExternalJobResultRequest,
|
||||||
|
ListApplicationsRequest,
|
||||||
|
FetchUrlRequest,
|
||||||
GetJobLogsRequest,
|
GetJobLogsRequest,
|
||||||
JobIdRequest,
|
JobIdRequest,
|
||||||
PendingIdRequest,
|
PendingIdRequest,
|
||||||
PrepareSubmitJobRequest,
|
PrepareSubmitJobRequest,
|
||||||
ReadJobFileRequest,
|
ReadJobFileRequest,
|
||||||
SaveConnectionRequest,
|
SaveConnectionRequest,
|
||||||
|
UpdateConnectionRequest,
|
||||||
UpdateJobFileRequest,
|
UpdateJobFileRequest,
|
||||||
UpdatePendingJobRequest,
|
UpdatePendingJobRequest,
|
||||||
)
|
)
|
||||||
@@ -342,6 +348,34 @@ def _get_external_job_result(req: ExternalJobResultRequest):
|
|||||||
return get_external_job_result(req.application_id, req.connection_name)
|
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(
|
||||||
@@ -356,12 +390,20 @@ def _get_external_job_result(req: ExternalJobResultRequest):
|
|||||||
"yarn_rm_url is required for get_external_* to work** — spark-submit "
|
"yarn_rm_url is required for get_external_* to work** — spark-submit "
|
||||||
"can discover the RM for submissions, but direct YARN REST queries "
|
"can discover the RM for submissions, but direct YARN REST queries "
|
||||||
"need an explicit URL. Saving with an existing name overwrites the "
|
"need an explicit URL. Saving with an existing name overwrites the "
|
||||||
"record in place (no version history)."
|
"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):
|
||||||
|
fields = req.model_dump(exclude_none=True)
|
||||||
|
name = fields.pop("name")
|
||||||
|
try:
|
||||||
|
get_connection(name)
|
||||||
|
except KeyError:
|
||||||
# exclude_none so we don't overwrite the function's default with explicit None
|
# exclude_none so we don't overwrite the function's default with explicit None
|
||||||
return save_connection(**req.model_dump(exclude_none=True))
|
return save_connection(name=name, **fields)
|
||||||
|
return update_connection(name=name, **fields)
|
||||||
|
|
||||||
|
|
||||||
@app.post(
|
@app.post(
|
||||||
@@ -394,6 +436,28 @@ 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",
|
||||||
@@ -467,3 +531,26 @@ def _read_job_file(req: ReadJobFileRequest):
|
|||||||
)
|
)
|
||||||
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"
|
||||||
|
"30s timeout, redirects followed."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
def _fetch_url(req: FetchUrlRequest):
|
||||||
|
return fetch_url(req.url, req.connection_name)
|
||||||
|
|||||||
@@ -4,47 +4,94 @@
|
|||||||
@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]
|
||||||
|
) -> 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 []
|
||||||
|
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 [],
|
||||||
|
}
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
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()]
|
||||||
|
|||||||
@@ -10,8 +10,13 @@ YARN application_id and the name of a saved Connection.
|
|||||||
import json
|
import json
|
||||||
|
|
||||||
from common.logging import logger
|
from common.logging import logger
|
||||||
from spark_executor.core.yarn_client import YarnClientConfig, get_application_status, get_application_logs
|
from spark_executor.core.yarn_client import (
|
||||||
from spark_executor.models import JobStatus, JobResult
|
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
|
from spark_executor.tools.connections import store as conn_store
|
||||||
|
|
||||||
|
|
||||||
@@ -63,3 +68,53 @@ def get_external_job_result(application_id: str, connection_name: str) -> JobRes
|
|||||||
)
|
)
|
||||||
logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}")
|
logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}")
|
||||||
return result
|
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,122 @@
|
|||||||
|
# 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):
|
||||||
|
resp = httpx.get(
|
||||||
|
url,
|
||||||
|
auth=auth,
|
||||||
|
verify=verify,
|
||||||
|
timeout=_REQUEST_TIMEOUT_SECONDS,
|
||||||
|
follow_redirects=False,
|
||||||
|
)
|
||||||
|
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})"
|
||||||
|
)
|
||||||
@@ -122,6 +122,18 @@ class SaveConnectionRequest(BaseModel):
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
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."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
class PrepareSubmitJobRequest(BaseModel):
|
class PrepareSubmitJobRequest(BaseModel):
|
||||||
connection: str = Field(
|
connection: str = Field(
|
||||||
@@ -378,3 +390,126 @@ class ExternalJobResultRequest(BaseModel):
|
|||||||
...,
|
...,
|
||||||
description="Name of a saved Connection pointing at the YARN cluster.",
|
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 yarn_rm_url defines the allowed host domain. "
|
||||||
|
"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."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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)."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -64,12 +64,15 @@ def test_seventeen_tool_routes_registered():
|
|||||||
assert "/update_pending_job" in paths
|
assert "/update_pending_job" in paths
|
||||||
|
|
||||||
|
|
||||||
def test_twenty_tool_routes_registered():
|
def test_twenty_three_tool_routes_registered():
|
||||||
paths = {r.path for r in app.routes}
|
paths = {r.path for r in app.routes}
|
||||||
for path in (
|
for path in (
|
||||||
"/get_external_job_logs",
|
"/get_external_job_logs",
|
||||||
"/get_external_job_status",
|
"/get_external_job_status",
|
||||||
"/get_external_job_result",
|
"/get_external_job_result",
|
||||||
|
"/list_applications",
|
||||||
|
"/fetch_url",
|
||||||
|
"/update_connection",
|
||||||
):
|
):
|
||||||
assert path in paths, f"missing MCP tool route: {path}"
|
assert path in paths, f"missing MCP tool route: {path}"
|
||||||
|
|
||||||
|
|||||||
@@ -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,77 @@ 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"
|
||||||
|
|||||||
@@ -5,8 +5,10 @@ from unittest.mock import patch
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from spark_executor.core import connection_store
|
from spark_executor.core import connection_store
|
||||||
from spark_executor.models import Connection
|
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 import connections, external_jobs
|
||||||
|
from spark_executor.tools.requests import ListApplicationsRequest
|
||||||
|
|
||||||
|
|
||||||
def _fresh_stores():
|
def _fresh_stores():
|
||||||
@@ -17,6 +19,13 @@ def _fresh_stores():
|
|||||||
external_jobs.conn_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():
|
def test_get_external_job_logs_returns_tailed():
|
||||||
_fresh_stores()
|
_fresh_stores()
|
||||||
external_jobs.conn_store.save(
|
external_jobs.conn_store.save(
|
||||||
@@ -131,3 +140,166 @@ def test_get_external_job_result_raises_when_connection_missing():
|
|||||||
external_jobs.get_external_job_result(
|
external_jobs.get_external_job_result(
|
||||||
application_id="application_1", connection_name="missing"
|
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,329 @@
|
|||||||
|
# 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)
|
||||||
Reference in New Issue
Block a user