diff --git a/src/drs/commands/setup.py b/src/drs/commands/setup.py index fb77ea2..fead31c 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,43 @@ err_console = Console(stderr=True) +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_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_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_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_list}, 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 +226,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..7931e65 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,68 @@ 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_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" + 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"