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

feat(yarn_client): auth config for YARN REST (none/simple/basic/kerberos)

Browse Source 
Add per-Connection authentication so CDH 5 / Kerberos / HTTP Basic
clusters can be queried.

- auth_type: none / simple / basic / kerberos
- auth_user / auth_password for HTTP Basic
- auth_principal / auth_keytab stored for audit/display; actual SPNEGO
  handled by httpx-kerberos using the system Kerberos credential cache
- YarnClientConfig.auth_for_httpx() returns the right httpx.Auth object
- _request passes auth= through to httpx.request alongside verify=

New dependency: httpx-kerberos.

Tests cover none/simple (no auth object), Basic auth header,
Kerberos auth object, missing basic user, invalid auth_type, and
tool-layer config propagation.
This commit is contained in:
Claude
2026-06-26 13:57:57 +08:00
parent ad557b5984
commit 10f075ab7f
9 changed files with 256 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/core/yarn_client.py
+32 -8
View File
@@ -21,6 +21,7 @@ import json
from dataclasses import dataclass
import httpx
import httpx_kerberos
from common.config import settings
from common.logging import logger
@@ -47,6 +48,11 @@ class YarnClientConfig:
yarn_rm_url: str | 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 # display/audit only for kerberos
auth_keytab: str | None = None # display/audit only for kerberos
@classmethod
def from_connection(cls, conn: Connection) -> "YarnClientConfig":
@@ -54,6 +60,11 @@ class YarnClientConfig:
yarn_rm_url=conn.yarn_rm_url,
ssl_verify=conn.ssl_verify,
ssl_ca_bundle=conn.ssl_ca_bundle,
auth_type=conn.auth_type,
auth_user=conn.auth_user,
auth_password=conn.auth_password,
auth_principal=conn.auth_principal,
auth_keytab=conn.auth_keytab,
)
def verify_for_httpx(self) -> bool | str:
@@ -70,6 +81,18 @@ class YarnClientConfig:
return self.ssl_verify
return settings.ssl_verify_default
def auth_for_httpx(self) -> httpx.Auth | None:
"""Return the httpx.Auth object implied by this config, if any."""
if self.auth_type in ("none", "simple"):
return None
if self.auth_type == "basic":
if not self.auth_user:
raise YarnConfigError("auth_type='basic' requires auth_user")
return httpx.BasicAuth(self.auth_user, self.auth_password or "")
if self.auth_type == "kerberos":
return httpx_kerberos.HTTPKerberosAuth()
raise YarnConfigError(f"Unknown auth_type: {self.auth_type!r}")
def _base_url(yarn_rm_url: str | None) -> str:
"""Resolve and validate the RM URL. Raises YarnConfigError if unusable."""
@@ -89,10 +112,11 @@ 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) -> httpx.Response:
timeout: float = 30.0, verify: bool | str = True,
auth: httpx.Auth | None = None) -> 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, verify=verify)
resp = httpx.request(method, url, json=json_body, timeout=timeout, verify=verify, auth=auth)
except httpx.HTTPError as exc:
logger.error(f"YARN {method} {url} failed: {exc}")
raise YarnError(f"YARN connection failed: {exc}") from exc
@@ -106,7 +130,7 @@ def _request(method: str, url: str, *, json_body: dict | None = None,
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(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
resp = _request("GET", url, verify=config.verify_for_httpx())
resp = _request("GET", url, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
if resp.status_code == 404:
raise YarnError(f"YARN application {application_id!r} not found")
if resp.status_code >= 400:
@@ -140,7 +164,7 @@ def _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -
# Latest (or all) app attempts.
attempts_url = f"{app_url}/appattempts"
resp = _request("GET", attempts_url, timeout=30.0, verify=config.verify_for_httpx())
resp = _request("GET", attempts_url, timeout=30.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
attempts = resp.json().get("appAttempts", {}).get("appAttempt", [])
@@ -154,7 +178,7 @@ def _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -
if not attempt_id:
continue
containers_url = f"{app_url}/appattempts/{attempt_id}/containers"
resp = _request("GET", containers_url, timeout=30.0, verify=config.verify_for_httpx())
resp = _request("GET", containers_url, timeout=30.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
containers.extend(resp.json().get("containers", {}).get("container", []))
@@ -171,7 +195,7 @@ def _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -
continue
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())
resp = _request("GET", nm_url, timeout=60.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
if resp.status_code >= 400:
raise _logs_unavailable_error(application_id)
parts.append(f"=== container: {container_id} ===\n{resp.text}")
@@ -184,7 +208,7 @@ def _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -
def get_application_logs(application_id: str, config: YarnClientConfig) -> str:
"""Return aggregated container logs for an application as text."""
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())
resp = _request("GET", url, timeout=60.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
if resp.status_code in (404, 501):
return _fetch_logs_via_nodemanager(application_id, config)
if resp.status_code >= 400:
@@ -197,7 +221,7 @@ def get_application_logs(application_id: str, config: YarnClientConfig) -> str:
def kill_application(application_id: str, config: YarnClientConfig) -> None:
"""PUT state=KILLED to /ws/v1/cluster/apps/{appid}/state."""
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())
resp = _request("PUT", url, json_body={"state": "KILLED"}, verify=config.verify_for_httpx(), auth=config.auth_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}")