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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/test_connection_tools.py
+14
View File
@@ -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):
tests/unit/test_kill_tool.py
+35 -5
View File
@@ -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")
tests/unit/test_logs_tool.py
+31 -3
View File
@@ -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")
tests/unit/test_models.py
+2
View File
@@ -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():
tests/unit/test_result_tool.py
+41 -30
View File
@@ -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")
tests/unit/test_status_tool.py
+37 -7
View File
@@ -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")
tests/unit/test_yarn_client.py
+94 -15
View File
@@ -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