From ad557b59844e29f034211b277cca23c37e05e99b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Jun 2026 13:50:19 +0800 Subject: [PATCH] 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. --- common/config.py | 13 ++++ spark_executor/core/yarn_client.py | 77 +++++++++++++++----- spark_executor/models.py | 4 + spark_executor/tools/connections.py | 4 + spark_executor/tools/kill.py | 9 ++- spark_executor/tools/logs.py | 9 ++- spark_executor/tools/requests.py | 2 + spark_executor/tools/result.py | 9 ++- spark_executor/tools/status.py | 8 +- tests/unit/test_connection_tools.py | 14 ++++ tests/unit/test_kill_tool.py | 40 ++++++++-- tests/unit/test_logs_tool.py | 34 ++++++++- tests/unit/test_models.py | 2 + tests/unit/test_result_tool.py | 71 ++++++++++-------- tests/unit/test_status_tool.py | 44 +++++++++-- tests/unit/test_yarn_client.py | 109 ++++++++++++++++++++++++---- 16 files changed, 362 insertions(+), 87 deletions(-) diff --git a/common/config.py b/common/config.py index a55b1c0..19afa64 100644 --- a/common/config.py +++ b/common/config.py @@ -76,6 +76,10 @@ class Settings: # The binary is resolved against $PATH (set by docker-entrypoint.sh). 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 def from_env(cls) -> "Settings": data_dir = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") @@ -95,6 +99,13 @@ class Settings: spark_submit_bin=os.environ.get( "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": @@ -110,6 +121,8 @@ class Settings: self.yarn_resource_manager_url = fresh.yarn_resource_manager_url self.log_level = fresh.log_level 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 diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index 8636bc2..d574fc7 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -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}") diff --git a/spark_executor/models.py b/spark_executor/models.py index b643593..b22f5bb 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -49,6 +49,10 @@ class Connection(BaseModel): deploy_mode: str = "cluster" yarn_rm_url: str | None = None 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") @classmethod diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index f76c8f9..c29f91a 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -15,6 +15,8 @@ def save_connection( 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, ) -> dict[str, str]: logger.debug( f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " @@ -26,6 +28,8 @@ def save_connection( 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, ) store.save(conn) return {"name": name, "status": "SAVED"} diff --git a/spark_executor/tools/kill.py b/spark_executor/tools/kill.py index a35de04..a922e9b 100644 --- a/spark_executor/tools/kill.py +++ b/spark_executor/tools/kill.py @@ -5,7 +5,8 @@ """ from common.logging import logger 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() @@ -15,7 +16,11 @@ def kill_job(job_id: str) -> dict[str, str]: job = store.get(job_id) if job is None: 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}") return { "job_id": job_id, diff --git a/spark_executor/tools/logs.py b/spark_executor/tools/logs.py index f06dea6..7456c56 100644 --- a/spark_executor/tools/logs.py +++ b/spark_executor/tools/logs.py @@ -5,7 +5,8 @@ """ from common.logging import logger 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() @@ -15,7 +16,11 @@ def get_job_logs(job_id: str, tail_chars: int = 5000) -> str: job = store.get(job_id) if job is None: 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 logger.info( f"get_job_logs ok job_id={job_id} application_id={job.application_id} " diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index f7d5eb2..d58cb4c 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -22,6 +22,8 @@ class SaveConnectionRequest(BaseModel): 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 class PrepareSubmitJobRequest(BaseModel): diff --git a/spark_executor/tools/result.py b/spark_executor/tools/result.py index 9c7b505..7e14265 100644 --- a/spark_executor/tools/result.py +++ b/spark_executor/tools/result.py @@ -7,7 +7,8 @@ import json from common.logging import logger 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 store = JobStore() @@ -18,7 +19,11 @@ def get_job_result(job_id: str) -> JobResult: job = store.get(job_id) if job is None: 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", {}) result = JobResult( application_id=job.application_id, diff --git a/spark_executor/tools/status.py b/spark_executor/tools/status.py index 2323765..34ffbea 100644 --- a/spark_executor/tools/status.py +++ b/spark_executor/tools/status.py @@ -4,9 +4,11 @@ @Author :tao.chen """ 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.yarn_client import get_application_status from spark_executor.models import JobStatus +from spark_executor.tools.connections import store as conn_store store = JobStore() @@ -16,6 +18,10 @@ def get_job_status(job_id: str) -> JobStatus: job = store.get(job_id) if job is None: 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}") return JobStatus(application_id=job.application_id, state=state, raw=raw) diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py index 4307622..890820d 100644 --- a/tests/unit/test_connection_tools.py +++ b/tests/unit/test_connection_tools.py @@ -65,6 +65,20 @@ def test_get_connection_returns_dict(_fresh_store): assert out["deploy_mode"] == "cluster" assert out["yarn_rm_url"] is None 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): diff --git a/tests/unit/test_kill_tool.py b/tests/unit/test_kill_tool.py index bdc851d..de1c7f3 100644 --- a/tests/unit/test_kill_tool.py +++ b/tests/unit/test_kill_tool.py @@ -2,13 +2,26 @@ from datetime import datetime 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.models import Job +from spark_executor.models import Connection, Job 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(): - kill.store = JobStore() + _fresh_stores() + kill.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")) kill.store.put( Job( job_id="abc", @@ -22,7 +35,9 @@ def test_kill_job_calls_yarn_kill(): ) with patch("spark_executor.tools.kill.kill_application") as m: 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 == { "job_id": "abc", "application_id": "application_1", @@ -31,7 +46,22 @@ def test_kill_job_calls_yarn_kill(): def test_kill_job_raises_for_unknown_job(): - kill.store = JobStore() - import pytest + _fresh_stores() with pytest.raises(KeyError): 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") diff --git a/tests/unit/test_logs_tool.py b/tests/unit/test_logs_tool.py index 1b518a9..2e30f2c 100644 --- a/tests/unit/test_logs_tool.py +++ b/tests/unit/test_logs_tool.py @@ -2,13 +2,26 @@ from datetime import datetime 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.models import Job +from spark_executor.models import Connection, Job 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"): - logs.store = JobStore() + _fresh_stores() + logs.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")) logs.store.put( Job( 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(): _seed() - import pytest with pytest.raises(KeyError): 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") diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py index 21ea59c..4b9b053 100644 --- a/tests/unit/test_models.py +++ b/tests/unit/test_models.py @@ -48,6 +48,8 @@ def test_connection_defaults(): assert c.deploy_mode == "cluster" assert c.yarn_rm_url is None assert c.spark_conf == {} + assert c.ssl_verify is None + assert c.ssl_ca_bundle is None def test_connection_master_defaults_to_yarn(): diff --git a/tests/unit/test_result_tool.py b/tests/unit/test_result_tool.py index 83549be..ca7333f 100644 --- a/tests/unit/test_result_tool.py +++ b/tests/unit/test_result_tool.py @@ -4,17 +4,28 @@ 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.models import Job +from spark_executor.models import Connection, Job 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() + + +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( Job( - job_id="abc", - application_id="application_1", + job_id=job_id, + application_id=app_id, script_path="/tmp/j.py", queue="default", submit_time=datetime(2026, 6, 24), @@ -22,6 +33,10 @@ def test_result_returns_parsed_fields(): yarn_rm_url="http://rm:8088", ) ) + + +def test_result_returns_parsed_fields(): + _seed() raw = { "app": { "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.started_time == 1700000000000 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(): - result.store = JobStore() - 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", - ) - ) + _seed(job_id="running", app_id="application_2") raw = { "app": { "id": "application_2", @@ -87,18 +93,7 @@ def test_result_handles_running_job(): def test_result_handles_missing_optional_fields(): - result.store = JobStore() - 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", - ) - ) + _seed(job_id="accepted", app_id="application_3") raw = {"app": {"id": "application_3", "state": "ACCEPTED"}} import json @@ -118,6 +113,22 @@ def test_result_handles_missing_optional_fields(): def test_result_raises_keyerror_for_unknown_job(): - result.store = JobStore() + _fresh_stores() with pytest.raises(KeyError, match="Unknown job_id"): 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") diff --git a/tests/unit/test_status_tool.py b/tests/unit/test_status_tool.py index fc49273..e3bc30b 100644 --- a/tests/unit/test_status_tool.py +++ b/tests/unit/test_status_tool.py @@ -2,13 +2,26 @@ from datetime import datetime from unittest.mock import patch +import pytest + from spark_executor.core.job_store import JobStore -from spark_executor.models import Job -from spark_executor.tools import status +from spark_executor.core import connection_store +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(): - status.store = JobStore() + _fresh_stores() + status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")) status.store.put( Job( job_id="abc", @@ -28,12 +41,29 @@ def test_get_job_status_returns_state(): assert out.application_id == "application_1" assert out.state == "RUNNING" assert "RUNNING" in out.raw - # yarn_rm_url is forwarded to the REST client - assert m.call_args.args == ("application_1", "http://rm:8088") + # resolved YarnClientConfig is forwarded to the REST client + 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(): - status.store = JobStore() - import pytest + _fresh_stores() with pytest.raises(KeyError): 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") diff --git a/tests/unit/test_yarn_client.py b/tests/unit/test_yarn_client.py index f340795..2ed8d2f 100644 --- a/tests/unit/test_yarn_client.py +++ b/tests/unit/test_yarn_client.py @@ -5,9 +5,11 @@ import httpx import pytest from common import config +from spark_executor.models import Connection from spark_executor.core.yarn_client import ( YarnConfigError, YarnError, + YarnClientConfig, get_application_logs, get_application_status, kill_application, @@ -15,6 +17,7 @@ from spark_executor.core.yarn_client import ( RM = "http://rm:8088" +HTTPS_RM = "https://rm:8088" @pytest.fixture(autouse=True) @@ -26,12 +29,16 @@ def _restore_settings(): jobs_dir=config.settings.jobs_dir, yarn_resource_manager_url=config.settings.yarn_resource_manager_url, 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 config.settings.data_dir = snapshot.data_dir config.settings.jobs_dir = snapshot.jobs_dir config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url 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: @@ -45,17 +52,18 @@ def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Resp def test_status_parses_app_state(): 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: - state, raw = get_application_status("application_1", RM) + state, raw = get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM)) assert state == "RUNNING" assert "RUNNING" in raw args = m.call_args.args 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(): with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)): 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(): @@ -64,34 +72,34 @@ def test_status_raises_on_5xx(): return_value=_resp(503, text="upstream down"), ): 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(): fake = _resp(200, json_data={"app": {"id": "application_1"}}) with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake): 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(): config.settings.yarn_resource_manager_url = None 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(): config.settings.yarn_resource_manager_url = "http://env-rm:8088" fake = _resp(200, json_data={"app": {"state": "FINISHED"}}) 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 "env-rm:8088" in m.call_args.args[1] def test_status_rejects_non_http_url(): 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 --- @@ -99,12 +107,13 @@ def test_status_rejects_non_http_url(): def test_logs_returns_text_on_h3_aggregated_endpoint(): 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: - 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 m.call_args.args == ( "GET", 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(): @@ -130,7 +139,7 @@ def test_logs_falls_back_to_nodemanager_on_404(): with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=responses ) 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 log content" in out 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 ): 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(): @@ -190,7 +199,7 @@ def test_logs_handles_multiple_containers(): with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=responses ) 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 two" in out 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"), ) as m: 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 @@ -220,10 +229,11 @@ def test_logs_5xx_on_aggregated_endpoint_raises_immediately(): def test_kill_sends_put_with_killed_state(): fake = _resp(200, json_data={"app": {"state": "KILLED"}}) 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 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["verify"] is True def test_kill_raises_on_5xx(): @@ -232,7 +242,7 @@ def test_kill_raises_on_5xx(): return_value=_resp(403, text="forbidden"), ): with pytest.raises(YarnError, match="403"): - kill_application("application_1", RM) + kill_application("application_1", YarnClientConfig(yarn_rm_url=RM)) # --- connection errors --- @@ -243,4 +253,73 @@ def test_status_wraps_httpx_errors_as_yarn_error(): side_effect=httpx.ConnectError("connection refused"), ): 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