From f80c6479c4251c3853a27477e5074f2d77c08b28 Mon Sep 17 00:00:00 2001 From: Vando Pereira Date: Fri, 5 Jun 2026 09:14:52 +0100 Subject: [PATCH 1/2] Warn setup users about env token overrides --- src/drs/commands/setup.py | 41 ++++++++++++++++++++++++++ tests/test_commands/test_setup.py | 49 +++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/src/drs/commands/setup.py b/src/drs/commands/setup.py index fb77ea2..9edebd5 100644 --- a/src/drs/commands/setup.py +++ b/src/drs/commands/setup.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import os import sys from pathlib import Path from typing import Any @@ -41,6 +42,44 @@ err_console = Console(stderr=True) +def _active_token_env_var() -> str | None: + """Return the token env var that load_config() will use, if one is set.""" + if os.environ.get("DREMIO_TOKEN"): + return "DREMIO_TOKEN" + if os.environ.get("DREMIO_PAT"): + return "DREMIO_PAT" + return None + + +def _warn_token_env_override(config_path: Path) -> None: + """Warn when an env token will override the config written by setup.""" + token_env = _active_token_env_var() + if token_env is None: + return + + console.print() + console.print( + Panel( + f"[bold yellow]{token_env} is set in your environment.[/bold yellow]\n\n" + "Dremio CLI resolves credentials in this order:\n" + " 1. --token\n" + " 2. DREMIO_TOKEN / DREMIO_PAT\n" + f" 3. {config_path}\n\n" + "That means future [bold]dremio[/bold] commands will use the environment token, " + "not the PAT saved by this setup wizard.\n\n" + f"To use the saved config, run [bold]unset {token_env}[/bold] before using dremio. " + f"To keep using environment auth, update [bold]{token_env}[/bold] to the PAT you just entered.", + title="Environment Token Overrides Config", + border_style="yellow", + ) + ) + if not typer.confirm("Continue setup anyway?", default=False): + console.print( + f"Setup cancelled. Unset or update {token_env}, then run [bold cyan]dremio setup[/bold cyan] again." + ) + raise typer.Exit(1) + + async def validate_credentials(uri: str, pat: str, project_id: str) -> tuple[bool, str, dict[str, Any] | None]: """Test credentials by calling get_project(). Returns (success, message, project_data).""" config = DrsConfig(uri=uri, pat=pat, project_id=project_id) @@ -188,6 +227,8 @@ def setup_command( console.print("Setup cancelled.") raise typer.Exit(0) + _warn_token_env_override(config_path) + # Step 1: Region api_uri, app_url = _prompt_region() diff --git a/tests/test_commands/test_setup.py b/tests/test_commands/test_setup.py index 2bab050..274c2f0 100644 --- a/tests/test_commands/test_setup.py +++ b/tests/test_commands/test_setup.py @@ -31,6 +31,12 @@ runner = CliRunner() +@pytest.fixture(autouse=True) +def clear_dremio_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("DREMIO_TOKEN", "DREMIO_PAT", "DREMIO_PROJECT_ID", "DREMIO_URI"): + monkeypatch.delenv(name, raising=False) + + def test_write_config(tmp_path) -> None: config_path = tmp_path / "config.yaml" write_config("https://api.eu.dremio.cloud", "my-pat", "my-project", config_path) @@ -227,6 +233,49 @@ def test_setup_existing_config_overwrite(tmp_path) -> None: assert data["pat"] == "new-pat" +def test_setup_warns_when_dremio_pat_overrides_config(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + """A DREMIO_PAT env var should be called out before writing config.""" + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("DREMIO_PAT", "env-pat") + + with ( + patch("drs.commands.setup.sys") as mock_sys, + patch("drs.commands.setup.DEFAULT_CONFIG_PATH", config_path), + ): + mock_sys.stdin.isatty.return_value = True + result = runner.invoke(app, ["setup"], input="n\n") + + assert result.exit_code == 1 + assert "DREMIO_PAT is set" in result.output + assert "will use the environment token" in result.output + assert "unset DREMIO_PAT" in result.output + assert not config_path.exists() + + +def test_setup_can_continue_after_env_token_warning(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + """Accepting the env-token warning should continue through normal setup.""" + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("DREMIO_PAT", "env-pat") + + mock_client = AsyncMock() + mock_client.get_project = AsyncMock(return_value={"id": "p1", "name": "Test Project"}) + mock_client.close = AsyncMock() + + with ( + patch("drs.commands.setup.sys") as mock_sys, + patch("drs.commands.setup.DremioClient", return_value=mock_client), + patch("drs.commands.setup.DEFAULT_CONFIG_PATH", config_path), + ): + mock_sys.stdin.isatty.return_value = True + result = runner.invoke(app, ["setup"], input="y\n1\nfile-pat\nfile-project\n") + + assert result.exit_code == 0 + assert "DREMIO_PAT is set" in result.output + data = yaml.safe_load(config_path.read_text()) + assert data["pat"] == "file-pat" + assert data["project_id"] == "file-project" + + def test_setup_retry_then_abort(tmp_path) -> None: """Validation failure followed by declining retry should exit 1.""" config_path = tmp_path / "config.yaml" From 26999347144b13952f85ec0d174c99583ceccbdf Mon Sep 17 00:00:00 2001 From: Vando Pereira Date: Mon, 8 Jun 2026 09:58:46 +0100 Subject: [PATCH 2/2] Include all token env vars in setup guidance --- src/drs/commands/setup.py | 25 ++++++++++++------------- tests/test_commands/test_setup.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/src/drs/commands/setup.py b/src/drs/commands/setup.py index 9edebd5..fead31c 100644 --- a/src/drs/commands/setup.py +++ b/src/drs/commands/setup.py @@ -42,40 +42,39 @@ err_console = Console(stderr=True) -def _active_token_env_var() -> str | None: - """Return the token env var that load_config() will use, if one is set.""" - if os.environ.get("DREMIO_TOKEN"): - return "DREMIO_TOKEN" - if os.environ.get("DREMIO_PAT"): - return "DREMIO_PAT" - return None +def _set_token_env_vars() -> list[str]: + """Return token env vars that will override config, in resolution order.""" + return [name for name in ("DREMIO_TOKEN", "DREMIO_PAT") if os.environ.get(name)] def _warn_token_env_override(config_path: Path) -> None: """Warn when an env token will override the config written by setup.""" - token_env = _active_token_env_var() - if token_env is None: + token_envs = _set_token_env_vars() + if not token_envs: return + token_env_list = " and ".join(token_envs) + token_env_verb = "is" if len(token_envs) == 1 else "are" + unset_command = f"unset {' '.join(token_envs)}" console.print() console.print( Panel( - f"[bold yellow]{token_env} is set in your environment.[/bold yellow]\n\n" + f"[bold yellow]{token_env_list} {token_env_verb} set in your environment.[/bold yellow]\n\n" "Dremio CLI resolves credentials in this order:\n" " 1. --token\n" " 2. DREMIO_TOKEN / DREMIO_PAT\n" f" 3. {config_path}\n\n" "That means future [bold]dremio[/bold] commands will use the environment token, " "not the PAT saved by this setup wizard.\n\n" - f"To use the saved config, run [bold]unset {token_env}[/bold] before using dremio. " - f"To keep using environment auth, update [bold]{token_env}[/bold] to the PAT you just entered.", + f"To use the saved config, run [bold]{unset_command}[/bold] before using dremio. " + f"To keep using environment auth, update [bold]{token_env_list}[/bold] to the PAT you just entered.", title="Environment Token Overrides Config", border_style="yellow", ) ) if not typer.confirm("Continue setup anyway?", default=False): console.print( - f"Setup cancelled. Unset or update {token_env}, then run [bold cyan]dremio setup[/bold cyan] again." + f"Setup cancelled. Unset or update {token_env_list}, then run [bold cyan]dremio setup[/bold cyan] again." ) raise typer.Exit(1) diff --git a/tests/test_commands/test_setup.py b/tests/test_commands/test_setup.py index 274c2f0..7931e65 100644 --- a/tests/test_commands/test_setup.py +++ b/tests/test_commands/test_setup.py @@ -252,6 +252,25 @@ def test_setup_warns_when_dremio_pat_overrides_config(tmp_path, monkeypatch: pyt assert not config_path.exists() +def test_setup_warns_to_unset_both_token_env_vars(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + """If both token env vars are set, both must be included in unset guidance.""" + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("DREMIO_TOKEN", "env-token") + monkeypatch.setenv("DREMIO_PAT", "env-pat") + + with ( + patch("drs.commands.setup.sys") as mock_sys, + patch("drs.commands.setup.DEFAULT_CONFIG_PATH", config_path), + ): + mock_sys.stdin.isatty.return_value = True + result = runner.invoke(app, ["setup"], input="n\n") + + assert result.exit_code == 1 + assert "DREMIO_TOKEN and DREMIO_PAT are set" in result.output + assert "unset DREMIO_TOKEN DREMIO_PAT" in result.output + assert not config_path.exists() + + def test_setup_can_continue_after_env_token_warning(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: """Accepting the env-token warning should continue through normal setup.""" config_path = tmp_path / "config.yaml"