From 0a7d8dfc1b19b1a7539dfa265868cd70e8220456 Mon Sep 17 00:00:00 2001 From: Zemin Piao Date: Fri, 20 Mar 2026 00:42:02 +0100 Subject: [PATCH] Add basic auth functionality Signed-off-by: Zemin Piao --- spark_history_cli/cli.py | 39 +++++++++++++++++++++++++--- spark_history_cli/core/client.py | 10 ++++++- spark_history_cli/tests/test_core.py | 36 +++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 4 deletions(-) diff --git a/spark_history_cli/cli.py b/spark_history_cli/cli.py index e4bd053..3d62cf4 100644 --- a/spark_history_cli/cli.py +++ b/spark_history_cli/cli.py @@ -28,10 +28,16 @@ def __init__(self): self.client: SparkHistoryClient | None = None self.session: Session = Session() self.json_mode: bool = False + self.basic_auth_username: str | None = None + self.basic_auth_password: str | None = None def ensure_client(self) -> SparkHistoryClient: if self.client is None: - self.client = SparkHistoryClient(self.session.server_url) + self.client = SparkHistoryClient( + self.session.server_url, + basic_auth_username=self.basic_auth_username, + basic_auth_password=self.basic_auth_password, + ) return self.client def resolve_app_id(self, app_id: str | None) -> str: @@ -86,17 +92,39 @@ def _fetch_sql_jobs(client, app_id: str, sql_exec: dict) -> list[dict]: @click.option("--server", "-s", default="http://localhost:18080", envvar="SPARK_HISTORY_SERVER", help="Spark History Server URL (default: http://localhost:18080)") +@click.option("--basic-auth-user", default=None, + envvar="SPARK_HISTORY_BASIC_AUTH_USER", + help="Basic Auth username for Spark History Server") +@click.option("--basic-auth-password", default=None, + envvar="SPARK_HISTORY_BASIC_AUTH_PASSWORD", + help="Basic Auth password for Spark History Server (omit to prompt securely)") @click.option("--json", "json_mode", is_flag=True, default=False, help="Output in JSON format for machine consumption") @click.option("--app-id", "-a", default=None, help="Application ID to use (sets context for subcommands)") @click.version_option(__version__, prog_name="spark-history-cli") @click.pass_context -def cli(ctx, server: str, json_mode: bool, app_id: str | None): +def cli( + ctx, + server: str, + basic_auth_user: str | None, + basic_auth_password: str | None, + json_mode: bool, + app_id: str | None, +): """CLI for querying the Apache Spark History Server REST API.""" + if basic_auth_password is not None and basic_auth_user is None: + raise click.UsageError( + "--basic-auth-password requires --basic-auth-user." + ) + if basic_auth_user is not None and basic_auth_password is None: + basic_auth_password = click.prompt("Basic Auth password", hide_input=True) + state = CliState() state.session.server_url = server state.json_mode = json_mode + state.basic_auth_username = basic_auth_user + state.basic_auth_password = basic_auth_password if app_id: state.session.set_app(app_id) state.ensure_client() @@ -183,7 +211,11 @@ def repl(state: CliState): elif cmd == "server": if args: state.session.server_url = args[0] - state.client = SparkHistoryClient(args[0]) + state.client = SparkHistoryClient( + args[0], + basic_auth_username=state.basic_auth_username, + basic_auth_password=state.basic_auth_password, + ) client = state.client try: ver = client.get_version() @@ -196,6 +228,7 @@ def repl(state: CliState): elif cmd == "status": skin.status_block({ "Server": state.session.server_url, + "Basic Auth": "enabled" if state.basic_auth_username else "disabled", "Current App": state.session.current_app_id or "(none)", "Attempt": state.session.current_attempt_id or "(none)", }, title="Session") diff --git a/spark_history_cli/core/client.py b/spark_history_cli/core/client.py index 379fa45..ceaef44 100644 --- a/spark_history_cli/core/client.py +++ b/spark_history_cli/core/client.py @@ -24,11 +24,19 @@ def __init__(self, status_code: int, message: str, url: str = ""): class SparkHistoryClient: """Client for the Spark History Server REST API (/api/v1).""" - def __init__(self, server_url: str = "http://localhost:18080", timeout: int = 30): + def __init__( + self, + server_url: str = "http://localhost:18080", + timeout: int = 30, + basic_auth_username: str | None = None, + basic_auth_password: str | None = None, + ): self.server_url = server_url.rstrip("/") self.base_url = f"{self.server_url}/api/v1" self.timeout = timeout self._session = requests.Session() + if basic_auth_username is not None: + self._session.auth = (basic_auth_username, basic_auth_password or "") self._attempt_cache: dict[str, str | None] = {} def _resolve_attempt(self, app_id: str) -> str: diff --git a/spark_history_cli/tests/test_core.py b/spark_history_cli/tests/test_core.py index 174542d..fa40464 100644 --- a/spark_history_cli/tests/test_core.py +++ b/spark_history_cli/tests/test_core.py @@ -13,10 +13,12 @@ from unittest.mock import patch, MagicMock import pytest +from click.testing import CliRunner from spark_history_cli.core.client import SparkHistoryClient, HistoryServerError from spark_history_cli.core.session import Session from spark_history_cli.core import formatters as fmt +from spark_history_cli.cli import cli # ── Sample API Responses ────────────────────────────────────────────── @@ -227,6 +229,14 @@ def test_connection_error(self): client.get_version() assert "Cannot connect" in str(exc_info.value) + def test_basic_auth_is_configured(self): + client = SparkHistoryClient( + "http://test:18080", + basic_auth_username="alice", + basic_auth_password="secret", + ) + assert client._session.auth == ("alice", "secret") + def test_http_404(self): client = self._mock_client({"message": "not found"}, status_code=404) with pytest.raises(HistoryServerError) as exc_info: @@ -433,3 +443,29 @@ def test_install_skill(self): assert result.returncode == 0 assert os.path.exists(os.path.join(target, "SKILL.md")) assert "Installed Copilot skill" in result.stdout + + +class TestCLIOptions: + def test_basic_auth_options_are_accepted(self): + runner = CliRunner() + result = runner.invoke( + cli, + ["--basic-auth-user", "alice", "--basic-auth-password", "secret", "--help"], + ) + assert result.exit_code == 0 + + def test_basic_auth_password_requires_user(self): + runner = CliRunner() + result = runner.invoke(cli, ["--basic-auth-password", "secret", "apps"]) + assert result.exit_code != 0 + assert "--basic-auth-password requires --basic-auth-user" in result.output + + def test_basic_auth_user_prompts_for_password_hidden(self): + runner = CliRunner() + mock_client = MagicMock() + mock_client.get_version.return_value = {"spark": "4.0.0"} + with patch("spark_history_cli.cli.click.prompt", return_value="secret") as prompt_mock: + with patch("spark_history_cli.cli.CliState.ensure_client", return_value=mock_client): + result = runner.invoke(cli, ["--basic-auth-user", "alice", "version"]) + assert result.exit_code == 0 + prompt_mock.assert_called_once_with("Basic Auth password", hide_input=True)