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

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.
2026-06-26 13:57:57 +08:00
parent ad557b5984
commit 10f075ab7f
9 changed files with 256 additions and 8 deletions
@@ -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}")
@@ -53,6 +53,13 @@ class Connection(BaseModel):
    # overrides the global default for this connection.
    ssl_verify: bool | None = None
    ssl_ca_bundle: str | None = None
+    # Authentication for YARN REST calls.
+    auth_type: str = "none"  # "none" | "simple" | "basic" | "kerberos"
+    auth_user: str | None = None
+    auth_password: str | None = None
+    # Display/audit only for kerberos; actual SPNEGO uses the system cache.
+    auth_principal: str | None = None
+    auth_keytab: str | None = None

    @field_validator("master")
    @classmethod
@@ -67,6 +74,15 @@ class Connection(BaseModel):
            f"or 'local[/N]'; got {v!r}"
        )

+    @field_validator("auth_type")
+    @classmethod
+    def _check_auth_type(cls, v: str) -> str:
+        if v not in {"none", "simple", "basic", "kerberos"}:
+            raise ValueError(
+                f"auth_type must be one of none/simple/basic/kerberos; got {v!r}"
+            )
+        return v
+

 class PendingSubmission(BaseModel):
    pending_id: str
@@ -17,6 +17,11 @@ def save_connection(
    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]:
    logger.debug(
        f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
@@ -30,6 +35,11 @@ def save_connection(
        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,
    )
    store.save(conn)
    return {"name": name, "status": "SAVED"}
@@ -24,6 +24,11 @@ class SaveConnectionRequest(BaseModel):
    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


 class PrepareSubmitJobRequest(BaseModel):