Compare commits

...
Author SHA1 Message Date
ClaudeandClaude deafb5a26b refactor: drop fetch_url body cap and modernize datetime usage
- fetch_url: remove the 1MB body cap and the streaming helper. The
  manual redirect loop with allowlist re-check (the SSRF fix) is kept
  intact; the body is now read in full via resp.text. Drop the now-
  meaningless `truncated` field from FetchUrlResult and the tests that
  asserted on it. Switched from httpx.stream() back to httpx.get()
  for the redirect loop — cleaner without the body cap.

- datetime: replace deprecated datetime.utcnow() with
  datetime.now(timezone.utc) in submit.py (3 sites) and
  core/job_writer.py (1 site). Update the stale comment in
  core/pending_store.py that referenced the old call.

- Clean up an unused `from common.config import settings` import in
  submit.py that ruff flagged.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-09 14:35:16 +08:00
ClaudeandClaude 7d4e512cb0 fix: address review findings on fetch-url-tool
- fetch_url: revalidate allowlist on every redirect hop (fixes SSRF where
  302 to disallowed host / 169.254.169.254 / file:// bypassed the
  url_allowlist). Stream response body with iter_bytes and cap at 1MB
  so a multi-GB response from an allowlisted host cannot OOM the service.
  Reuses the manual-redirect-loop pattern from yarn_client.

- list_applications: stop swallowing 404 (YARN returns 200+empty for
  "no match"; 404 means the RM doesn't support the endpoint — surface
  the YarnError instead of hiding it as an empty result). Add
  Field(ge=1, le=10000) to ListApplicationsRequest.limit so a runaway
  limit is rejected at the Pydantic layer with 422.

- save_connection: PATCH semantics for existing records. Re-route to
  update_connection when the name already exists so partial updates
  (e.g. only master) no longer wipe url_allowlist back to []. Uses an
  _UNSET sentinel in the tool function to distinguish "omitted" from
  "None" without breaking the existing parameter list.

- README: drop leading space on 5 new connection-tool table rows that
  was breaking GitHub Flavored Markdown table continuity.

- Indentation: normalize connections.py and requests.py to 4-space
  indent (auth_password/auth_principal/auth_keytab were 3-space).

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-09 14:23:16 +08:00
ClaudeandClaude Fable 5 7fbad97a87 feat: add list_applications tool for YARN application enumeration
Add a new MCP tool that queries YARN's /ws/v1/cluster/apps endpoint
through a named Connection, returning a list of ApplicationSummary
records. Bypasses the local JobStore — useful for enumerating apps
that were not submitted through this service.

API:
  list_applications(
    connection_name: str,           # required, which YARN cluster
    state: str | None = None,       # YARN state filter: NEW/NEW_SAVING/
                                    # SUBMITTED/ACCEPTED/RUNNING/
                                    # FINISHED/FAILED/KILLED
    queue: str | None = None,       # YARN queue filter
    limit: int = 100,               # cap on returned apps (YARN has no
                                    # offset-based pagination; combine
                                    # state/queue filters for big clusters)
  ) -> list[ApplicationSummary]

Implementation:
  - yarn_client.list_applications(config, *, state, queue, limit) -> list[dict]
    Returns raw YARN app dicts; raises YarnError on 4xx/5xx; returns
    [] on 404 (no apps match). Uses the existing _request helper,
    which now accepts a "params" kwarg for query strings (one-line
    additive change).
  - external_jobs.list_applications(connection_name, state, queue, limit)
    -> list[ApplicationSummary]. Looks up the Connection, builds the
    YarnClientConfig, calls the yarn_client function, maps each raw
    YARN dict to ApplicationSummary (mirroring the manual field-mapping
    style of get_job_result). The yarn_client function is imported
    as "list_applications_yarn" to avoid name collision.
  - ApplicationSummary: 12-field Pydantic model with snake_case names
    (application_id, name, user, queue, state, final_status,
    application_type, application_tags, started_time, finished_time,
    tracking_url, progress). Unused YARN fields (memorySeconds,
    vcoreSeconds, preemptedResource*, etc.) are not exposed.
  - ListApplicationsRequest: Pydantic body model with Field(description=)
    for LLM-facing schema.
  - /list_applications route registered with operation_id=
    "list_applications", placed next to the other external YARN tools.

Tests:
  - 8 new unit tests in test_external_jobs.py (happy path, state/queue/
    limit pass-through, default limit, empty list, missing connection,
    full field mapping).
  - test_mcp_routes.py: assert 23 tool routes.
  - README: list_applications row added to the Spark Executor table.

Tests: 390 passed (was 382, +8 net).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 13:40:20 +08:00
ClaudeandClaude Fable 5 a5b9539663 refactor(fetch_url): trust the allowlist, rename to url_allowlist
Two changes:

1) Drop every check except the allowlist lookup.

  Old _validate_url_host did: scheme check, host-presence check,
  IP-literal check, empty-allowlist check, then glob match.

  New _validate_url_host does: parse host, return on glob match,
  raise on miss. That's it. The only remaining structural check is
  'the URL must have a host' (otherwise the glob has nothing to
  test against).

  Security implication: scheme (file://, gopher://, ftp://) and
  IP literals (10.0.0.1, ::1) are NO LONGER rejected by the
  validator. The allowlist is the single source of truth. If the
  user writes ['*.*.*.*'], they have opted in to 4-label hosts
  including IP literals; if they write ['ccam*'], they get ccam1-
  ccam99 and nothing else. The default ['ccam*'] / [] pattern is
  tight by construction.

  Removed: import ipaddress, the scheme/IP rejection branches, the
  'allowlist empty' explicit branch (the empty list naturally
  matches nothing).

2) Rename allowed_url_hosts -> url_allowlist.

  The previous name was a verbose double-negative ('allowed ... hosts').
  The new name is short, modern (allowlist > whitelist), and matches
  the pattern of the field (URL hosts allowed). Renamed in:
    - Connection (models.py)
    - SaveConnectionRequest, UpdateConnectionRequest, FetchUrlRequest
      (requests.py)
    - _validate_url_host, _host_matches_any_glob parameters
      (fetch_url.py)
    - save_connection / update_connection call sites and error
      messages (connections.py, fetch_url.py)
    - Route descriptions (server.py)
    - All test files
    - README

  No backward-compat alias: the field was added in 0291f36 and
  hasn't shipped, so no production migration. Local dev data
  (data/connections.json, gitignored) with allowed_url_hosts set
  will be silently dropped by Pydantic v2 (default for extra
  fields is ignore) — those connections lose their allowlist and
  fetch_url will reject everything until re-saved.

Test cleanup:
  - Removed 9 obsolete tests (scheme/IP/suffix rejection)
  - Renamed allowed_url_hosts -> url_allowlist in 11 surviving tests
  - Added 4 new tests documenting the 'allowlist is the only gate'
    model: IP literal accepted, HTTPS accepted, no-host rejected,
    empty allowlist rejected, error message mentions url_allowlist

  -1 obsolete test, net -4 from 386 -> 382 tests passing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 12:09:41 +08:00
ClaudeandClaude Fable 5 6d7387b022 feat: simplify fetch_url allowlist + add update_connection tool
Part 1 — fetch_url validation simplification
  Now that Connection.allowed_url_hosts exists, it's the ONLY check.
  The old 2-label suffix-overlap rule against yarn_rm_url is gone.

  - Connection.allowed_url_hosts: list[str] = Field(default_factory=list)
    (was list[str] | None = None). Default empty list means the
    connection has no fetch access; the user must explicitly opt in
    via save_connection or update_connection.
  - _validate_url_host drops the yarn_rm_url parameter and the
    _host_suffix_overlap helper. Now: scheme/host/IP checks, then
    empty-reject, then glob match. Reject everything else.
  - fetch_url's error message now points the user at
    Connection.allowed_url_hosts as the fix.
  - The label-by-label glob check from the previous commit is kept
    (SSRF guard: 'ccam*' matches 'ccam50' but NOT 'ccam50.evil.com').

Part 2 — update_connection tool
  PATCH-style update for an existing Connection record. Only the
  fields the caller provides are changed. Same pattern as
  update_pending_job: req.model_dump(exclude_none=True), with
  'name' popped before passing to the store.

  To CLEAR a field (e.g. drop auth_password), use delete_connection
  + save_connection. This is YAGNI; the alternative (model_fields_set
  to distinguish 'omitted' from 'null') adds surface for bugs.

  - ConnectionStore.update(name, **fields): get, model_copy(update=...),
    save. Lock + atomic write. Re-validates the patched record.
  - update_connection tool function: passes fields through to the
    store, logs which fields were changed.
  - UpdateConnectionRequest Pydantic model: 12 mutable fields + name.
  - /update_connection route with operation_id='update_connection'.
  - 9 new unit tests in test_connection_tools.py (PATCH, dict/list
    replace-not-merge, unknown name 404, validation of patched record,
    disk persistence).
  - test_fetch_url.py: existing 21 tests updated; 3 new tests for
    the empty/None/missing allowed_url_hosts cases.
  - test_mcp_routes.py: assert 22 tool routes.
  - README: update_connection row added; fetch_url row updated.

Tests: 386 passed (up from 377, +9 net).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 11:46:07 +08:00
ClaudeandClaude Fable 5 0291f36b01 feat(fetch_url): per-Connection allowed_url_hosts glob allowlist
The default 'URL host shares >= 2 labels of suffix with yarn_rm_url
host' rule is too strict for clusters whose hostnames are single-label
(e.g. 'ccam1' through 'ccam99'). The user's cluster is reachable at
http://ccam1:8088, and they want fetch_url to work for any ccamN
host — but ccam1 and ccam50 share 0 suffix labels, so the existing
rule rejects everything.

Add an opt-in allowlist field on Connection:

  allowed_url_hosts: list[str] | None

Each entry is an fnmatch glob pattern. The URL host is allowed if it
matches ANY pattern, regardless of the suffix rule. Save a Connection
with ['ccam*'] to allow ccam1, ccam2, ..., ccam99.

SSRF safety: the implementation does NOT use vanilla fnmatch on the
flat host string (because '*' in fnmatch crosses '.', so 'ccam*' would
match 'ccam50.evil.com' — a security hole). Instead, both the pattern
and the host are split on '.' and matched LABEL-BY-LABEL with the
label counts required to match exactly. So 'ccam*' matches 'ccam50'
but NOT 'ccam50.evil.com' (different label counts).

Test coverage:
- 8 new tests in tests/unit/test_fetch_url.py (glob allow, glob deny,
  dot-boundary SSRF test, multiple globs, empty list, None fallback,
  error message hint)
- All 21 fetch_url tests pass; 377 total.

- spark_executor/models.py: Connection.allowed_url_hosts (with description)
- spark_executor/tools/requests.py: SaveConnectionRequest.allowed_url_hosts
- spark_executor/tools/fetch_url.py: new _host_matches_any_glob helper;
  _validate_url_host now takes allowed_hosts and checks globs BEFORE
  the suffix rule
- tests/unit/test_fetch_url.py: 8 new tests
- README.md: fetch_url row mentions the allowlist

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 11:22:55 +08:00
ClaudeandClaude Fable 5 627e70f697 feat: add fetch_url tool for proxying HTTP GET to cluster-internal URLs
Add a new MCP tool that lets the agent fetch URLs on the cluster's
network (YARN tracking UI, Spark History Server, NodeManager web UIs)
when the agent is on a different network and cannot reach those hosts
directly.

The MCP service runs on the YARN RM node, so it can reach every host
the cluster knows about — the agent just needs a way to ask.

Security: SSRF guard via host suffix overlap
  - URL host must share >= 2 labels of suffix with the named
    Connection's yarn_rm_url host (e.g. yarn_rm_url='rm.prod.internal'
    allows 'http://nm01.prod.internal/...')
  - IP literals (10.0.0.1, ::1) rejected
  - Non-http(s) schemes (file://, gopher://, ftp://) rejected
  - Connection with no yarn_rm_url cannot use this tool
- Reuses Connection.auth_for_httpx() and verify_for_httpx() so the
  agent does not need cluster credentials
- Response body capped at 1 MB (truncated=true if larger)
- 30s timeout, follows redirects, loguru INFO audit log on every call

- spark_executor/tools/fetch_url.py: new tool + 2 helpers
  (_host_suffix_overlap, _validate_url_host)
- spark_executor/models.py: FetchUrlResult Pydantic model
- spark_executor/tools/requests.py: FetchUrlRequest with descriptions
- spark_executor/server.py: /fetch_url route, operation_id='fetch_url'
- tests/unit/test_fetch_url.py: 13 unit tests covering all guards,
  truncation, auth/SSL pass-through, redirect follow
- tests/integration/test_mcp_routes.py: assert 21 tool routes
- README.md: 1 row in Spark Executor 工具 table

Tests: 369 passed (up from 356).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 10:59:25 +08:00
tao.chen f8536b63ad Merge pull request 'feat: add external job tools + improve LLM-facing tool descriptions' (#4) from feat/external-job-tools into main
Reviewed-on: https://gitea-production-a772.up.railway.app/taochen/mcp-server/pulls/4
2026-07-08 12:04:57 +00:00
ClaudeandClaude Fable 5 6cf68439a2 feat: add external job tools + improve LLM-facing tool descriptions
Add 3 new MCP tools for inspecting YARN applications NOT submitted
through this service: get_external_job_logs, get_external_job_status,
get_external_job_result. Each takes application_id + connection_name
and queries YARN directly, bypassing the local JobStore.

- spark_executor/tools/external_jobs.py: 3 tool functions
- spark_executor/tools/requests.py: 3 new Pydantic body models
  (ExternalJobLogsRequest, ExternalJobStatusRequest,
  ExternalJobResultRequest)
- spark_executor/server.py: 3 new POST routes with explicit operation_id
- tests/unit/test_external_jobs.py: 7 unit tests
- tests/integration/test_mcp_routes.py: assert 20 tool routes
- README.md: list the 3 new tools

To make the LLM pick the right tool and not guess at field values,
also:

- Add Pydantic field descriptions for 22 fields across 8 request models
  (SaveConnectionRequest, UpdatePendingJobRequest, GetJobLogsRequest,
  JobIdRequest, PendingIdRequest, ConnectionNameRequest, plus the new
  ExternalJob*Request models).
- Update 12 route descriptions with cross-references, prerequisite
  context, and 400 behavior notes.
- Refactor _unknown_job_error: an input that looks like a YARN
  application_id (starts with 'application_') now returns HTTP 400
  (ValueError) with a hint message naming the right external tool;
  other not-found cases still return 404 (KeyError). This catches the
  common LLM mistake of passing application_id to the internal
  get_job_* / kill_job tools.
- 4 new unit tests for the 400 behavior.

Tests: 356 passed (up from 242).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-08 20:03:40 +08:00
24 changed files with 1871 additions and 98 deletions
+6
View File
@@ -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 工具
+23
View File
@@ -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()
+2 -2
View File
@@ -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)
+1 -1
View File
@@ -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:
+52 -3
View File
@@ -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 []
+44
View File
@@ -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:
+229 -23
View File
@@ -11,21 +11,35 @@ 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
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,
) )
@@ -117,7 +131,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 +144,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 +160,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 +214,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 +240,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 +265,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):
@@ -241,12 +287,95 @@ def _get_job_logs(req: GetJobLogsRequest):
"**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 +383,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):
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(
"/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 +425,51 @@ 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 Connection record, or 404 if not found. Useful to "
"verify a connection was saved correctly, or to inspect the "
"resolved yarn_rm_url / auth_type before submitting a job or "
"calling one of the get_external_* tools."
),
) )
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 +482,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):
@@ -348,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)
+75 -28
View File
@@ -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()]
+120
View File
@@ -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
+122
View File
@@ -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})"
)
+7 -1
View File
@@ -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}")
+25 -5
View File
@@ -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}")
+397 -26
View File
@@ -17,22 +17,135 @@ 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."
),
)
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 +185,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 +336,180 @@ 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 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)."
),
)
+4 -1
View File
@@ -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}")
+4 -1
View File
@@ -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}")
+4 -5
View File
@@ -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,
) )
+13
View File
@@ -64,6 +64,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
+75
View File
@@ -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"
+305
View File
@@ -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")
+329
View File
@@ -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)
+8
View File
@@ -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(
+8
View File
@@ -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(
+8
View File
@@ -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(
+8
View File
@@ -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(