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:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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} "
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user