tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(yarn_client): SSL/TLS config for YARN REST connections

Browse Source 
Add per-Connection ssl_verify / ssl_ca_bundle plus global defaults so
CDH 5 / on-prem clusters with self-signed certs or custom CA bundles can
be queried without patching code.

- Connection gets ssl_verify (bool|None) and ssl_ca_bundle (str|None)
- Settings gets ssl_verify_default and ssl_ca_bundle_default
- New YarnClientConfig dataclass carries the resolved verify= value
- _request passes verify= through to httpx.request
- All public yarn_client functions now take YarnClientConfig instead of
  a bare yarn_rm_url string; tool call sites resolve the Connection
- SaveConnectionRequest exposes the two new fields

Tests cover per-connection CA bundle, per-connection verify=False,
global default fallback, and connection-not-found error.
This commit is contained in:
Claude
2026-06-26 13:50:30 +08:00
parent 5566d55a7c
commit ad557b5984
16 changed files with 362 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/core/yarn_client.py
+57 -20
View File
@@ -18,11 +18,13 @@ var if unset. This matches the pattern the original Connection.yarn_rm_url
field was designed for, but no longer requires the `yarn` CLI to interpret it.
"""
import json
from dataclasses import dataclass
import httpx
from common.config import settings
from common.logging import logger
from spark_executor.models import Connection
class YarnError(Exception):
@@ -33,6 +35,42 @@ class YarnConfigError(YarnError):
"""Raised when the YARN ResourceManager URL is missing or malformed."""
@dataclass(frozen=True)
class YarnClientConfig:
"""Resolved per-connection + global-default configuration for YARN REST calls.
Carries the ResourceManager URL and the SSL/TLS verification value that
should be passed to httpx.request. Future authentication fields
(basic/kerberos/SPNEGO) can be added here without renaming the class.
"""
yarn_rm_url: str | None
ssl_verify: bool | None = None
ssl_ca_bundle: str | None = None
@classmethod
def from_connection(cls, conn: Connection) -> "YarnClientConfig":
return cls(
yarn_rm_url=conn.yarn_rm_url,
ssl_verify=conn.ssl_verify,
ssl_ca_bundle=conn.ssl_ca_bundle,
)
def verify_for_httpx(self) -> bool | str:
"""Return the value for httpx.request(verify=...).
Per-connection CA bundle wins over the global default. If neither is
set, fall back to the per-connection ssl_verify flag, then to the
global ssl_verify_default.
"""
ca = self.ssl_ca_bundle or settings.ssl_ca_bundle_default
if ca:
return ca
if self.ssl_verify is not None:
return self.ssl_verify
return settings.ssl_verify_default
def _base_url(yarn_rm_url: str | None) -> str:
"""Resolve and validate the RM URL. Raises YarnConfigError if unusable."""
url = yarn_rm_url or settings.yarn_resource_manager_url
@@ -51,10 +89,10 @@ def _base_url(yarn_rm_url: str | None) -> str:
def _request(method: str, url: str, *, json_body: dict | None = None,
timeout: float = 30.0) -> httpx.Response:
timeout: float = 30.0, verify: bool | str = True) -> httpx.Response:
logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
try:
resp = httpx.request(method, url, json=json_body, timeout=timeout)
resp = httpx.request(method, url, json=json_body, timeout=timeout, verify=verify)
except httpx.HTTPError as exc:
logger.error(f"YARN {method} {url} failed: {exc}")
raise YarnError(f"YARN connection failed: {exc}") from exc
@@ -65,10 +103,10 @@ def _request(method: str, url: str, *, json_body: dict | None = None,
return resp
def get_application_status(application_id: str, yarn_rm_url: str | None) -> tuple[str, str]:
def get_application_status(application_id: str, config: YarnClientConfig) -> tuple[str, str]:
"""Return (state, raw_json_text) for an application, or raise YarnError."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
resp = _request("GET", url)
url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
resp = _request("GET", url, verify=config.verify_for_httpx())
if resp.status_code == 404:
raise YarnError(f"YARN application {application_id!r} not found")
if resp.status_code >= 400:
@@ -92,17 +130,17 @@ def _logs_unavailable_error(application_id: str) -> YarnError:
)
def _fetch_logs_via_nodemanager(application_id: str, yarn_rm_url: str | None) -> str:
def _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -> str:
"""
Hadoop 2.x fallback: walk app attempts -> containers -> NodeManager
container logs and concatenate the results.
"""
base = _base_url(yarn_rm_url)
base = _base_url(config.yarn_rm_url)
app_url = f"{base}/ws/v1/cluster/apps/{application_id}"
# Latest (or all) app attempts.
attempts_url = f"{app_url}/appattempts"
resp = _request("GET", attempts_url, timeout=30.0)
resp = _request("GET", attempts_url, timeout=30.0, verify=config.verify_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
attempts = resp.json().get("appAttempts", {}).get("appAttempt", [])
@@ -116,7 +154,7 @@ def _fetch_logs_via_nodemanager(application_id: str, yarn_rm_url: str | None) ->
if not attempt_id:
continue
containers_url = f"{app_url}/appattempts/{attempt_id}/containers"
resp = _request("GET", containers_url, timeout=30.0)
resp = _request("GET", containers_url, timeout=30.0, verify=config.verify_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
containers.extend(resp.json().get("containers", {}).get("container", []))
@@ -131,10 +169,9 @@ def _fetch_logs_via_nodemanager(application_id: str, yarn_rm_url: str | None) ->
node_http_address = container.get("nodeHttpAddress")
if not all((container_id, user, node_http_address)):
continue
nm_url = (
f"http://{node_http_address}/node/containerlogs/{container_id}/{user}/"
)
resp = _request("GET", nm_url, timeout=60.0)
scheme = "https" if config.yarn_rm_url and config.yarn_rm_url.startswith("https://") else "http"
nm_url = f"{scheme}://{node_http_address}/node/containerlogs/{container_id}/{user}/"
resp = _request("GET", nm_url, timeout=60.0, verify=config.verify_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
parts.append(f"=== container: {container_id} ===\n{resp.text}")
@@ -144,12 +181,12 @@ def _fetch_logs_via_nodemanager(application_id: str, yarn_rm_url: str | None) ->
return "\n\n".join(parts)
def get_application_logs(application_id: str, yarn_rm_url: str | None) -> str:
def get_application_logs(application_id: str, config: YarnClientConfig) -> str:
"""Return aggregated container logs for an application as text."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs"
resp = _request("GET", url, timeout=60.0)
url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs"
resp = _request("GET", url, timeout=60.0, verify=config.verify_for_httpx())
if resp.status_code in (404, 501):
return _fetch_logs_via_nodemanager(application_id, yarn_rm_url)
return _fetch_logs_via_nodemanager(application_id, config)
if resp.status_code >= 400:
logger.error(f"YARN GET {url} -> {resp.status_code}: {resp.text[:500]}")
raise YarnError(f"YARN GET logs returned HTTP {resp.status_code}")
@@ -157,10 +194,10 @@ def get_application_logs(application_id: str, yarn_rm_url: str | None) -> str:
return resp.text
def kill_application(application_id: str, yarn_rm_url: str | None) -> None:
def kill_application(application_id: str, config: YarnClientConfig) -> None:
"""PUT state=KILLED to /ws/v1/cluster/apps/{appid}/state."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state"
resp = _request("PUT", url, json_body={"state": "KILLED"})
url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state"
resp = _request("PUT", url, json_body={"state": "KILLED"}, verify=config.verify_for_httpx())
if resp.status_code >= 400:
logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")