From c714a135fb7e9a105b78baa73f203a3875903c78 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Thu, 30 Apr 2026 09:22:54 +0000 Subject: [PATCH] Fix CLI token source --profile fallback with version detection Replace the broken error-based --profile detection with version-based CLI detection at init time. The --profile flag is a global Cobra flag that old CLIs (< v0.207.1) silently accept instead of reporting "unknown flag: --profile", making the previous --host fallback dead code. The SDK now runs `databricks version` at init time, parses the semver, and uses it to decide between --profile and --host. If version detection fails, the SDK falls back to --host (most conservative command). This also simplifies CliTokenSource to hold a single cmd with no runtime fallback logic, and enables future version-gated flags (e.g. --force-refresh) without additional subprocess calls. See: https://github.com/databricks/cli/pull/855 --- NEXT_CHANGELOG.md | 3 + databricks/sdk/credentials_provider.py | 222 +++++++++++-- tests/test_credentials_provider.py | 441 ++++++++++++++++--------- 3 files changed, 471 insertions(+), 195 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 150c60b2e..9d43c9aca 100755 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,12 +7,15 @@ ### Security ### Bug Fixes +* Fixed Databricks CLI `--profile` fallback by detecting the CLI version at init time. The previous error-based detection was broken because `--profile` is a global Cobra flag silently accepted by old CLIs. ### Documentation ### Breaking Changes ### Internal Changes +* Detect Databricks CLI version at init time via `databricks version`, enabling version-gated flag support without additional subprocess calls. +* Validate Databricks CLI configuration at `DatabricksCliTokenSource.__init__` time. Misconfiguration (missing profile and host, or `--profile`-unsupported CLI without a host fallback) now surfaces as `IOError` synchronously from construction rather than lazily from the first `refresh()` call. The exception type matches the previous `refresh()`-time behaviour, so callers who already catch `IOError` are unaffected. ### API Changes * Add [w.temporary_volume_credentials](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/catalog/temporary_volume_credentials.html) workspace-level service. diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index a992ea67a..090c54edd 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -1,5 +1,6 @@ import abc import base64 +import dataclasses import functools import io import json @@ -662,7 +663,47 @@ def refreshed_headers() -> Dict[str, str]: return OAuthCredentialsProvider(refreshed_headers, token) +@dataclasses.dataclass(order=True) +class CliVersion: + """Semver version triple of the Databricks CLI. + + Three sentinel states in the (major, minor, patch) tuple: + * `(-1, -1, -1)` — the default, meaning version detection failed. It + compares less than every real release so every feature gate fails. + * `(0, 0, 0)` — the CLI's default dev build, emitted when the binary + was built without version metadata. See `is_default_dev_build`. + * anything else — a real CLI version. + + Prerelease tags (e.g. "-rc.1", "-dev+commit") are deliberately ignored: + for our purposes the base triple is sufficient, and feature gates are + release-based so a prerelease of a version with a flag is assumed to + have the flag too. + """ + + major: int = -1 + minor: int = -1 + patch: int = -1 + + @property + def is_default_dev_build(self) -> bool: + """True when the CLI reports (0, 0, 0), i.e. its default dev build. + + Narrowly matches the CLI's "no version injected" marker: its version + metadata stays at the zero defaults. A version the user explicitly + set (e.g. v1.0.0-dev) is intentional and is not flagged here. + """ + return (self.major, self.minor, self.patch) == (0, 0, 0) + + def __str__(self) -> str: + if self == CliVersion(): + return "unknown" + if self.is_default_dev_build: + return "v0.0.0-dev" + return f"v{self.major}.{self.minor}.{self.patch}" + + class CliTokenSource(oauth.Refreshable): + def __init__( self, cmd: List[str], @@ -670,15 +711,9 @@ def __init__( access_token_field: str, expiry_field: str, disable_async: bool = True, - fallback_cmd: Optional[List[str]] = None, ): super().__init__(disable_async=disable_async) self._cmd = cmd - # fallback_cmd is tried when the primary command fails with "unknown flag: --profile", - # indicating the CLI is too old to support --profile. Can be removed once support - # for CLI versions predating --profile is dropped. - # See: https://github.com/databricks/databricks-sdk-go/pull/1497 - self._fallback_cmd = fallback_cmd self._token_type_field = token_type_field self._access_token_field = access_token_field self._expiry_field = expiry_field @@ -713,16 +748,7 @@ def _exec_cli_command(self, cmd: List[str]) -> oauth.Token: raise IOError(f"cannot get access token: {message}") from e def refresh(self) -> oauth.Token: - try: - return self._exec_cli_command(self._cmd) - except IOError as e: - if self._fallback_cmd is not None and "unknown flag: --profile" in str(e): - logger.warning( - "Databricks CLI does not support --profile flag. Falling back to --host. " - "Please upgrade your CLI to the latest version." - ) - return self._exec_cli_command(self._fallback_cmd) - raise + return self._exec_cli_command(self._cmd) def _run_subprocess( @@ -911,16 +937,7 @@ def __init__(self, cfg: "Config"): elif cli_path.count("/") == 0: cli_path = self.__class__._find_executable(cli_path) - fallback_cmd = None - if cfg.profile: - # When profile is set, use --profile as the primary command. - # The profile contains the full config (host, account_id, etc.). - args = ["auth", "token", "--profile", cfg.profile] - # Build a --host fallback for older CLIs that don't support --profile. - if cfg.host: - fallback_cmd = [cli_path, *self.__class__._build_host_args(cfg)] - else: - args = self.__class__._build_host_args(cfg) + cmd = self.__class__._resolve_cli_command(cli_path, cfg) # get_scopes() defaults to ["all-apis"] when nothing is configured, which would # cause false-positive mismatches against every token that wasn't issued with @@ -930,12 +947,11 @@ def __init__(self, cfg: "Config"): self._host = cfg.host super().__init__( - cmd=[cli_path, *args], + cmd=cmd, token_type_field="token_type", access_token_field="access_token", expiry_field="expiry", disable_async=cfg.disable_async_token_refresh, - fallback_cmd=fallback_cmd, ) def refresh(self) -> oauth.Token: @@ -993,15 +1009,151 @@ def _validate_token_scopes(self, token: oauth.Token): f"Scopes default to all-apis." ) + # --profile support added in CLI v0.207.1: https://github.com/databricks/cli/pull/855 + _CLI_VERSION_FOR_PROFILE = CliVersion(0, 207, 1) + + @staticmethod + def _parse_cli_version(output: str) -> CliVersion: + """Parse the JSON output of `databricks version --output json`. + + Takes Major/Minor/Patch from the JSON's pre-parsed numeric fields. The + `Prerelease` JSON field and the `Version` string are intentionally + ignored: for our feature-gate purposes the base triple is sufficient, + and the (0, 0, 0) case already identifies the default dev build (a + CLI built without version metadata leaves these fields at their zero + defaults). + + Returns CliVersion() (unknown, (-1, -1, -1)) on failure so that an + unparseable version disables every feature gate. + """ + try: + data = json.loads(output) + return CliVersion( + int(data["Major"]), + int(data["Minor"]), + int(data["Patch"]), + ) + except (json.JSONDecodeError, KeyError, TypeError, ValueError) as e: + logger.debug(f"Failed to parse Databricks CLI version from output: {output!r} ({e})") + return CliVersion() + + @staticmethod + @functools.lru_cache(maxsize=8) + def _probe_cli_version(cli_path: str) -> "CliVersion": + """Run `databricks version --output json` and return the parsed CliVersion. + + Cached by `cli_path` for the Python process lifetime: a successful probe + against an immutable binary doesn't need to be repeated. Subprocess-level + failures (timeout, OSError, CalledProcessError) raise out instead of + returning a sentinel, so transient errors are NOT cached and the next + caller retries — see `_get_cli_version` for the catching wrapper. + + The 5-second timeout prevents a hung CLI (first-run disk scan, antivirus, + stdin wait) from wedging SDK init indefinitely. + """ + out = _run_subprocess( + [cli_path, "version", "--output", "json"], + capture_output=True, + check=True, + timeout=5, + ) + return DatabricksCliTokenSource._parse_cli_version(out.stdout.decode()) + + @staticmethod + def _get_cli_version(cli_path: str) -> "CliVersion": + """Probe wrapper that catches subprocess failures and returns CliVersion(). + + The catch is intentionally outside the cached helper so a transient failure + (timeout, AV scan, sandbox quota, etc.) does not pin every subsequent + DatabricksCliTokenSource to the conservative fallback for the rest of the + process lifetime. + """ + try: + return DatabricksCliTokenSource._probe_cli_version(cli_path) + except Exception as e: + logger.warning(f"Failed to detect Databricks CLI version: {e}. Falling back to conservative flag set.") + return CliVersion() + + @staticmethod + def _resolve_cli_command(cli_path: str, cfg: "Config") -> List[str]: + """Detect CLI version and build the auth-token command. + + Raises IOError with a precise message when no usable command can be + built (missing profile/host, or --profile-unsupported CLI with no host + fallback configured). IOError matches the exception type already + handled by the `databricks_cli()` credential-chain strategy, so + graceful-skip behaviour is preserved. + """ + version = DatabricksCliTokenSource._get_cli_version(cli_path) + if version.is_default_dev_build: + # A default-marker dev build has no injected version, so every + # feature gate fails via at_least. Surface an informational hint so + # users know why their feature flags aren't taking effect. + logger.info( + f"Databricks CLI {version} is a development build; feature detection will use " + "conservative fallbacks. Rebuild the CLI with an explicit version to enable " + "capability-based flag selection." + ) + try: + return DatabricksCliTokenSource._build_cli_command(cli_path, cfg, version) + except IOError as e: + raise IOError(f"cannot configure CLI token source: {e}") from e + + @staticmethod + def _build_cli_command(cli_path: str, cfg: "Config", version: CliVersion) -> List[str]: + """Build the `auth token` command. + + Falls back to --host when --profile is either not configured or not + supported by the installed CLI. Raises IOError describing which + piece of configuration is missing when no command can be built. + """ + if not cfg.profile: + try: + return DatabricksCliTokenSource._build_host_command(cli_path, cfg) + except IOError as e: + raise IOError(f"neither profile nor host is configured: {e}") from e + + # Flag --profile is a global flag and is recognized for all + # commands even the ones that do not support it. Only use --profile + # in CLI versions that are known to support it in `auth token`. + if version < DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE: + # For the unknown (detection failed) and dev-build (no version + # metadata) cases we can't actually prove the CLI lacks --profile; + # we just failed to confirm it. Log accordingly. + if version == CliVersion() or version.is_default_dev_build: + logger.warning( + f"Could not confirm --profile support for Databricks CLI {version} " + f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}). " + "Falling back to --host." + ) + else: + logger.warning( + f"Databricks CLI {version} does not support --profile " + f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}). " + "Falling back to --host." + ) + try: + return DatabricksCliTokenSource._build_host_command(cli_path, cfg) + except IOError as e: + raise IOError( + f"Databricks CLI {version} does not support --profile " + f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}) " + f"and no host fallback is configured: {e}" + ) from e + + return [cli_path, "auth", "token", "--profile", cfg.profile] + @staticmethod - def _build_host_args(cfg: "Config") -> List[str]: - """Build CLI arguments using --host (legacy path).""" - args = ["auth", "token", "--host", cfg.host] - # This is here to support older versions of the Databricks CLI, so we need to keep the client type check. - # This won't work for unified hosts, but it is not supposed to. + def _build_host_command(cli_path: str, cfg: "Config") -> List[str]: + """Build the --host based auth token command. Raises when cfg.host is unset.""" + if not cfg.host: + raise IOError("host is not set") + cmd = [cli_path, "auth", "token", "--host", cfg.host] + # Older Databricks CLIs need --account-id for account-scoped calls; + # unified hosts don't use this path. if cfg.client_type == ClientType.ACCOUNT: - args += ["--account-id", cfg.account_id] - return args + cmd += ["--account-id", cfg.account_id] + return cmd @staticmethod def _find_executable(name) -> str: diff --git a/tests/test_credentials_provider.py b/tests/test_credentials_provider.py index a6eae3018..eb5c0c9dc 100644 --- a/tests/test_credentials_provider.py +++ b/tests/test_credentials_provider.py @@ -306,169 +306,290 @@ def mock_exchange_id_token(id_token: oidc.IdToken): assert headers == {"Authorization": "Bearer exchanged-test-jwt-token"} -# Tests for DatabricksCliTokenSource CLI argument construction -class TestDatabricksCliTokenSourceArgs: - """Tests that DatabricksCliTokenSource constructs correct CLI arguments.""" - - def test_account_client_passes_account_id(self, mocker): - """Non-unified account client should pass --account-id.""" - mock_init = mocker.patch.object( - credentials_provider.CliTokenSource, - "__init__", - return_value=None, - ) +_CV = credentials_provider.CliVersion - mock_cfg = Mock() - mock_cfg.profile = None - mock_cfg.host = "https://accounts.cloud.databricks.com" - - mock_cfg.account_id = "test-account-id" - mock_cfg.client_type = ClientType.ACCOUNT - mock_cfg.databricks_cli_path = "/path/to/databricks" - mock_cfg.disable_async_token_refresh = False - - credentials_provider.DatabricksCliTokenSource(mock_cfg) - - call_kwargs = mock_init.call_args - cmd = call_kwargs.kwargs["cmd"] - - assert "--experimental-is-unified-host" not in cmd - assert "--account-id" in cmd - assert "test-account-id" in cmd - assert "--workspace-id" not in cmd - - def test_profile_uses_profile_flag_with_host_fallback(self, mocker): - """When profile is set, --profile is used as primary and --host as fallback.""" - mock_init = mocker.patch.object( - credentials_provider.CliTokenSource, - "__init__", - return_value=None, - ) - mock_cfg = Mock() - mock_cfg.profile = "my-profile" - mock_cfg.host = "https://workspace.databricks.com" - - mock_cfg.databricks_cli_path = "/path/to/databricks" - mock_cfg.disable_async_token_refresh = False - - credentials_provider.DatabricksCliTokenSource(mock_cfg) - - call_kwargs = mock_init.call_args - cmd = call_kwargs.kwargs["cmd"] - host_cmd = call_kwargs.kwargs["fallback_cmd"] - - assert cmd == ["/path/to/databricks", "auth", "token", "--profile", "my-profile"] - assert host_cmd is not None - assert "--host" in host_cmd - assert "https://workspace.databricks.com" in host_cmd - assert "--profile" not in host_cmd - - def test_profile_without_host_no_fallback(self, mocker): - """When profile is set but host is absent, no fallback is built.""" - mock_init = mocker.patch.object( - credentials_provider.CliTokenSource, - "__init__", - return_value=None, - ) +@pytest.fixture(autouse=True) +def _clear_cli_version_cache(): + # `_probe_cli_version` is `@lru_cache`-decorated, so a value cached by one + # test would leak into the next. Clear before every test. + credentials_provider.DatabricksCliTokenSource._probe_cli_version.cache_clear() - mock_cfg = Mock() - mock_cfg.profile = "my-profile" - mock_cfg.host = None - mock_cfg.databricks_cli_path = "/path/to/databricks" - mock_cfg.disable_async_token_refresh = False - - credentials_provider.DatabricksCliTokenSource(mock_cfg) - - call_kwargs = mock_init.call_args - cmd = call_kwargs.kwargs["cmd"] - host_cmd = call_kwargs.kwargs["fallback_cmd"] - - assert cmd == ["/path/to/databricks", "auth", "token", "--profile", "my-profile"] - assert host_cmd is None - - -# Tests for CliTokenSource fallback on unknown --profile flag -class TestCliTokenSourceFallback: - """Tests that CliTokenSource falls back to --host when CLI doesn't support --profile.""" - - def _make_token_source(self, fallback_cmd=None): - ts = credentials_provider.CliTokenSource.__new__(credentials_provider.CliTokenSource) - ts._cmd = ["databricks", "auth", "token", "--profile", "my-profile"] - ts._fallback_cmd = fallback_cmd - ts._token_type_field = "token_type" - ts._access_token_field = "access_token" - ts._expiry_field = "expiry" - return ts - - def _make_process_error(self, stderr: str, stdout: str = ""): - import subprocess - - err = subprocess.CalledProcessError(1, ["databricks"]) - err.stdout = stdout.encode() - err.stderr = stderr.encode() - return err - - def test_fallback_on_unknown_profile_flag(self, mocker): - """When --profile fails with 'unknown flag: --profile', falls back to --host command.""" - import json - - expiry = (datetime.now() + timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%S") - valid_response = json.dumps({"access_token": "fallback-token", "token_type": "Bearer", "expiry": expiry}) - - mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess") - mock_run.side_effect = [ - self._make_process_error("Error: unknown flag: --profile"), - Mock(stdout=valid_response.encode()), - ] - - fallback_cmd = ["databricks", "auth", "token", "--host", "https://workspace.databricks.com"] - ts = self._make_token_source(fallback_cmd=fallback_cmd) - token = ts.refresh() - assert token.access_token == "fallback-token" - assert mock_run.call_count == 2 - assert mock_run.call_args_list[1][0][0] == fallback_cmd - - def test_fallback_triggered_when_unknown_flag_in_stderr_only(self, mocker): - """Fallback triggers even when CLI also writes usage text to stdout.""" - import json - - expiry = (datetime.now() + timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%S") - valid_response = json.dumps({"access_token": "fallback-token", "token_type": "Bearer", "expiry": expiry}) - - mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess") - mock_run.side_effect = [ - self._make_process_error(stderr="Error: unknown flag: --profile", stdout="Usage: databricks auth token"), - Mock(stdout=valid_response.encode()), - ] - - fallback_cmd = ["databricks", "auth", "token", "--host", "https://workspace.databricks.com"] - ts = self._make_token_source(fallback_cmd=fallback_cmd) - token = ts.refresh() - assert token.access_token == "fallback-token" - - def test_no_fallback_on_real_auth_error(self, mocker): - """When --profile fails with a real error (not unknown flag), no fallback is attempted.""" - mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess") - mock_run.side_effect = self._make_process_error("cache: databricks OAuth is not configured for this host") - - fallback_cmd = ["databricks", "auth", "token", "--host", "https://workspace.databricks.com"] - ts = self._make_token_source(fallback_cmd=fallback_cmd) - with pytest.raises(IOError) as exc_info: - ts.refresh() - assert "databricks OAuth is not configured" in str(exc_info.value) - assert mock_run.call_count == 1 - - def test_no_fallback_when_fallback_cmd_not_set(self, mocker): - """When fallback_cmd is None and --profile fails, the original error is raised.""" - mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess") - mock_run.side_effect = self._make_process_error("Error: unknown flag: --profile") - - ts = self._make_token_source(fallback_cmd=None) - with pytest.raises(IOError) as exc_info: - ts.refresh() - assert "unknown flag: --profile" in str(exc_info.value) - assert mock_run.call_count == 1 + +@pytest.mark.parametrize( + "output,expected", + [ + # Stable releases. + ('{"Major": 0, "Minor": 207, "Patch": 1}', _CV(0, 207, 1)), + ('{"Major": 0, "Minor": 296, "Patch": 0}', _CV(0, 296, 0)), + ('{"Major": 1, "Minor": 2, "Patch": 3}', _CV(1, 2, 3)), + # RC release — we intentionally ignore the prerelease tag; + # the base triple alone is what we gate features on. + ('{"Major": 0, "Minor": 296, "Patch": 0, "Prerelease": "rc.1"}', _CV(0, 296, 0)), + # Nightly snapshot. + ('{"Major": 0, "Minor": 295, "Patch": 1, "Prerelease": "dev"}', _CV(0, 295, 1)), + # Default dev build: numeric fields stay at their "0" defaults when + # the CLI is built without version metadata. (0, 0, 0) is the sentinel. + ('{"Major": 0, "Minor": 0, "Patch": 0}', _CV(0, 0, 0)), + # User-chosen dev version (intentional — treated as v1.0.0). + ('{"Major": 1, "Minor": 0, "Patch": 0, "Prerelease": "dev"}', _CV(1, 0, 0)), + # Full real-world payload with the additional fields the CLI emits. + ( + '{"ProjectName":"cli","Version":"0.295.0","Branch":"HEAD","Tag":"v0.295.0",' + '"Major":0,"Minor":295,"Patch":0,"Prerelease":"","IsSnapshot":false}', + _CV(0, 295, 0), + ), + # Failure cases — all fall back to the unknown CliVersion() (-1,-1,-1). + ("not json", _CV()), + ("", _CV()), + # Old CLIs that don't support --output json emit text — parse fails. + ("Databricks CLI v0.207.1", _CV()), + # Missing a numeric field. + ('{"Minor": 207, "Patch": 1}', _CV()), + # Wrong type on a numeric field. + ('{"Major": "oops", "Minor": 207, "Patch": 1}', _CV()), + ], +) +def test_parse_cli_version(output, expected): + assert credentials_provider.DatabricksCliTokenSource._parse_cli_version(output) == expected + + +@pytest.mark.parametrize( + "a,b,ordering", + [ + (_CV(0, 207, 1), _CV(0, 207, 1), "=="), + (_CV(0, 207, 2), _CV(0, 207, 1), ">"), + (_CV(0, 207, 0), _CV(0, 207, 1), "<"), + (_CV(0, 208, 0), _CV(0, 207, 1), ">"), + (_CV(0, 206, 9), _CV(0, 207, 1), "<"), + (_CV(1, 0, 0), _CV(0, 207, 1), ">"), + (_CV(0, 999, 0), _CV(1, 0, 0), "<"), + (_CV(0, 0, 0), _CV(0, 0, 0), "=="), + (_CV(0, 0, 0), _CV(0, 207, 1), "<"), + # Unknown (-1, -1, -1) compares less than every real version so all + # feature gates fail for it. + (_CV(), _CV(0, 207, 1), "<"), + (_CV(), _CV(0, 0, 0), "<"), + (_CV(), _CV(), "=="), + ], +) +def test_cli_version_total_order(a, b, ordering): + # Lock in all six operators so a future refactor that replaces + # @dataclass(order=True) with a custom __lt__ can't introduce + # asymmetries (e.g. `a > b` True while `b < a` False). + lt, eq, gt = ordering == "<", ordering == "==", ordering == ">" + assert (a < b) is lt + assert (a == b) is eq + assert (a > b) is gt + assert (a <= b) is (lt or eq) + assert (a >= b) is (gt or eq) + assert (a != b) is not eq + + +@pytest.mark.parametrize( + "version,expected", + [ + # Default dev build: the CLI's "no version injected" sentinel. + (_CV(0, 0, 0), True), + # Regular releases. + (_CV(0, 207, 1), False), + (_CV(0, 296, 0), False), + (_CV(1, 0, 0), False), + # Unknown (detection failure) is distinct from the dev-build sentinel. + (_CV(), False), + ], +) +def test_cli_version_is_default_dev_build(version, expected): + assert version.is_default_dev_build is expected + + +_CLI = "/path/to/databricks" +_HOST = "https://workspace.databricks.com" +_ACCT_HOST = "https://accounts.cloud.databricks.com" + + +def _make_cfg(*, profile=None, host=None, account_id=None): + cfg = Mock() + cfg.profile = profile + cfg.host = host + cfg.account_id = account_id + cfg.client_type = ClientType.ACCOUNT if (host and "accounts" in host) else ClientType.WORKSPACE + return cfg + + +@pytest.mark.parametrize( + "name,cfg,version,expected", + [ + ("host only", _make_cfg(host=_HOST), _CV(0, 200, 0), [_CLI, "auth", "token", "--host", _HOST]), + ( + "account host", + _make_cfg(host=_ACCT_HOST, account_id="acct-123"), + _CV(0, 200, 0), + [_CLI, "auth", "token", "--host", _ACCT_HOST, "--account-id", "acct-123"], + ), + ( + "profile with new CLI", + _make_cfg(profile="my-profile", host=_HOST), + _CV(0, 207, 1), + [_CLI, "auth", "token", "--profile", "my-profile"], + ), + ( + "profile with old CLI falls back to host", + _make_cfg(profile="my-profile", host=_HOST), + _CV(0, 200, 0), + [_CLI, "auth", "token", "--host", _HOST], + ), + ( + "unknown version falls back to host", + _make_cfg(profile="my-profile", host=_HOST), + _CV(), + [_CLI, "auth", "token", "--host", _HOST], + ), + ( + "dev-build version falls back to host", + _make_cfg(profile="my-profile", host=_HOST), + _CV(0, 0, 0), + [_CLI, "auth", "token", "--host", _HOST], + ), + ], +) +def test_build_cli_command(name, cfg, version, expected): + assert credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, version) == expected + + +@pytest.mark.parametrize( + "name,cfg,version,match", + [ + ( + "neither profile nor host", + _make_cfg(), + _CV(0, 207, 1), + r"neither profile nor host is configured", + ), + ( + "profile only with old CLI has no host fallback", + _make_cfg(profile="my-profile"), + _CV(0, 200, 0), + r"does not support --profile .* and no host fallback is configured", + ), + ], +) +def test_build_cli_command_errors(name, cfg, version, match): + with pytest.raises(IOError, match=match): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, version) + + +def test_build_cli_command_old_cli_logs_warning(caplog): + import logging + + cfg = _make_cfg(profile="my-profile", host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk"): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, _CV(0, 200, 0)) + assert any("does not support --profile" in rec.message and rec.levelname == "WARNING" for rec in caplog.records) + + +@pytest.mark.parametrize( + "version", + [ + # Detection failed: we don't actually know the CLI lacks --profile. + _CV(), + # Default dev build: no version metadata injected, same story. + _CV(0, 0, 0), + ], +) +def test_build_cli_command_unconfirmed_profile_softens_warning(caplog, version): + import logging + + cfg = _make_cfg(profile="my-profile", host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk"): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, version) + # Softer phrasing for states where --profile support wasn't proven absent. + assert any( + "Could not confirm --profile support" in rec.message and rec.levelname == "WARNING" for rec in caplog.records + ) + assert not any("does not support --profile" in rec.message for rec in caplog.records) + + +def _stub_version_output(mocker, output: str): + """Mock `_run_subprocess` so `_get_cli_version` returns a controlled version.""" + return mocker.patch( + "databricks.sdk.credentials_provider._run_subprocess", + return_value=Mock(stdout=output.encode()), + ) + + +def test_resolve_cli_command_dev_build_logs_info_and_falls_back(mocker, caplog): + import logging + + _stub_version_output( + mocker, + '{"Version": "0.0.0-dev+abcdef123456", "Major": 0, "Minor": 0, "Patch": 0}', + ) + cfg = _make_cfg(profile="my-profile", host=_HOST) + with caplog.at_level(logging.INFO, logger="databricks.sdk.credentials_provider"): + cmd = credentials_provider.DatabricksCliTokenSource._resolve_cli_command(_CLI, cfg) + # Dev build reports as zero version, so --profile is disabled and we fall + # back to --host. + assert cmd == [_CLI, "auth", "token", "--host", _HOST] + assert any("development build" in rec.message and rec.levelname == "INFO" for rec in caplog.records) + + +def test_resolve_cli_command_version_detection_failure_logs_warning(mocker, caplog): + import logging + + mocker.patch( + "databricks.sdk.credentials_provider._run_subprocess", + side_effect=OSError("boom"), + ) + cfg = _make_cfg(host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk.credentials_provider"): + cmd = credentials_provider.DatabricksCliTokenSource._resolve_cli_command(_CLI, cfg) + assert cmd == [_CLI, "auth", "token", "--host", _HOST] + assert any( + "Failed to detect Databricks CLI version" in rec.message and rec.levelname == "WARNING" + for rec in caplog.records + ) + + +def test_get_cli_version_does_not_cache_subprocess_failures(mocker): + # Regression: a transient subprocess failure (timeout, OSError) must not + # be cached. Otherwise a one-off blip pins every later DatabricksCliTokenSource + # to the conservative fallback for the rest of the process lifetime. + mock_run = mocker.patch( + "databricks.sdk.credentials_provider._run_subprocess", + side_effect=[ + OSError("transient"), + Mock(stdout=b'{"Major": 0, "Minor": 207, "Patch": 1}'), + ], + ) + assert credentials_provider.DatabricksCliTokenSource._get_cli_version(_CLI) == _CV() + assert credentials_provider.DatabricksCliTokenSource._get_cli_version(_CLI) == _CV(0, 207, 1) + assert mock_run.call_count == 2 + + +def test_resolve_cli_command_wraps_missing_config_error(mocker): + _stub_version_output( + mocker, + '{"Version": "0.207.1", "Major": 0, "Minor": 207, "Patch": 1}', + ) + cfg = _make_cfg() + with pytest.raises( + IOError, + match=r"cannot configure CLI token source: neither profile nor host is configured", + ): + credentials_provider.DatabricksCliTokenSource._resolve_cli_command(_CLI, cfg) + + +def test_resolve_cli_command_new_cli_uses_profile(mocker): + # Happy path: post-v0.207.1 CLI + profile+host cfg produces a --profile + # command. Exercises the primary code path this PR enables end-to-end. + _stub_version_output( + mocker, + '{"Version": "0.207.1", "Major": 0, "Minor": 207, "Patch": 1}', + ) + cfg = _make_cfg(profile="my-profile", host=_HOST) + cmd = credentials_provider.DatabricksCliTokenSource._resolve_cli_command(_CLI, cfg) + assert cmd == [_CLI, "auth", "token", "--profile", "my-profile"] # Tests for cloud-agnostic hosts and removed cloud checks