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

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
+13
View File
@@ -76,6 +76,10 @@ class Settings:
# The binary is resolved against $PATH (set by docker-entrypoint.sh). # The binary is resolved against $PATH (set by docker-entrypoint.sh).
spark_submit_bin: str = "spark-submit" spark_submit_bin: str = "spark-submit"
# --- SSL/TLS defaults for YARN REST calls ---
ssl_verify_default: bool = True
ssl_ca_bundle_default: str | None = None
@classmethod @classmethod
def from_env(cls) -> "Settings": def from_env(cls) -> "Settings":
data_dir = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") data_dir = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data")
@@ -95,6 +99,13 @@ class Settings:
spark_submit_bin=os.environ.get( spark_submit_bin=os.environ.get(
"SPARK_EXECUTOR_SPARK_SUBMIT_BIN", "spark-submit" "SPARK_EXECUTOR_SPARK_SUBMIT_BIN", "spark-submit"
), ),
ssl_verify_default=(
os.environ.get("SPARK_EXECUTOR_SSL_VERIFY_DEFAULT", "true").lower()
not in ("false", "0", "no", "off")
),
ssl_ca_bundle_default=(
os.environ.get("SPARK_EXECUTOR_SSL_CA_BUNDLE_DEFAULT") or None
),
) )
def reload(self) -> "Settings": def reload(self) -> "Settings":
@@ -110,6 +121,8 @@ class Settings:
self.yarn_resource_manager_url = fresh.yarn_resource_manager_url self.yarn_resource_manager_url = fresh.yarn_resource_manager_url
self.log_level = fresh.log_level self.log_level = fresh.log_level
self.spark_submit_bin = fresh.spark_submit_bin self.spark_submit_bin = fresh.spark_submit_bin
self.ssl_verify_default = fresh.ssl_verify_default
self.ssl_ca_bundle_default = fresh.ssl_ca_bundle_default
return self return self
+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. field was designed for, but no longer requires the `yarn` CLI to interpret it.
""" """
import json import json
from dataclasses import dataclass
import httpx import httpx
from common.config import settings from common.config import settings
from common.logging import logger from common.logging import logger
from spark_executor.models import Connection
class YarnError(Exception): class YarnError(Exception):
@@ -33,6 +35,42 @@ class YarnConfigError(YarnError):
"""Raised when the YARN ResourceManager URL is missing or malformed.""" """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: def _base_url(yarn_rm_url: str | None) -> str:
"""Resolve and validate the RM URL. Raises YarnConfigError if unusable.""" """Resolve and validate the RM URL. Raises YarnConfigError if unusable."""
url = yarn_rm_url or settings.yarn_resource_manager_url 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, 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 "")) logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
try: 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: except httpx.HTTPError as exc:
logger.error(f"YARN {method} {url} failed: {exc}") logger.error(f"YARN {method} {url} failed: {exc}")
raise YarnError(f"YARN connection failed: {exc}") from 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 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.""" """Return (state, raw_json_text) for an application, or raise YarnError."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}" url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
resp = _request("GET", url) resp = _request("GET", url, verify=config.verify_for_httpx())
if resp.status_code == 404: if resp.status_code == 404:
raise YarnError(f"YARN application {application_id!r} not found") raise YarnError(f"YARN application {application_id!r} not found")
if resp.status_code >= 400: 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 Hadoop 2.x fallback: walk app attempts -> containers -> NodeManager
container logs and concatenate the results. 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}" app_url = f"{base}/ws/v1/cluster/apps/{application_id}"
# Latest (or all) app attempts. # Latest (or all) app attempts.
attempts_url = f"{app_url}/appattempts" 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: if resp.status_code >= 400:
raise _logs_unavailable_error(application_id) raise _logs_unavailable_error(application_id)
attempts = resp.json().get("appAttempts", {}).get("appAttempt", []) 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: if not attempt_id:
continue continue
containers_url = f"{app_url}/appattempts/{attempt_id}/containers" 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: if resp.status_code >= 400:
raise _logs_unavailable_error(application_id) raise _logs_unavailable_error(application_id)
containers.extend(resp.json().get("containers", {}).get("container", [])) 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") node_http_address = container.get("nodeHttpAddress")
if not all((container_id, user, node_http_address)): if not all((container_id, user, node_http_address)):
continue continue
nm_url = ( scheme = "https" if config.yarn_rm_url and config.yarn_rm_url.startswith("https://") else "http"
f"http://{node_http_address}/node/containerlogs/{container_id}/{user}/" 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)
if resp.status_code >= 400: if resp.status_code >= 400:
raise _logs_unavailable_error(application_id) raise _logs_unavailable_error(application_id)
parts.append(f"=== container: {container_id} ===\n{resp.text}") 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) 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.""" """Return aggregated container logs for an application as text."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs" url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs"
resp = _request("GET", url, timeout=60.0) resp = _request("GET", url, timeout=60.0, verify=config.verify_for_httpx())
if resp.status_code in (404, 501): 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: if resp.status_code >= 400:
logger.error(f"YARN GET {url} -> {resp.status_code}: {resp.text[:500]}") logger.error(f"YARN GET {url} -> {resp.status_code}: {resp.text[:500]}")
raise YarnError(f"YARN GET logs returned HTTP {resp.status_code}") 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 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.""" """PUT state=KILLED to /ws/v1/cluster/apps/{appid}/state."""
url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state" url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state"
resp = _request("PUT", url, json_body={"state": "KILLED"}) resp = _request("PUT", url, json_body={"state": "KILLED"}, verify=config.verify_for_httpx())
if resp.status_code >= 400: if resp.status_code >= 400:
logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}") logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}") raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
+4
View File
@@ -49,6 +49,10 @@ class Connection(BaseModel):
deploy_mode: str = "cluster" deploy_mode: str = "cluster"
yarn_rm_url: str | None = None yarn_rm_url: str | None = None
spark_conf: dict[str, str] = Field(default_factory=dict) spark_conf: dict[str, str] = Field(default_factory=dict)
# None means "fall back to Settings.ssl_verify_default". Explicit True/False
# overrides the global default for this connection.
ssl_verify: bool | None = None
ssl_ca_bundle: str | None = None
@field_validator("master") @field_validator("master")
@classmethod @classmethod
+4
View File
@@ -15,6 +15,8 @@ def save_connection(
deploy_mode: str = "cluster", deploy_mode: str = "cluster",
yarn_rm_url: str | None = None, yarn_rm_url: str | None = None,
spark_conf: dict[str, str] | None = None, spark_conf: dict[str, str] | None = None,
ssl_verify: bool | None = None,
ssl_ca_bundle: str | None = None,
) -> dict[str, str]: ) -> dict[str, str]:
logger.debug( logger.debug(
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
@@ -26,6 +28,8 @@ def save_connection(
deploy_mode=deploy_mode, deploy_mode=deploy_mode,
yarn_rm_url=yarn_rm_url, yarn_rm_url=yarn_rm_url,
spark_conf=spark_conf or {}, spark_conf=spark_conf or {},
ssl_verify=ssl_verify,
ssl_ca_bundle=ssl_ca_bundle,
) )
store.save(conn) store.save(conn)
return {"name": name, "status": "SAVED"} return {"name": name, "status": "SAVED"}
+7 -2
View File
@@ -5,7 +5,8 @@
""" """
from common.logging import logger from common.logging import logger
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import kill_application from spark_executor.core.yarn_client import YarnClientConfig, kill_application
from spark_executor.tools.connections import store as conn_store
store = JobStore() store = JobStore()
@@ -15,7 +16,11 @@ def kill_job(job_id: str) -> dict[str, str]:
job = store.get(job_id) job = store.get(job_id)
if job is None: if job is None:
raise KeyError(f"Unknown job_id: {job_id}") raise KeyError(f"Unknown job_id: {job_id}")
kill_application(job.application_id, job.yarn_rm_url) conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
config = YarnClientConfig.from_connection(conn)
kill_application(job.application_id, config)
logger.info(f"kill_job ok job_id={job_id} application_id={job.application_id}") logger.info(f"kill_job ok job_id={job_id} application_id={job.application_id}")
return { return {
"job_id": job_id, "job_id": job_id,
+7 -2
View File
@@ -5,7 +5,8 @@
""" """
from common.logging import logger from common.logging import logger
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import get_application_logs from spark_executor.core.yarn_client import YarnClientConfig, get_application_logs
from spark_executor.tools.connections import store as conn_store
store = JobStore() store = JobStore()
@@ -15,7 +16,11 @@ def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
job = store.get(job_id) job = store.get(job_id)
if job is None: if job is None:
raise KeyError(f"Unknown job_id: {job_id}") raise KeyError(f"Unknown job_id: {job_id}")
full = get_application_logs(job.application_id, job.yarn_rm_url) conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
config = YarnClientConfig.from_connection(conn)
full = get_application_logs(job.application_id, config)
tailed = full[-tail_chars:] if len(full) > tail_chars else full tailed = full[-tail_chars:] if len(full) > tail_chars else full
logger.info( logger.info(
f"get_job_logs ok job_id={job_id} application_id={job.application_id} " f"get_job_logs ok job_id={job_id} application_id={job.application_id} "
+2
View File
@@ -22,6 +22,8 @@ class SaveConnectionRequest(BaseModel):
deploy_mode: str = "cluster" deploy_mode: str = "cluster"
yarn_rm_url: str | None = None yarn_rm_url: str | None = None
spark_conf: dict[str, str] | None = None spark_conf: dict[str, str] | None = None
ssl_verify: bool | None = None
ssl_ca_bundle: str | None = None
class PrepareSubmitJobRequest(BaseModel): class PrepareSubmitJobRequest(BaseModel):
+7 -2
View File
@@ -7,7 +7,8 @@ import json
from common.logging import logger from common.logging import logger
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import get_application_status from spark_executor.core.yarn_client import YarnClientConfig, get_application_status
from spark_executor.tools.connections import store as conn_store
from spark_executor.models import JobResult from spark_executor.models import JobResult
store = JobStore() store = JobStore()
@@ -18,7 +19,11 @@ def get_job_result(job_id: str) -> JobResult:
job = store.get(job_id) job = store.get(job_id)
if job is None: if job is None:
raise KeyError(f"Unknown job_id: {job_id}") raise KeyError(f"Unknown job_id: {job_id}")
state, raw = get_application_status(job.application_id, job.yarn_rm_url) conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
config = YarnClientConfig.from_connection(conn)
state, raw = get_application_status(job.application_id, config)
app = json.loads(raw).get("app", {}) app = json.loads(raw).get("app", {})
result = JobResult( result = JobResult(
application_id=job.application_id, application_id=job.application_id,
+7 -1
View File
@@ -4,9 +4,11 @@
@Author :tao.chen @Author :tao.chen
""" """
from common.logging import logger from common.logging import logger
from spark_executor.core.yarn_client import YarnClientConfig
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import get_application_status from spark_executor.core.yarn_client import get_application_status
from spark_executor.models import JobStatus from spark_executor.models import JobStatus
from spark_executor.tools.connections import store as conn_store
store = JobStore() store = JobStore()
@@ -16,6 +18,10 @@ def get_job_status(job_id: str) -> JobStatus:
job = store.get(job_id) job = store.get(job_id)
if job is None: if job is None:
raise KeyError(f"Unknown job_id: {job_id}") raise KeyError(f"Unknown job_id: {job_id}")
state, raw = get_application_status(job.application_id, job.yarn_rm_url) conn = conn_store.get(job.connection)
if conn is None:
raise KeyError(f"Connection not found: {job.connection}")
config = YarnClientConfig.from_connection(conn)
state, raw = get_application_status(job.application_id, config)
logger.info(f"get_job_status ok job_id={job_id} application_id={job.application_id} state={state}") logger.info(f"get_job_status ok job_id={job_id} application_id={job.application_id} state={state}")
return JobStatus(application_id=job.application_id, state=state, raw=raw) return JobStatus(application_id=job.application_id, state=state, raw=raw)
+14
View File
@@ -65,6 +65,20 @@ def test_get_connection_returns_dict(_fresh_store):
assert out["deploy_mode"] == "cluster" assert out["deploy_mode"] == "cluster"
assert out["yarn_rm_url"] is None assert out["yarn_rm_url"] is None
assert out["spark_conf"] == {} assert out["spark_conf"] == {}
assert out["ssl_verify"] is None
assert out["ssl_ca_bundle"] is None
def test_save_connection_preserves_ssl_fields(_fresh_store):
connections.save_connection(
name="secure",
master="yarn",
ssl_verify=False,
ssl_ca_bundle="/tmp/ca.crt",
)
out = connections.get_connection("secure")
assert out["ssl_verify"] is False
assert out["ssl_ca_bundle"] == "/tmp/ca.crt"
def test_get_connection_unknown_raises(_fresh_store): def test_get_connection_unknown_raises(_fresh_store):
+35 -5
View File
@@ -2,13 +2,26 @@
from datetime import datetime from datetime import datetime
from unittest.mock import patch from unittest.mock import patch
import pytest
from spark_executor.core import connection_store
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.models import Job from spark_executor.models import Connection, Job
from spark_executor.tools import kill from spark_executor.tools import kill
from spark_executor.tools import connections
def _fresh_stores():
store = connection_store.ConnectionStore()
connection_store.store = store
connections.store = store
kill.conn_store = store
kill.store = JobStore()
def test_kill_job_calls_yarn_kill(): def test_kill_job_calls_yarn_kill():
kill.store = JobStore() _fresh_stores()
kill.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
kill.store.put( kill.store.put(
Job( Job(
job_id="abc", job_id="abc",
@@ -22,7 +35,9 @@ def test_kill_job_calls_yarn_kill():
) )
with patch("spark_executor.tools.kill.kill_application") as m: with patch("spark_executor.tools.kill.kill_application") as m:
result = kill.kill_job("abc") result = kill.kill_job("abc")
m.assert_called_once_with("application_1", "http://rm:8088") args = m.call_args.args
assert args[0] == "application_1"
assert args[1].yarn_rm_url == "http://rm:8088"
assert result == { assert result == {
"job_id": "abc", "job_id": "abc",
"application_id": "application_1", "application_id": "application_1",
@@ -31,7 +46,22 @@ def test_kill_job_calls_yarn_kill():
def test_kill_job_raises_for_unknown_job(): def test_kill_job_raises_for_unknown_job():
kill.store = JobStore() _fresh_stores()
import pytest
with pytest.raises(KeyError): with pytest.raises(KeyError):
kill.kill_job("missing") kill.kill_job("missing")
def test_kill_job_raises_when_connection_missing():
_fresh_stores()
kill.store.put(
Job(
job_id="abc",
application_id="application_1",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="missing",
)
)
with pytest.raises(KeyError, match="Connection not found"):
kill.kill_job("abc")
+31 -3
View File
@@ -2,13 +2,26 @@
from datetime import datetime from datetime import datetime
from unittest.mock import patch from unittest.mock import patch
import pytest
from spark_executor.core import connection_store
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.models import Job from spark_executor.models import Connection, Job
from spark_executor.tools import logs from spark_executor.tools import logs
from spark_executor.tools import connections
def _fresh_stores():
store = connection_store.ConnectionStore()
connection_store.store = store
connections.store = store
logs.conn_store = store
logs.store = JobStore()
def _seed(job_id="abc", app_id="application_1"): def _seed(job_id="abc", app_id="application_1"):
logs.store = JobStore() _fresh_stores()
logs.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
logs.store.put( logs.store.put(
Job( Job(
job_id=job_id, job_id=job_id,
@@ -43,6 +56,21 @@ def test_get_job_logs_respects_custom_tail_chars():
def test_get_job_logs_raises_for_unknown_job(): def test_get_job_logs_raises_for_unknown_job():
_seed() _seed()
import pytest
with pytest.raises(KeyError): with pytest.raises(KeyError):
logs.get_job_logs("missing") logs.get_job_logs("missing")
def test_get_job_logs_raises_when_connection_missing():
_fresh_stores()
logs.store.put(
Job(
job_id="abc",
application_id="application_1",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="missing",
)
)
with pytest.raises(KeyError, match="Connection not found"):
logs.get_job_logs("abc")
+2
View File
@@ -48,6 +48,8 @@ def test_connection_defaults():
assert c.deploy_mode == "cluster" assert c.deploy_mode == "cluster"
assert c.yarn_rm_url is None assert c.yarn_rm_url is None
assert c.spark_conf == {} assert c.spark_conf == {}
assert c.ssl_verify is None
assert c.ssl_ca_bundle is None
def test_connection_master_defaults_to_yarn(): def test_connection_master_defaults_to_yarn():
+41 -30
View File
@@ -4,17 +4,28 @@ from unittest.mock import patch
import pytest import pytest
from spark_executor.core import connection_store
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.models import Job from spark_executor.models import Connection, Job
from spark_executor.tools import result from spark_executor.tools import result
from spark_executor.tools import connections
def test_result_returns_parsed_fields(): def _fresh_stores():
store = connection_store.ConnectionStore()
connection_store.store = store
connections.store = store
result.conn_store = store
result.store = JobStore() result.store = JobStore()
def _seed(job_id="abc", app_id="application_1"):
_fresh_stores()
result.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
result.store.put( result.store.put(
Job( Job(
job_id="abc", job_id=job_id,
application_id="application_1", application_id=app_id,
script_path="/tmp/j.py", script_path="/tmp/j.py",
queue="default", queue="default",
submit_time=datetime(2026, 6, 24), submit_time=datetime(2026, 6, 24),
@@ -22,6 +33,10 @@ def test_result_returns_parsed_fields():
yarn_rm_url="http://rm:8088", yarn_rm_url="http://rm:8088",
) )
) )
def test_result_returns_parsed_fields():
_seed()
raw = { raw = {
"app": { "app": {
"id": "application_1", "id": "application_1",
@@ -48,22 +63,13 @@ def test_result_returns_parsed_fields():
assert out.tracking_url == "http://nm:8088/proxy/application_1" assert out.tracking_url == "http://nm:8088/proxy/application_1"
assert out.started_time == 1700000000000 assert out.started_time == 1700000000000
assert out.finished_time == 1700000123000 assert out.finished_time == 1700000123000
assert m.call_args.args == ("application_1", "http://rm:8088") args = m.call_args.args
assert args[0] == "application_1"
assert args[1].yarn_rm_url == "http://rm:8088"
def test_result_handles_running_job(): def test_result_handles_running_job():
result.store = JobStore() _seed(job_id="running", app_id="application_2")
result.store.put(
Job(
job_id="running",
application_id="application_2",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="prod",
yarn_rm_url="http://rm:8088",
)
)
raw = { raw = {
"app": { "app": {
"id": "application_2", "id": "application_2",
@@ -87,18 +93,7 @@ def test_result_handles_running_job():
def test_result_handles_missing_optional_fields(): def test_result_handles_missing_optional_fields():
result.store = JobStore() _seed(job_id="accepted", app_id="application_3")
result.store.put(
Job(
job_id="accepted",
application_id="application_3",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="prod",
yarn_rm_url="http://rm:8088",
)
)
raw = {"app": {"id": "application_3", "state": "ACCEPTED"}} raw = {"app": {"id": "application_3", "state": "ACCEPTED"}}
import json import json
@@ -118,6 +113,22 @@ def test_result_handles_missing_optional_fields():
def test_result_raises_keyerror_for_unknown_job(): def test_result_raises_keyerror_for_unknown_job():
result.store = JobStore() _fresh_stores()
with pytest.raises(KeyError, match="Unknown job_id"): with pytest.raises(KeyError, match="Unknown job_id"):
result.get_job_result("missing") result.get_job_result("missing")
def test_result_raises_when_connection_missing():
_fresh_stores()
result.store.put(
Job(
job_id="abc",
application_id="application_1",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="missing",
)
)
with pytest.raises(KeyError, match="Connection not found"):
result.get_job_result("abc")
+37 -7
View File
@@ -2,13 +2,26 @@
from datetime import datetime from datetime import datetime
from unittest.mock import patch from unittest.mock import patch
import pytest
from spark_executor.core.job_store import JobStore from spark_executor.core.job_store import JobStore
from spark_executor.models import Job from spark_executor.core import connection_store
from spark_executor.tools import status from spark_executor.models import Connection, Job
from spark_executor.tools import connections, status
def _fresh_stores():
"""Reset job store and connection store singletons for a single test."""
store = connection_store.ConnectionStore()
connection_store.store = store
connections.store = store
status.conn_store = store
status.store = JobStore()
def test_get_job_status_returns_state(): def test_get_job_status_returns_state():
status.store = JobStore() _fresh_stores()
status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
status.store.put( status.store.put(
Job( Job(
job_id="abc", job_id="abc",
@@ -28,12 +41,29 @@ def test_get_job_status_returns_state():
assert out.application_id == "application_1" assert out.application_id == "application_1"
assert out.state == "RUNNING" assert out.state == "RUNNING"
assert "RUNNING" in out.raw assert "RUNNING" in out.raw
# yarn_rm_url is forwarded to the REST client # resolved YarnClientConfig is forwarded to the REST client
assert m.call_args.args == ("application_1", "http://rm:8088") args = m.call_args.args
assert args[0] == "application_1"
assert args[1].yarn_rm_url == "http://rm:8088"
def test_get_job_status_raises_for_unknown_job(): def test_get_job_status_raises_for_unknown_job():
status.store = JobStore() _fresh_stores()
import pytest
with pytest.raises(KeyError): with pytest.raises(KeyError):
status.get_job_status("missing") status.get_job_status("missing")
def test_get_job_status_raises_when_connection_missing():
_fresh_stores()
status.store.put(
Job(
job_id="abc",
application_id="application_1",
script_path="/tmp/j.py",
queue="default",
submit_time=datetime(2026, 6, 24),
connection="missing",
)
)
with pytest.raises(KeyError, match="Connection not found"):
status.get_job_status("abc")
+94 -15
View File
@@ -5,9 +5,11 @@ import httpx
import pytest import pytest
from common import config from common import config
from spark_executor.models import Connection
from spark_executor.core.yarn_client import ( from spark_executor.core.yarn_client import (
YarnConfigError, YarnConfigError,
YarnError, YarnError,
YarnClientConfig,
get_application_logs, get_application_logs,
get_application_status, get_application_status,
kill_application, kill_application,
@@ -15,6 +17,7 @@ from spark_executor.core.yarn_client import (
RM = "http://rm:8088" RM = "http://rm:8088"
HTTPS_RM = "https://rm:8088"
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
@@ -26,12 +29,16 @@ def _restore_settings():
jobs_dir=config.settings.jobs_dir, jobs_dir=config.settings.jobs_dir,
yarn_resource_manager_url=config.settings.yarn_resource_manager_url, yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
log_level=config.settings.log_level, log_level=config.settings.log_level,
ssl_verify_default=config.settings.ssl_verify_default,
ssl_ca_bundle_default=config.settings.ssl_ca_bundle_default,
) )
yield yield
config.settings.data_dir = snapshot.data_dir config.settings.data_dir = snapshot.data_dir
config.settings.jobs_dir = snapshot.jobs_dir config.settings.jobs_dir = snapshot.jobs_dir
config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
config.settings.log_level = snapshot.log_level config.settings.log_level = snapshot.log_level
config.settings.ssl_verify_default = snapshot.ssl_verify_default
config.settings.ssl_ca_bundle_default = snapshot.ssl_ca_bundle_default
def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response: def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response:
@@ -45,17 +52,18 @@ def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Resp
def test_status_parses_app_state(): def test_status_parses_app_state():
fake = _resp(200, json_data={"app": {"id": "application_1", "state": "RUNNING"}}) fake = _resp(200, json_data={"app": {"id": "application_1", "state": "RUNNING"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
state, raw = get_application_status("application_1", RM) state, raw = get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM))
assert state == "RUNNING" assert state == "RUNNING"
assert "RUNNING" in raw assert "RUNNING" in raw
args = m.call_args.args args = m.call_args.args
assert args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1") assert args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1")
assert m.call_args.kwargs["verify"] is True
def test_status_raises_on_404(): def test_status_raises_on_404():
with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)): with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)):
with pytest.raises(YarnError, match="not found"): with pytest.raises(YarnError, match="not found"):
get_application_status("application_x", RM) get_application_status("application_x", YarnClientConfig(yarn_rm_url=RM))
def test_status_raises_on_5xx(): def test_status_raises_on_5xx():
@@ -64,34 +72,34 @@ def test_status_raises_on_5xx():
return_value=_resp(503, text="upstream down"), return_value=_resp(503, text="upstream down"),
): ):
with pytest.raises(YarnError, match="503"): with pytest.raises(YarnError, match="503"):
get_application_status("application_1", RM) get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM))
def test_status_raises_when_state_field_missing(): def test_status_raises_when_state_field_missing():
fake = _resp(200, json_data={"app": {"id": "application_1"}}) fake = _resp(200, json_data={"app": {"id": "application_1"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake): with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake):
with pytest.raises(YarnError, match="Could not parse YARN state"): with pytest.raises(YarnError, match="Could not parse YARN state"):
get_application_status("application_1", RM) get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM))
def test_status_requires_rm_url(): def test_status_requires_rm_url():
config.settings.yarn_resource_manager_url = None config.settings.yarn_resource_manager_url = None
with pytest.raises(YarnConfigError): with pytest.raises(YarnConfigError):
get_application_status("application_1", None) get_application_status("application_1", YarnClientConfig(yarn_rm_url=None))
def test_status_falls_back_to_settings_yarn_rm_url(): def test_status_falls_back_to_settings_yarn_rm_url():
config.settings.yarn_resource_manager_url = "http://env-rm:8088" config.settings.yarn_resource_manager_url = "http://env-rm:8088"
fake = _resp(200, json_data={"app": {"state": "FINISHED"}}) fake = _resp(200, json_data={"app": {"state": "FINISHED"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
state, _ = get_application_status("application_1", None) state, _ = get_application_status("application_1", YarnClientConfig(yarn_rm_url=None))
assert state == "FINISHED" assert state == "FINISHED"
assert "env-rm:8088" in m.call_args.args[1] assert "env-rm:8088" in m.call_args.args[1]
def test_status_rejects_non_http_url(): def test_status_rejects_non_http_url():
with pytest.raises(YarnConfigError, match="must start with"): with pytest.raises(YarnConfigError, match="must start with"):
get_application_status("application_1", "rm:8088") get_application_status("application_1", YarnClientConfig(yarn_rm_url="rm:8088"))
# --- get_application_logs --- # --- get_application_logs ---
@@ -99,12 +107,13 @@ def test_status_rejects_non_http_url():
def test_logs_returns_text_on_h3_aggregated_endpoint(): def test_logs_returns_text_on_h3_aggregated_endpoint():
fake = _resp(200, text="log line 1\nlog line 2\n") fake = _resp(200, text="log line 1\nlog line 2\n")
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
out = get_application_logs("application_1", RM) out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
assert out == "log line 1\nlog line 2\n" assert out == "log line 1\nlog line 2\n"
assert m.call_args.args == ( assert m.call_args.args == (
"GET", "GET",
f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs", f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs",
) )
assert m.call_args.kwargs["verify"] is True
def test_logs_falls_back_to_nodemanager_on_404(): def test_logs_falls_back_to_nodemanager_on_404():
@@ -130,7 +139,7 @@ def test_logs_falls_back_to_nodemanager_on_404():
with patch( with patch(
"spark_executor.core.yarn_client.httpx.request", side_effect=responses "spark_executor.core.yarn_client.httpx.request", side_effect=responses
) as m: ) as m:
out = get_application_logs("application_1", RM) out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
assert "container_1" in out assert "container_1" in out
assert "container log content" in out assert "container log content" in out
calls = m.call_args_list calls = m.call_args_list
@@ -158,7 +167,7 @@ def test_logs_raises_on_404_when_nm_also_empty():
"spark_executor.core.yarn_client.httpx.request", side_effect=responses "spark_executor.core.yarn_client.httpx.request", side_effect=responses
): ):
with pytest.raises(YarnError, match="log-aggregation-enable"): with pytest.raises(YarnError, match="log-aggregation-enable"):
get_application_logs("application_1", RM) get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
def test_logs_handles_multiple_containers(): def test_logs_handles_multiple_containers():
@@ -190,7 +199,7 @@ def test_logs_handles_multiple_containers():
with patch( with patch(
"spark_executor.core.yarn_client.httpx.request", side_effect=responses "spark_executor.core.yarn_client.httpx.request", side_effect=responses
) as m: ) as m:
out = get_application_logs("application_1", RM) out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
assert "log one" in out assert "log one" in out
assert "log two" in out assert "log two" in out
assert out.index("log one") < out.index("log two") assert out.index("log one") < out.index("log two")
@@ -211,7 +220,7 @@ def test_logs_5xx_on_aggregated_endpoint_raises_immediately():
return_value=_resp(500, text="boom"), return_value=_resp(500, text="boom"),
) as m: ) as m:
with pytest.raises(YarnError, match="500"): with pytest.raises(YarnError, match="500"):
get_application_logs("application_1", RM) get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
assert m.call_count == 1 assert m.call_count == 1
@@ -220,10 +229,11 @@ def test_logs_5xx_on_aggregated_endpoint_raises_immediately():
def test_kill_sends_put_with_killed_state(): def test_kill_sends_put_with_killed_state():
fake = _resp(200, json_data={"app": {"state": "KILLED"}}) fake = _resp(200, json_data={"app": {"state": "KILLED"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
kill_application("application_1", RM) kill_application("application_1", YarnClientConfig(yarn_rm_url=RM))
args = m.call_args.args args = m.call_args.args
assert args == ("PUT", f"{RM}/ws/v1/cluster/apps/application_1/state") assert args == ("PUT", f"{RM}/ws/v1/cluster/apps/application_1/state")
assert m.call_args.kwargs["json"] == {"state": "KILLED"} assert m.call_args.kwargs["json"] == {"state": "KILLED"}
assert m.call_args.kwargs["verify"] is True
def test_kill_raises_on_5xx(): def test_kill_raises_on_5xx():
@@ -232,7 +242,7 @@ def test_kill_raises_on_5xx():
return_value=_resp(403, text="forbidden"), return_value=_resp(403, text="forbidden"),
): ):
with pytest.raises(YarnError, match="403"): with pytest.raises(YarnError, match="403"):
kill_application("application_1", RM) kill_application("application_1", YarnClientConfig(yarn_rm_url=RM))
# --- connection errors --- # --- connection errors ---
@@ -243,4 +253,73 @@ def test_status_wraps_httpx_errors_as_yarn_error():
side_effect=httpx.ConnectError("connection refused"), side_effect=httpx.ConnectError("connection refused"),
): ):
with pytest.raises(YarnError, match="connection failed"): with pytest.raises(YarnError, match="connection failed"):
get_application_status("application_1", RM) get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM))
# --- SSL/TLS configuration ---
def test_request_passes_ssl_verify_true():
config.settings.ssl_verify_default = True
fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
get_application_status(
"application_1",
YarnClientConfig(yarn_rm_url=RM, ssl_verify=True),
)
assert m.call_args.kwargs["verify"] is True
def test_request_passes_ssl_ca_bundle_path():
fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
get_application_status(
"application_1",
YarnClientConfig(
yarn_rm_url=HTTPS_RM,
ssl_ca_bundle="/path/to/ca.crt",
),
)
assert m.call_args.kwargs["verify"] == "/path/to/ca.crt"
def test_request_uses_global_ssl_default_when_connection_unset():
config.settings.ssl_verify_default = True
config.settings.ssl_ca_bundle_default = "/global/ca.crt"
config.settings.yarn_resource_manager_url = "https://rm:8088"
fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
conn = Connection(name="prod", master="yarn", ssl_verify=None, ssl_ca_bundle=None)
cfg = YarnClientConfig.from_connection(conn)
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
get_application_status("application_1", cfg)
assert m.call_args.kwargs["verify"] == "/global/ca.crt"
def test_request_per_connection_ssl_overrides_global():
config.settings.ssl_verify_default = True
config.settings.ssl_ca_bundle_default = "/global/ca.crt"
config.settings.yarn_resource_manager_url = "https://rm:8088"
fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
conn = Connection(
name="prod",
master="yarn",
ssl_verify=False,
ssl_ca_bundle="/conn/ca.crt",
)
cfg = YarnClientConfig.from_connection(conn)
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
get_application_status("application_1", cfg)
assert m.call_args.kwargs["verify"] == "/conn/ca.crt"
def test_ssl_verify_false_without_ca_bundle_uses_global_default():
"""If a connection only sets ssl_verify=False and no CA bundle is
configured globally, httpx should skip certificate verification."""
config.settings.ssl_verify_default = True
config.settings.ssl_ca_bundle_default = None
config.settings.yarn_resource_manager_url = "https://rm:8088"
fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
conn = Connection(name="prod", master="yarn", ssl_verify=False)
cfg = YarnClientConfig.from_connection(conn)
with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
get_application_status("application_1", cfg)
assert m.call_args.kwargs["verify"] is False