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` | 列出所有连接配置 |
| `get_connection` | 按名称读取连接配置 |
| `delete_connection` | 删除连接配置 |
| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) |
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
| `read_job_file` | 读取已存在的 PySpark 脚本内容 |
| `update_job_file` | 覆盖更新已存在的 PySpark 脚本 |
@@ -117,6 +118,11 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
| `get_job_result` | 查询终态结果视图 |
| `get_job_logs` | 拉取 YARN 聚合日志 |
| `kill_job` | Kill YARN application |
| `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态(按 `application_id` + `connection_name` |
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
| `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 |
| `list_applications` | 列出 YARN 上所有应用(按 `state` / `queue` / `limit` 过滤),绕过 JobStore |
| `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.url_allowlist` glob allowlist 约束, 空则全拒) |
### Files MCP 工具
+23
View File
@@ -77,6 +77,29 @@ class ConnectionStore:
self._dump(records)
logger.info(f"connection saved name={conn.name} master={conn.master}")
def update(self, name: str, **fields) -> Connection:
"""Apply `fields` to the existing Connection identified by `name` and persist.
PATCH semantics: only fields explicitly passed in `fields` are changed.
Use Pydantic's `model_copy(update=fields)` to apply the patch.
Raises KeyError if no Connection with `name` exists.
Raises pydantic.ValidationError if the patched Connection is invalid
(e.g. `master='http://...'` fails the master validator).
"""
with self._lock:
records = self._load()
if name not in records:
raise KeyError(f"Connection not found: {name}")
existing = records[name]
patched = Connection.model_validate(existing.model_copy(update=fields).model_dump())
records[name] = patched
self._dump(records)
logger.info(
f"connection updated name={name} fields={sorted(fields.keys())}"
)
return patched
def delete(self, name: str) -> bool:
with self._lock:
records = self._load()
+2 -2
View File
@@ -14,7 +14,7 @@ Resolution order for the output directory:
"""
import os
import secrets
from datetime import datetime
from datetime import datetime, timezone
from common.config import settings
from common.logging import logger
@@ -63,7 +63,7 @@ def write_job_file(code: str, jobs_dir: str | None = None) -> str:
effective_dir = resolve_jobs_dir(jobs_dir)
os.makedirs(effective_dir, exist_ok=True)
stamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
name = f"job_{stamp}_{secrets.token_hex(3)}.py"
path = os.path.join(effective_dir, name)
abs_path = os.path.abspath(path)
+1 -1
View File
@@ -37,7 +37,7 @@ class PendingStore:
return Path(self._data_dir) / self._dir_name
def _date_path(self, created_at: datetime) -> Path:
# created_at is datetime.utcnow(), so the shard date is a UTC date.
# created_at is timezone-aware UTC, so the shard date is a UTC date.
return self.dir_path / f"{created_at.date().isoformat()}.json"
def _legacy_path(self) -> Path:
+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,
timeout: float = 30.0, verify: bool | str = True,
auth: httpx.Auth | None = None) -> httpx.Response:
params: dict[str, str] | None = None, timeout: float = 30.0,
verify: bool | str = True, auth: httpx.Auth | None = None) -> httpx.Response:
headers = {"Accept": "application/json"}
logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
try:
resp = httpx.request(
method, url, json=json_body, headers=headers, timeout=timeout, verify=verify, auth=auth
method, url, json=json_body, params=params, headers=headers,
timeout=timeout, verify=verify, auth=auth
)
except httpx.HTTPError as exc:
logger.error(f"YARN {method} {url} failed: {exc}")
@@ -257,3 +258,51 @@ def kill_application(application_id: str, config: YarnClientConfig) -> None:
logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
logger.info(f"YARN kill {application_id} -> ok")
def list_applications(
config: YarnClientConfig,
*,
state: str | None = None,
queue: str | None = None,
limit: int | None = None,
) -> list[dict]:
"""List YARN applications, optionally filtered.
YARN endpoint: GET /ws/v1/cluster/apps?state=...&queue=...&limit=...
Filters:
- state: YARN application state. Common values:
"NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING",
"FINISHED", "FAILED", "KILLED".
Note: "FINISHED" is the umbrella state covering SUCCEEDED/FAILED/KILLED.
- queue: YARN queue name
- limit: cap on number of returned apps (YARN has no pagination;
callers that need a full enumeration should make multiple
calls with state=... filters or accept the cap)
Returns a list of YARN app dicts (each with id, name, user, queue,
state, finalStatus, applicationType, startedTime, finishedTime,
trackingUrl, progress, etc). Empty list if no apps match.
Raises YarnError on transport / 4xx / 5xx.
"""
params: dict[str, str] = {}
if state is not None:
params["state"] = state
if queue is not None:
params["queue"] = queue
if limit is not None:
params["limit"] = str(limit)
url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps"
resp = _request("GET", url, params=params,
verify=config.verify_for_httpx(),
auth=config.auth_for_httpx())
if resp.status_code >= 400:
raise YarnError(
f"YARN list applications failed: {resp.status_code} {resp.text[:200]}"
)
data = resp.json()
apps_container = data.get("apps") or {}
return apps_container.get("app", []) or []
+44
View File
@@ -40,6 +40,35 @@ class SubmitResult(BaseModel):
tracking_url: str | None = None
class FetchUrlResult(BaseModel):
url: str
status_code: int
content_type: str
body: str
class ApplicationSummary(BaseModel):
"""A YARN application summary from /ws/v1/cluster/apps.
Field names are mapped from the YARN JSON keys to clearer
snake_case names by the tool function. Unused YARN fields
(memorySeconds, vcoreSeconds, preemptedResource*, etc.) are
not exposed — the LLM doesn't need them.
"""
application_id: str
name: str
user: str
queue: str
state: str
final_status: str | None = None
application_type: str | None = None
application_tags: str = ""
started_time: int = 0
finished_time: int = 0
tracking_url: str | None = None
progress: float | None = None
class Connection(BaseModel):
name: str
# Defaults to "yarn" because that's the literal string spark-submit wants
@@ -61,6 +90,21 @@ class Connection(BaseModel):
auth_principal: str | None = None
auth_keytab: str | None = None
url_allowlist: list[str] = Field(
default_factory=list,
description=(
"List of fnmatch glob patterns for hosts the fetch_url tool may access. "
"The list is mandatory-opt-in: an empty list (the default) denies all "
"hosts, so you must populate it before fetch_url can access any URL. "
"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")
@classmethod
def _check_master(cls, v: str) -> str:
+231 -25
View File
@@ -11,21 +11,35 @@ from spark_executor.tools.connections import (
get_connection,
list_connections,
save_connection,
update_connection,
)
from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_job_file
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.external_jobs import (
get_external_job_logs,
get_external_job_status,
get_external_job_result,
list_applications,
)
from spark_executor.tools.fetch_url import fetch_url
from spark_executor.tools.requests import (
ConnectionNameRequest,
EmptyRequest,
WriteJobFileRequest,
ExternalJobLogsRequest,
ExternalJobStatusRequest,
ExternalJobResultRequest,
ListApplicationsRequest,
FetchUrlRequest,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
PrepareSubmitJobRequest,
ReadJobFileRequest,
SaveConnectionRequest,
UpdateConnectionRequest,
UpdateJobFileRequest,
UpdatePendingJobRequest,
)
@@ -117,7 +131,9 @@ def _prepare_submit_job(req: PrepareSubmitJobRequest):
"Actually invoke spark-submit for the PendingSubmission identified "
"by pending_id. Requires status=PENDING. On success, transitions the "
"pending entry to SUBMITTED and creates a Job record. On failure, "
"marks the entry FAILED and re-raises."
"marks the entry FAILED and re-raises. A FAILED pending can be "
"re-confirmed — it resets to PENDING for a single fresh attempt — so "
"transient failures (e.g. YARN RM was down) are recoverable."
),
)
def _confirm_submit_job(req: PendingIdRequest):
@@ -128,7 +144,13 @@ def _confirm_submit_job(req: PendingIdRequest):
"/list_pending_jobs",
operation_id="list_pending_jobs",
summary="List all pending submissions",
description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).",
description=(
"Return every PendingSubmission in any status (PENDING, SUBMITTED, "
"CANCELLED, FAILED). Call this before prepare_submit_job to check if a "
"submission with the same parameters is already in flight, or after a "
"batch of confirm_submit_job calls to inspect the lifecycle of recent "
"submissions."
),
)
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
return list_pending_jobs()
@@ -138,7 +160,13 @@ def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
"/get_pending_job",
operation_id="get_pending_job",
summary="Get a single pending submission",
description="Return the PendingSubmission identified by pending_id, including its current status and outcome fields.",
description=(
"Return the PendingSubmission identified by pending_id, including its "
"current status and outcome fields. Use this to inspect a pending "
"submission between prepare_submit_job and confirm_submit_job (e.g. "
"to confirm the snapshotted connection), or to read the error field "
"of a FAILED submission before re-confirming."
),
)
def _get_pending_job(req: PendingIdRequest):
return get_pending_job(req.pending_id)
@@ -186,7 +214,13 @@ def _cancel_pending_job(req: PendingIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_status(application_id, connection_name)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_status(req: JobIdRequest):
@@ -206,7 +240,13 @@ def _get_job_status(req: JobIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_result(application_id, connection_name)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_result(req: JobIdRequest):
@@ -225,7 +265,13 @@ def _get_job_result(req: JobIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_logs(application_id, connection_name, tail_chars)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_logs(req: GetJobLogsRequest):
@@ -240,13 +286,96 @@ def _get_job_logs(req: GetJobLogsRequest):
"PUT state=KILLED to YARN REST API for the job's application_id. "
"**job_id accepts BOTH identifiers** returned by "
"confirm_submit_job: the local job_id (12-char hex) and the YARN "
"application_id. The lookup is by job_id first, then by application_id."
"application_id. The lookup is by job_id first, then by application_id. "
"**If you pass a YARN application_id and the app is NOT in the "
"local JobStore, this tool returns HTTP 400** (not 404) with a "
"hint pointing to the YARN CLI / UI. **This tool only works for "
"jobs submitted through this service**; there is no external "
"equivalent. For YARN applications you did not submit here, use "
"the YARN CLI / UI directly to kill them."
),
)
def _kill_job(req: JobIdRequest):
return kill_job(req.job_id)
# --- External YARN job tools (bypass JobStore) ---
@app.post(
"/get_external_job_logs",
operation_id="get_external_job_logs",
summary="Query YARN logs for an application not submitted through this service",
description=(
"Fetch aggregated container logs for a YARN application using its "
"application_id and a saved Connection. This bypasses the local JobStore, "
"so it works for jobs submitted outside this MCP service. "
"application_id format is 'application_<14-digit-timestamp>_<sequence>'. "
"For jobs submitted via this service, use get_job_logs(job_id=...) instead."
),
)
def _get_external_job_logs(req: ExternalJobLogsRequest):
return get_external_job_logs(req.application_id, req.connection_name, req.tail_chars)
@app.post(
"/get_external_job_status",
operation_id="get_external_job_status",
summary="Query YARN status for an application not submitted through this service",
description=(
"Return the YARN application state and raw REST response for an "
"application using its application_id and a saved Connection. "
"This bypasses the local JobStore, so it works for jobs submitted "
"outside this MCP service. For jobs submitted via this service, "
"use get_job_status(job_id=...) instead."
),
)
def _get_external_job_status(req: ExternalJobStatusRequest):
return get_external_job_status(req.application_id, req.connection_name)
@app.post(
"/get_external_job_result",
operation_id="get_external_job_result",
summary="Query YARN terminal result for an application not submitted through this service",
description=(
"Return a terminal-oriented view (final_status, diagnostics, tracking_url, "
"started_time, finished_time) for a YARN application using its application_id "
"and a saved Connection. This bypasses the local JobStore. "
"For jobs submitted via this service, use get_job_result(job_id=...) instead."
),
)
def _get_external_job_result(req: ExternalJobResultRequest):
return get_external_job_result(req.application_id, req.connection_name)
@app.post(
"/list_applications",
operation_id="list_applications",
summary="List YARN applications on a cluster, optionally filtered",
description=(
"Query YARN's /ws/v1/cluster/apps endpoint through the named "
"Connection, returning a list of ApplicationSummary records. "
"Bypasses the local JobStore — useful for enumerating apps that "
"were not submitted through this service.\n\n"
"**Filters:** state (YARN state, e.g. 'RUNNING', 'FINISHED', "
"'FAILED'), queue (YARN queue name), limit (default 100, max ~10000). "
"YARN has no offset-based pagination, so for large clusters combine "
"state/queue filters to scope the result. The `FINISHED` state "
"covers SUCCEEDED/FAILED/KILLED.\n\n"
"Returns an empty list if no apps match. The Connection's auth_type "
"/ auth_user / auth_password / ssl_verify / ssl_ca_bundle are reused "
"for the request."
),
)
def _list_applications(req: ListApplicationsRequest):
return list_applications(
req.connection_name,
state=req.state,
queue=req.queue,
limit=req.limit,
)
# --- Connection management tools ---
@app.post(
@@ -254,21 +383,39 @@ def _kill_job(req: JobIdRequest):
operation_id="save_connection",
summary="Save or update a named Spark connection",
description=(
"Upsert a Connection record (master URL, deploy mode, optional YARN RM URL, "
"spark_conf K/V) keyed by name. Used by prepare_submit_job via the "
"connection parameter."
"Upsert a Connection record (master URL, deploy mode, optional YARN "
"RM URL, spark_conf K/V) keyed by name. Referenced by "
"prepare_submit_job via the connection parameter, and by the 3 "
"get_external_* tools via connection_name. **The Connection's "
"yarn_rm_url is required for get_external_* to work** — spark-submit "
"can discover the RM for submissions, but direct YARN REST queries "
"need an explicit URL. Saving with an existing name overwrites the "
"record in place (no version history). When the name already exists, "
"only the provided fields are changed (PATCH semantics); omitted "
"fields keep their previous values."
),
)
def _save_connection(req: SaveConnectionRequest):
# exclude_none so we don't overwrite the function's default with explicit None
return save_connection(**req.model_dump(exclude_none=True))
fields = req.model_dump(exclude_none=True)
name = fields.pop("name")
try:
get_connection(name)
except KeyError:
# exclude_none so we don't overwrite the function's default with explicit None
return save_connection(name=name, **fields)
return update_connection(name=name, **fields)
@app.post(
"/list_connections",
operation_id="list_connections",
summary="List all saved Spark connections",
description="Return every Connection in the registry (model_dump form).",
description=(
"Return every Connection in the registry (model_dump form). Call "
"this before save_connection to see existing names (saving with an "
"existing name overwrites), or after save_connection to verify the "
"record you just stored."
),
)
def _list_connections(_req: EmptyRequest = EmptyRequest()):
return list_connections()
@@ -278,17 +425,51 @@ def _list_connections(_req: EmptyRequest = EmptyRequest()):
"/get_connection",
operation_id="get_connection",
summary="Get a single connection by name",
description="Return the Connection record, or 404 if not found.",
description=(
"Return the 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):
return get_connection(req.name)
@app.post(
"/update_connection",
operation_id="update_connection",
summary="Update an existing connection's fields",
description=(
"Apply a partial update (PATCH) to an existing Connection record. "
"Only the fields you provide are changed; the rest are kept as-is. "
"The `name` is the immutable identifier (use delete_connection + "
"save_connection to rename).\n\n"
"To CLEAR an optional field (e.g. remove `yarn_rm_url`), use "
"delete_connection followed by save_connection with the field omitted. "
"This tool cannot clear fields — only replace them.\n\n"
"Returns the full updated Connection record. 404 if no Connection "
"with the given name exists."
),
)
def _update_connection(req: UpdateConnectionRequest):
fields = req.model_dump(exclude_none=True)
fields.pop("name", None) # name is the identity, not a field to patch
return update_connection(name=req.name, **fields)
@app.post(
"/delete_connection",
operation_id="delete_connection",
summary="Delete a saved connection",
description="Remove a Connection by name. 404 if not found.",
description=(
"Remove a Connection by name. 404 if not found. Deleting a "
"Connection does NOT affect any pending submission or running job "
"that already references it (the connection details are snapshotted "
"at prepare_submit_job time, and YARN holds the live submission "
"state). New prepare_submit_job calls will fail until you re-save "
"the connection with the same name."
),
)
def _delete_connection(req: ConnectionNameRequest):
return delete_connection(req.name)
@@ -301,16 +482,18 @@ def _delete_connection(req: ConnectionNameRequest):
operation_id="write_job_file",
summary="Write LLM-authored PySpark code to disk",
description=(
"Takes a PySpark code string the LLM has already composed in its "
"context and writes it to a timestamped file under "
"SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
"path for use as the script_path argument of prepare_submit_job — the "
"two-step pattern means the LLM writes the file, the user can review "
"it (via read_job_file), and only then is the job submitted.\n\n"
"Note: this tool does NOT generate PySpark code. The calling LLM is "
"expected to have already written the code; this tool only persists "
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
"only) before being written forbidden statements cause a 400."
"Persist PySpark code you've already written in your context to a "
"timestamped file under SPARK_EXECUTOR_JOBS_DIR (default "
"./data/jobs/). Returns the absolute path to pass as the "
"script_path argument of prepare_submit_job. The two-step pattern "
"(write the file, then prepare) means the user can review the "
"file via read_job_file before anything runs.\n\n"
"Prerequisite: you should have already composed the PySpark code "
"in your own context before calling this tool — it only persists "
"code, it does not generate it. Code is run through the SQL safety "
"policy (SELECT/INSERT only) before being written; forbidden "
"statements cause a 400 with details about which line broke the "
"policy."
),
)
def _write_job_file(req: WriteJobFileRequest):
@@ -348,3 +531,26 @@ def _read_job_file(req: ReadJobFileRequest):
)
def _update_job_file(req: UpdateJobFileRequest):
return update_job_file(req.script_path, req.content)
# --- HTTP fetch proxy (host allowlist via Connection.yarn_rm_url) ---
@app.post(
"/fetch_url",
operation_id="fetch_url",
summary="Fetch a URL on the cluster's network and return the body",
description=(
"Proxy an HTTP GET to a URL on the cluster's network, returning the "
"response body. Useful when the agent is on a different network from "
"the cluster and cannot reach YARN tracking pages, Spark History "
"Server, or NodeManager web UIs directly.\n\n"
"**Security constraints:** the URL host must match one of the fnmatch "
"glob patterns in the named Connection's url_allowlist. An empty or "
"omitted allowlist denies every host. There are no scheme or IP-literal "
"guardrails — the allowlist is the only gate — so keep it tight. The "
"Connection's saved auth is reused, so the agent does not need cluster "
"credentials.\n\n"
"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
"""
from common.logging import logger
from spark_executor.core.connection_store import ConnectionStore, store
from spark_executor.core.connection_store import store
from spark_executor.models import Connection
_UNSET = object()
def save_connection(
*,
name: str,
master: str,
deploy_mode: str = "cluster",
yarn_rm_url: str | None = None,
spark_conf: dict[str, str] | None = None,
ssl_verify: bool | None = None,
ssl_ca_bundle: str | None = None,
auth_type: str = "none",
auth_user: str | None = None,
auth_password: str | None = None,
auth_principal: str | None = None,
auth_keytab: str | None = None,
) -> dict[str, str]:
deploy_mode: str = _UNSET, # type: ignore[assignment]
yarn_rm_url: str | None = _UNSET, # type: ignore[assignment]
spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment]
ssl_verify: bool | None = _UNSET, # type: ignore[assignment]
ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment]
auth_type: str = _UNSET, # type: ignore[assignment]
auth_user: str | None = _UNSET, # type: ignore[assignment]
auth_password: str | None = _UNSET, # type: ignore[assignment]
auth_principal: str | None = _UNSET, # type: ignore[assignment]
auth_keytab: str | None = _UNSET, # type: ignore[assignment]
url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment]
) -> dict[str, object]:
logger.debug(
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}"
)
conn = Connection(
name=name,
master=master,
deploy_mode=deploy_mode,
yarn_rm_url=yarn_rm_url,
spark_conf=spark_conf or {},
ssl_verify=ssl_verify,
ssl_ca_bundle=ssl_ca_bundle,
auth_type=auth_type,
auth_user=auth_user,
auth_password=auth_password,
auth_principal=auth_principal,
auth_keytab=auth_keytab,
f"save_connection enter name={name} master={master} "
f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}"
)
existing = store.get(name)
if existing is not None:
fields: dict[str, object] = {"master": master}
if deploy_mode is not _UNSET:
fields["deploy_mode"] = deploy_mode
if yarn_rm_url is not _UNSET:
fields["yarn_rm_url"] = yarn_rm_url
if spark_conf is not _UNSET:
fields["spark_conf"] = spark_conf or {}
if ssl_verify is not _UNSET:
fields["ssl_verify"] = ssl_verify
if ssl_ca_bundle is not _UNSET:
fields["ssl_ca_bundle"] = ssl_ca_bundle
if auth_type is not _UNSET:
fields["auth_type"] = auth_type
if auth_user is not _UNSET:
fields["auth_user"] = auth_user
if auth_password is not _UNSET:
fields["auth_password"] = auth_password
if auth_principal is not _UNSET:
fields["auth_principal"] = auth_principal
if auth_keytab is not _UNSET:
fields["auth_keytab"] = auth_keytab
if url_allowlist is not _UNSET:
fields["url_allowlist"] = url_allowlist or []
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)
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]]:
logger.debug("list_connections enter")
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}")
job = store.get_either(job_id)
if job is None:
raise _unknown_job_error(job_id)
raise _unknown_job_error(
job_id,
external_tool_hint=(
"the YARN CLI (`yarn application -kill <app_id>`) or the YARN "
"UI directly — this service has no external kill tool"
),
)
conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
+25 -5
View File
@@ -11,14 +11,31 @@ from spark_executor.tools.connections import store as conn_store
store = JobStore()
def _unknown_job_error(uid: str) -> KeyError:
"""Standard "we tried both IDs and found nothing" message.
def _unknown_job_error(uid: str, external_tool_hint: str | None = None) -> Exception:
"""Build the right error for a not-found job.
The agent gets this from confirm_submit_job's response:
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
and routinely confuses which to pass here. Spelling out that BOTH
IDs were tried (and what they look like) saves a round trip.
and routinely confuses which to pass here. We differentiate two cases:
- `uid` looks like a YARN application_id (starts with 'application_')
AND the caller passed an `external_tool_hint`: the YARN app likely
exists, this tool just can't serve it because it was not submitted
through this service. Raise ValueError (-> 400 via the FastAPI
handler) pointing the agent at the right external tool.
- Otherwise: no local record of either form of id. Raise KeyError
(-> 404). Spelling out what both IDs look like saves a round trip.
"""
if uid.startswith("application_") and external_tool_hint:
return ValueError(
f"job_id={uid!r} looks like a YARN application_id (starts with "
f"'application_'), but this tool only works for jobs submitted "
f"through this MCP service (no local JobStore record). For YARN "
f"applications not submitted here, use {external_tool_hint} "
f"instead. (If you actually submitted this job through this "
f"service, pass the local job_id — it is a 12-char hex like "
f"'a1b2c3d4e5f6'.)"
)
return KeyError(
f"No Job found for id={uid!r} (neither as job_id nor as "
f"application_id). Pass the job_id from confirm_submit_job's "
@@ -37,7 +54,10 @@ def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
job = store.get_either(job_id)
if job is None:
raise _unknown_job_error(job_id)
raise _unknown_job_error(
job_id,
external_tool_hint="get_external_job_logs(application_id, connection_name, tail_chars)",
)
conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
+397 -26
View File
@@ -17,22 +17,135 @@ class EmptyRequest(BaseModel):
class SaveConnectionRequest(BaseModel):
name: str
master: str
deploy_mode: str = "cluster"
yarn_rm_url: str | None = None
spark_conf: dict[str, str] | None = None
ssl_verify: bool | None = None
ssl_ca_bundle: str | None = None
auth_type: str = "none"
auth_user: str | None = None
auth_password: str | None = None
auth_principal: str | None = None
auth_keytab: str | None = None
name: str = Field(
...,
description=(
"Unique connection name. Referenced by prepare_submit_job.connection "
"and get_external_*.connection_name. Saving with an existing name "
"overwrites that record."
),
)
master: str = Field(
...,
description=(
"Spark master URL. **Required** at the MCP layer (even though "
"the underlying Connection model has a default of 'yarn'). "
"Other valid values: 'yarn' (default for YARN), 'spark://host:port' "
"(Standalone), 'k8s://...', 'mesos://...', 'local[N]' or 'local[*]'. "
"The validator rejects 'yarn-cluster' and bare 'http://...' URLs — "
"those are common typos. The literal 'yarn' (not 'yarn-cluster') is "
"what spark-submit wants for --master."
),
)
deploy_mode: str = Field(
default="cluster",
description=(
"Spark deploy mode. 'cluster' (default, driver runs in YARN) or "
"'client' (driver runs where spark-submit is invoked). Most YARN "
"production submissions use 'cluster'."
),
)
yarn_rm_url: str | None = Field(
default=None,
description=(
"YARN ResourceManager REST base URL, e.g. 'http://rm-host:8088'. "
"**Required** for the 3 get_external_* tools to query YARN "
"directly. Optional for prepare_submit_job — spark-submit "
"discovers the RM via the cluster config when this is unset."
),
)
spark_conf: dict[str, str] | None = Field(
default=None,
description=(
"Dict of Spark conf key→value pairs, passed as --conf flags to "
"spark-submit. Example: {'spark.executor.memory': '4g', "
"'spark.sql.shuffle.partitions': '200'}. None or empty means "
"no extra --conf flags."
),
)
ssl_verify: bool | None = Field(
default=None,
description=(
"Whether to verify the YARN RM TLS certificate. None (default) "
"falls back to the global setting; explicit True/False overrides "
"the global default for this connection. Set False only for "
"self-signed dev clusters."
),
)
ssl_ca_bundle: str | None = Field(
default=None,
description=(
"Absolute path to a CA bundle file for YARN RM TLS verification. "
"Only relevant when the RM uses a private CA. Ignored when "
"ssl_verify=False."
),
)
auth_type: str = Field(
default="none",
description=(
"Authentication mode for YARN REST calls. One of: 'none' "
"(default, no auth header), 'simple' (pseudo-auth, "
"auth_user required), 'basic' (HTTP Basic, auth_user + "
"auth_password required), 'kerberos' (SPNEGO via the system "
"ticket cache — run kinit beforehand)."
),
)
auth_user: str | None = Field(
default=None,
description=(
"Username for auth_type='simple' or 'basic'. Ignored when "
"auth_type='none' or 'kerberos'."
),
)
auth_password: str | None = Field(
default=None,
description=(
"Password for auth_type='basic'. Ignored otherwise. Sent on "
"every YARN REST request — store with care."
),
)
auth_principal: str | None = Field(
default=None,
description=(
"Kerberos principal (e.g. 'user@REALM'). Display/audit only; "
"the actual SPNEGO handshake uses the system ticket cache. "
"Run `kinit <principal>` on the host before invoking the tools."
),
)
auth_keytab: str | None = Field(
default=None,
description=(
"Absolute path to a Kerberos keytab file. Optional convenience "
"for 'kinit -kt' workflows. The service does NOT auto-initialize "
"from the keytab — you must `kinit -kt <auth_keytab> <auth_principal>` "
"yourself before calling the tools."
),
)
url_allowlist: list[str] | None = Field(
default=None,
description=(
"Optional list of fnmatch glob patterns for hosts the fetch_url tool "
"may access. See Connection.url_allowlist for full semantics. "
"Example for single-label host clusters: ['ccam*'] allows any host "
"starting with 'ccam' (ccam1, ccam2, ..., ccam99). If omitted, None, "
"or empty, the saved connection will have url_allowlist=[] (the "
"default), meaning fetch_url will reject every URL until the list is "
"populated via update_connection."
),
)
class PrepareSubmitJobRequest(BaseModel):
connection: str
connection: str = Field(
...,
description=(
"Name of a saved Connection (call save_connection first, or "
"list_connections to see available names). The Connection's "
"master / deploy_mode / spark_conf / yarn_rm_url are snapshotted "
"into the pending submission at prepare time, so editing the "
"Connection afterwards does NOT retarget this pending job."
),
)
app_name: str = Field(
...,
description="Human-readable application name for tracking the pending submission.",
@@ -72,34 +185,115 @@ class PrepareSubmitJobRequest(BaseModel):
class PendingIdRequest(BaseModel):
pending_id: str
pending_id: str = Field(
...,
description=(
"ID of a pending submission. Format: 'p_' + 12 hex chars, "
"e.g. 'p_a1b2c3d4e5f6'. Returned by prepare_submit_job; "
"visible via list_pending_jobs."
),
)
class UpdatePendingJobRequest(BaseModel):
pending_id: str
pending_id: str = Field(
...,
description=(
"ID of the pending submission to modify. Format: 'p_' + 12 hex "
"chars, e.g. 'p_a1b2c3d4e5f6'. Only PENDING submissions can "
"be updated — once SUBMITTED, CANCELLED, or FAILED, the "
"pending is terminal."
),
)
script_path: str | None = Field(
default=None,
description="Optional new absolute path to the PySpark script. If provided, the file must exist and pass SQL guard.",
description=(
"Optional new absolute path to the PySpark script. If provided, "
"the file must exist and pass the SQL guard. Omit to keep the "
"current script_path."
),
)
queue: str | None = Field(
default=None,
description=(
"New YARN queue name. Omit to keep the current value (PATCH "
"semantics: only fields you provide are changed)."
),
)
executor_memory: str | None = Field(
default=None,
description=(
"New executor memory, e.g. '4G'. Omit to keep the current value."
),
)
executor_cores: int | None = Field(
default=None,
description="New cores per executor. Omit to keep the current value.",
)
num_executors: int | None = Field(
default=None,
description="New total executor count. Omit to keep the current value.",
)
app_name: str | None = Field(
default=None,
description=(
"New human-readable application name (visible in YARN UI). "
"Omit to keep the current value."
),
)
extra_args: dict[str, str] | None = Field(
default=None,
description=(
"Replacement dict of additional spark-submit flags (e.g. "
"{'jars': '/path/to.jar'}). Unlike the scalar fields, providing "
"this REPLACES the entire dict — it is not deep-merged. Omit to "
"keep the current value."
),
)
queue: str | None = None
executor_memory: str | None = None
executor_cores: int | None = None
num_executors: int | None = None
app_name: str | None = None
extra_args: dict[str, str] | None = None
class JobIdRequest(BaseModel):
job_id: str
job_id: str = Field(
...,
description=(
"Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') or "
"the YARN application_id (e.g. 'application_17400000001_0001') of "
"a job submitted through this service. The lookup tries job_id "
"first, then application_id. For YARN applications not submitted "
"here, use the get_external_* tools instead."
),
)
class GetJobLogsRequest(BaseModel):
job_id: str
tail_chars: int = 5000
job_id: str = Field(
...,
description=(
"Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') "
"or the YARN application_id (e.g. 'application_17400000001_0001') "
"of a job submitted through this service. For YARN applications "
"not submitted here, use get_external_job_logs instead."
),
)
tail_chars: int = Field(
default=5000,
description=(
"Return only the last N characters of the aggregated container "
"logs. Default 5000. Use a larger value if the head of the log "
"(stack traces, driver errors) is being truncated."
),
)
class ConnectionNameRequest(BaseModel):
name: str
name: str = Field(
...,
description=(
"Name of a saved Connection. Use list_connections to see "
"available names. Saving with this name updates an existing "
"record (see save_connection)."
),
)
class WriteJobFileRequest(BaseModel):
@@ -142,3 +336,180 @@ class UpdateJobFileRequest(BaseModel):
"Maximum 1 MB to keep the MCP response bounded."
),
)
class ExternalJobLogsRequest(BaseModel):
application_id: str = Field(
...,
description=(
"YARN application_id of the external job, format "
"'application_<14-digit-timestamp>_<sequence>' (e.g. "
"'application_1740000000001_0001'). This tool is for jobs "
"NOT submitted through this MCP service — for those, use "
"get_job_logs(job_id=...) instead."
),
)
connection_name: str = Field(
...,
description=(
"Name of a saved Connection (see list_connections) pointing "
"at the YARN cluster where the application ran."
),
)
tail_chars: int = Field(
default=5000,
description="Return only the last N characters of the aggregated container logs.",
)
class ExternalJobStatusRequest(BaseModel):
application_id: str = Field(
...,
description=(
"YARN application_id of the external job, format "
"'application_<14-digit-timestamp>_<sequence>'. Use "
"get_job_status(job_id=...) for jobs submitted through this service."
),
)
connection_name: str = Field(
...,
description="Name of a saved Connection pointing at the YARN cluster.",
)
class ExternalJobResultRequest(BaseModel):
application_id: str = Field(
...,
description=(
"YARN application_id of the external job, format "
"'application_<14-digit-timestamp>_<sequence>'. Use "
"get_job_result(job_id=...) for jobs submitted through this service."
),
)
connection_name: str = Field(
...,
description="Name of a saved Connection pointing at the YARN cluster.",
)
class FetchUrlRequest(BaseModel):
url: str = Field(
...,
description=(
"Absolute URL to fetch. The only access control is the named "
"Connection's url_allowlist: the URL host must match one of the "
"fnmatch glob patterns in that list. An empty or omitted allowlist "
"denies every host. IP literals and non-HTTP schemes are allowed "
"if and only if they are matched by the allowlist. The Connection's "
"saved auth is reused for the outbound request — the agent does not "
"need cluster credentials."
),
)
connection_name: str = Field(
...,
description=(
"Name of a saved Connection (see list_connections). The "
"Connection's 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}")
job = store.get_either(job_id)
if job is None:
raise _unknown_job_error(job_id)
raise _unknown_job_error(
job_id,
external_tool_hint="get_external_job_result(application_id, connection_name)",
)
conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
+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}")
job = store.get_either(job_id)
if job is None:
raise _unknown_job_error(job_id)
raise _unknown_job_error(
job_id,
external_tool_hint="get_external_job_status(application_id, connection_name)",
)
conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
+4 -5
View File
@@ -6,9 +6,8 @@
import os
import secrets
import uuid
from datetime import datetime
from datetime import datetime, timezone
from common.config import settings
from common.logging import logger
from common.sql_guard import validate_pyspark_code
from spark_executor.core.connection_store import store as conn_store
@@ -124,7 +123,7 @@ def prepare_submit_job(
num_executors=num_executors,
spark_conf=dict(conn.spark_conf),
extra_args=dict(extra_args or {}),
created_at=datetime.utcnow(),
created_at=datetime.now(timezone.utc),
status="PENDING",
)
pending_store.save(pending)
@@ -248,7 +247,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
application_id=application_id,
script_path=pending.script_path,
queue=pending.queue,
submit_time=datetime.utcnow(),
submit_time=datetime.now(timezone.utc),
connection=pending.connection,
yarn_rm_url=pending.yarn_rm_url,
)
@@ -282,7 +281,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult:
application_id=application_id,
script_path=pending.script_path,
queue=pending.queue,
submit_time=datetime.utcnow(),
submit_time=datetime.now(timezone.utc),
connection=pending.connection,
yarn_rm_url=pending.yarn_rm_url,
)
+13
View File
@@ -64,6 +64,19 @@ def test_seventeen_tool_routes_registered():
assert "/update_pending_job" in paths
def test_twenty_three_tool_routes_registered():
paths = {r.path for r in app.routes}
for path in (
"/get_external_job_logs",
"/get_external_job_status",
"/get_external_job_result",
"/list_applications",
"/fetch_url",
"/update_connection",
):
assert path in paths, f"missing MCP tool route: {path}"
# --- operation_id: pin clean MCP tool names (no auto-generated suffixes) ---
#
# fastapi-mcp uses each route's OpenAPI `operationId` as the MCP tool name
+75
View File
@@ -1,4 +1,5 @@
# coding=utf-8
import json
from pathlib import Path
import pytest
@@ -120,3 +121,77 @@ def test_delete_connection_returns_status(_fresh_store):
def test_delete_connection_unknown_raises(_fresh_store):
with pytest.raises(KeyError):
connections.delete_connection("missing")
def test_save_connection_preserves_url_allowlist_on_existing_record(_fresh_store):
connections.save_connection(name="prod", master="yarn", url_allowlist=["ccam*"])
out = connections.save_connection(name="prod", master="spark://new:7077")
assert out["master"] == "spark://new:7077"
assert out["url_allowlist"] == ["ccam*"]
def test_update_connection_changes_specified_field(_fresh_store):
connections.save_connection(name="prod", master="yarn")
out = connections.update_connection(name="prod", master="spark://new:7077")
assert out["master"] == "spark://new:7077"
assert _fresh_store.get("prod").master == "spark://new:7077"
def test_update_connection_keeps_omitted_fields(_fresh_store):
connections.save_connection(
name="prod",
master="yarn",
deploy_mode="client",
yarn_rm_url="http://rm:8088",
)
out = connections.update_connection(name="prod", deploy_mode="cluster")
assert out["deploy_mode"] == "cluster"
assert out["master"] == "yarn"
assert out["yarn_rm_url"] == "http://rm:8088"
def test_update_connection_with_no_fields_is_noop(_fresh_store):
connections.save_connection(name="prod", master="yarn", deploy_mode="client")
out = connections.update_connection(name="prod")
assert out["master"] == "yarn"
assert out["deploy_mode"] == "client"
def test_update_connection_changes_url_allowlist(_fresh_store):
connections.save_connection(name="prod", master="yarn")
out = connections.update_connection(name="prod", url_allowlist=["ccam*"])
assert out["url_allowlist"] == ["ccam*"]
def test_update_connection_replaces_not_merges_dict(_fresh_store):
connections.save_connection(name="prod", master="yarn", spark_conf={"a": "1"})
out = connections.update_connection(name="prod", spark_conf={"b": "2"})
assert out["spark_conf"] == {"b": "2"}
def test_update_connection_replaces_not_merges_url_allowlist(_fresh_store):
connections.save_connection(
name="prod",
master="yarn",
url_allowlist=["ccam*"],
)
out = connections.update_connection(name="prod", url_allowlist=["nm*"])
assert out["url_allowlist"] == ["nm*"]
def test_update_connection_raises_for_unknown_name(_fresh_store):
with pytest.raises(KeyError, match="Connection not found"):
connections.update_connection(name="missing", master="yarn")
def test_update_connection_validates_patched_connection(_fresh_store):
connections.save_connection(name="prod", master="yarn")
with pytest.raises(ValueError, match="master must be"):
connections.update_connection(name="prod", master="http://bad")
def test_update_connection_persists_to_disk(_fresh_store, tmp_path):
connections.save_connection(name="prod", master="yarn")
connections.update_connection(name="prod", master="spark://new:7077")
raw = json.loads((tmp_path / "connections.json").read_text())
assert raw["prod"]["master"] == "spark://new:7077"
+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")
def test_kill_job_raises_400_for_external_application_id(fresh_stores):
"""Input that looks like a YARN application_id but is not in the local
JobStore must raise ValueError (-> 400) — kill_job has no external
equivalent, so the error points to the YARN CLI / UI."""
with pytest.raises(ValueError, match="YARN CLI"):
kill.kill_job("application_17400000001_0001")
def test_kill_job_raises_when_connection_missing():
_fresh_stores()
kill.store.put(
+8
View File
@@ -68,6 +68,14 @@ def test_get_job_logs_raises_for_unknown_job():
logs.get_job_logs("missing")
def test_get_job_logs_raises_400_for_external_application_id(fresh_stores):
"""Input that looks like a YARN application_id but is not in the local
JobStore must raise ValueError (-> 400) with a hint to use the
external tool, NOT a generic KeyError (-> 404)."""
with pytest.raises(ValueError, match="get_external_job_logs"):
logs.get_job_logs("application_17400000001_0001")
def test_get_job_logs_raises_when_connection_missing(fresh_stores):
logs.store.put(
Job(
+8
View File
@@ -136,6 +136,14 @@ def test_result_raises_keyerror_for_unknown_job():
assert "application_id" in msg
def test_result_raises_400_for_external_application_id(fresh_stores):
"""Input that looks like a YARN application_id but is not in the local
JobStore must raise ValueError (-> 400) with a hint to use the
external tool, NOT a generic KeyError (-> 404)."""
with pytest.raises(ValueError, match="get_external_job_result"):
result.get_job_result("application_17400000001_0001")
def test_result_raises_when_connection_missing():
_fresh_stores()
result.store.put(
+8
View File
@@ -66,6 +66,14 @@ def test_get_job_status_raises_for_unknown_job():
status.get_job_status("missing")
def test_get_job_status_raises_400_for_external_application_id(fresh_stores):
"""Input that looks like a YARN application_id but is not in the local
JobStore must raise ValueError (-> 400) with a hint to use the
external tool, NOT a generic KeyError (-> 404)."""
with pytest.raises(ValueError, match="get_external_job_status"):
status.get_job_status("application_17400000001_0001")
def test_get_job_status_raises_when_connection_missing():
_fresh_stores()
status.store.put(