Fix CLI token source --profile fallback with version detection

[mihaimitrea-db]

🥞 Stacked PR

Use this link to review incremental changes.

• stack/cli-force-refresh [Files changed]
  • stack/cli-progressive-token-commands [Files changed]

---

Summary

Replace the broken error-based --profile fallback in CliTokenSource with version-based CLI detection at init time. Mirrors databricks/databricks-sdk-go#1605.

Why

--profile on databricks auth token is a global Cobra flag, so old CLIs (< v0.207.1) silently accept it and then fail with "cannot fetch credentials" instead of "unknown flag: --profile". The existing retry check was matching on the latter and never fired — the --host fallback it gated was effectively dead code. Switching to databricks version + a minimum-version constant makes the fallback reliable and sets up future capability-gated flags (e.g. --force-refresh in #1378) without additional subprocess calls.

What changed

Interface changes

None. CliTokenSource and DatabricksCliTokenSource are not part of the public API surface.

Behavioral changes

• cfg.profile + CLI < v0.207.1 now correctly falls back to --host (previously broken).
• databricks version failures log a WARNING and fall back to the most conservative command.
• A default dev build (v0.0.0-dev[+commit]) logs an INFO explaining why feature gates are conservative.
• Error messages for missing configuration are precise ("neither profile nor host is configured" / "…does not support --profile … and no host fallback is configured") instead of a single generic message.

AzureCliTokenSource is untouched.

Internal changes

• New CliVersion dataclass with semver-aware comparison operators and an is_default_dev_build property.
• CliTokenSource simplified to a single cmd; the fallback_cmd parameter and its retry logic are removed.
• DatabricksCliTokenSource gains _get_cli_version, _resolve_cli_command, _build_cli_command, and _build_host_command helpers.

How is this tested?

Unit tests in tests/test_credentials_provider.py cover version parsing, comparison operators, dev-build detection, command assembly for every profile/host/version combination, precise error messages, and logging behavior on detection failure and dev builds. Existing tests/test_core.py::test_databricks_cli_token_source_* tests pass unchanged.

[mihaimitrea-db]

Range-diff: main (1d03f6d -> 9e82d18)

NEXT_CHANGELOG.md

@@ -0,0 +1,10 @@
+diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
+--- a/NEXT_CHANGELOG.md
++++ b/NEXT_CHANGELOG.md
+ ## Release v0.102.0
+ 
+ ### New Features and Improvements
++* Pass `--force-refresh` to the Databricks CLI `auth token` command so the SDK always receives a freshly minted token instead of a potentially stale cached one. Falls back gracefully on older CLIs that do not support the flag.
+ 
+ ### Security
+ 
\ No newline at end of file

databricks/sdk/credentials_provider.py

@@ -63,8 +63,7 @@
 +            flag = self._get_unsupported_flag(e)
 +            if flag in self._KNOWN_CLI_FLAGS:
 +                logger.warning(
-+                    "Databricks CLI does not support %s. "
-+                    "Please upgrade your CLI to the latest version.",
++                    "Databricks CLI does not support %s. " "Please upgrade your CLI to the latest version.",
 +                    flag,
 +                )
 +                token = super().refresh()

pyrefly.toml

@@ -0,0 +1,15 @@
+diff --git a/pyrefly.toml b/pyrefly.toml
+new file mode 100644
+--- /dev/null
++++ b/pyrefly.toml
++project-includes = [
++   "**/*.py"
++]
++
++project-excludes = []
++
++search-path = []
++
++disable-search-path-heuristics = true
++ignore-missing-imports = ["*"]
++ignore-errors-in-generated-code = true
\ No newline at end of file

tests/__init__.pyc

@@ -0,0 +1,3 @@
+diff --git a/tests/__init__.pyc b/tests/__init__.pyc
+new file mode 100644
+Binary files /dev/null and b/tests/__init__.pyc differ
\ No newline at end of file

tests/test_credentials_provider.py

@@ -121,9 +121,7 @@
 +        ts = self._make_token_source()
 +
 +        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"
-+        )
++        mock_run.side_effect = self._make_process_error("cache: databricks OAuth is not configured for this host")
 +
 +        with pytest.raises(IOError) as exc_info:
 +            ts.refresh()

_{Reproduce locally: git range-diff 34d6184..1d03f6d 34d6184..9e82d18 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (9e82d18 -> 3195f5d)

NEXT_CHANGELOG.md

@@ -1,9 +1,9 @@
 diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
 --- a/NEXT_CHANGELOG.md
 +++ b/NEXT_CHANGELOG.md
- ## Release v0.102.0
- 
  ### New Features and Improvements
+ * Add support for unified hosts. A single configuration profile can now be used for both account-level and workspace-level operations when the host supports it and both `account_id` and `workspace_id` are available. The `experimental_is_unified_host` flag has been removed; unified host detection is now automatic.
+ * Accept `DATABRICKS_OIDC_TOKEN_FILEPATH` environment variable for consistency with other Databricks SDKs (Go, CLI, Terraform). The previous `DATABRICKS_OIDC_TOKEN_FILE` is still supported as an alias.
 +* Pass `--force-refresh` to the Databricks CLI `auth token` command so the SDK always receives a freshly minted token instead of a potentially stale cached one. Falls back gracefully on older CLIs that do not support the flag.
  
  ### Security

_{Reproduce locally: git range-diff 34d6184..9e82d18 78914fa..3195f5d | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (3195f5d -> 4115f37)

databricks/sdk/credentials_provider.py

@@ -1,34 +1,14 @@
 diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py
 --- a/databricks/sdk/credentials_provider.py
 +++ b/databricks/sdk/credentials_provider.py
- import os
- import pathlib
- import platform
-+import re
- import subprocess
- import sys
- import threading
  
  
  class CliTokenSource(oauth.Refreshable):
-+    _UNKNOWN_FLAG_RE = re.compile(r"unknown flag: (--[a-z-]+)")
 +
      def __init__(
          self,
          cmd: List[str],
-             message = "\n".join(filter(None, [stdout, stderr]))
-             raise IOError(f"cannot get access token: {message}") from e
  
-+    @staticmethod
-+    def _get_unsupported_flag(error: IOError) -> Optional[str]:
-+        """Extract the flag name if the error is an 'unknown flag' CLI rejection."""
-+        match = CliTokenSource._UNKNOWN_FLAG_RE.search(str(error))
-+        return match.group(1) if match else None
-+
-     def refresh(self) -> oauth.Token:
-         try:
-             return self._exec_cli_command(self._cmd)
- 
          fallback_cmd = None
          if cfg.profile:
 -            # When profile is set, use --profile as the primary command.
@@ -45,11 +25,8 @@
          # get_scopes() defaults to ["all-apis"] when nothing is configured, which would
          # cause false-positive mismatches against every token that wasn't issued with
          # exactly ["all-apis"]. Only validate when scopes are explicitly set (either
-             fallback_cmd=fallback_cmd,
          )
  
-+    _KNOWN_CLI_FLAGS = {"--force-refresh", "--profile"}
-+
      def refresh(self) -> oauth.Token:
 -        # The scope validation lives in refresh() because this is the only method that
 -        # produces new tokens (see Refreshable._token assignments). By overriding here,
@@ -60,11 +37,11 @@
 +        try:
 +            token = self._exec_cli_command(self._force_cmd)
 +        except IOError as e:
-+            flag = self._get_unsupported_flag(e)
-+            if flag in self._KNOWN_CLI_FLAGS:
++            err_msg = str(e)
++            if "unknown flag: --force-refresh" in err_msg or "unknown flag: --profile" in err_msg:
 +                logger.warning(
-+                    "Databricks CLI does not support %s. " "Please upgrade your CLI to the latest version.",
-+                    flag,
++                    "Databricks CLI does not support --force-refresh. "
++                    "Please upgrade your CLI to the latest version."
 +                )
 +                token = super().refresh()
 +            else:

tests/test_credentials_provider.py

@@ -128,13 +128,6 @@
 +        assert "databricks OAuth is not configured" in str(exc_info.value)
 +        assert mock_run.call_count == 1
 +
-+    def test_get_unsupported_flag_extracts_flag(self):
-+        """The classifier correctly parses the flag name from CLI error output."""
-+        get = credentials_provider.CliTokenSource._get_unsupported_flag
-+        assert get(IOError("Error: unknown flag: --force-refresh")) == "--force-refresh"
-+        assert get(IOError("Error: unknown flag: --profile")) == "--profile"
-+        assert get(IOError("some other error")) is None
-+
 +
  # Tests for cloud-agnostic hosts and removed cloud checks
  class TestCloudAgnosticHosts:

_{Reproduce locally: git range-diff 78914fa..3195f5d 78914fa..4115f37 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (4115f37 -> 469cb44)

pyrefly.toml

@@ -1,15 +0,0 @@
-diff --git a/pyrefly.toml b/pyrefly.toml
-new file mode 100644
---- /dev/null
-+++ b/pyrefly.toml
-+project-includes = [
-+   "**/*.py"
-+]
-+
-+project-excludes = []
-+
-+search-path = []
-+
-+disable-search-path-heuristics = true
-+ignore-missing-imports = ["*"]
-+ignore-errors-in-generated-code = true
\ No newline at end of file

tests/__init__.pyc

@@ -1,3 +0,0 @@
-diff --git a/tests/__init__.pyc b/tests/__init__.pyc
-new file mode 100644
-Binary files /dev/null and b/tests/__init__.pyc differ
\ No newline at end of file

_{Reproduce locally: git range-diff 78914fa..4115f37 78914fa..469cb44 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (469cb44 -> 32c2cbd)

databricks/sdk/credentials_provider.py

@@ -8,23 +8,19 @@
      def __init__(
          self,
          cmd: List[str],
+             cli_path = self.__class__._find_executable(cli_path)
  
          fallback_cmd = None
++        self._force_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.
++            self._force_cmd = [cli_path, *args, "--force-refresh"]
              if cfg.host:
                  fallback_cmd = [cli_path, *self.__class__._build_host_args(cfg)]
          else:
-             args = self.__class__._build_host_args(cfg)
- 
-+        self._force_cmd = [cli_path, *args, "--force-refresh"]
-+
-         # get_scopes() defaults to ["all-apis"] when nothing is configured, which would
-         # cause false-positive mismatches against every token that wasn't issued with
-         # exactly ["all-apis"]. Only validate when scopes are explicitly set (either
          )
  
      def refresh(self) -> oauth.Token:
@@ -34,18 +30,21 @@
 -        # when the cached token expires. This catches cases where a user re-authenticates
 -        # mid-session with different scopes.
 -        token = super().refresh()
-+        try:
-+            token = self._exec_cli_command(self._force_cmd)
-+        except IOError as e:
-+            err_msg = str(e)
-+            if "unknown flag: --force-refresh" in err_msg or "unknown flag: --profile" in err_msg:
-+                logger.warning(
-+                    "Databricks CLI does not support --force-refresh. "
-+                    "Please upgrade your CLI to the latest version."
-+                )
-+                token = super().refresh()
-+            else:
-+                raise
++        if self._force_cmd is None:
++            token = super().refresh()
++        else:
++            try:
++                token = self._exec_cli_command(self._force_cmd)
++            except IOError as e:
++                err_msg = str(e)
++                if "unknown flag: --force-refresh" in err_msg or "unknown flag: --profile" in err_msg:
++                    logger.warning(
++                        "Databricks CLI does not support --force-refresh. "
++                        "Please upgrade your CLI to the latest version."
++                    )
++                    token = super().refresh()
++                else:
++                    raise
          if self._requested_scopes:
              self._validate_token_scopes(token)
          return token
\ No newline at end of file

pyrefly.toml

@@ -0,0 +1,15 @@
+diff --git a/pyrefly.toml b/pyrefly.toml
+new file mode 100644
+--- /dev/null
++++ b/pyrefly.toml
++project-includes = [
++   "**/*.py"
++]
++
++project-excludes = []
++
++search-path = []
++
++disable-search-path-heuristics = true
++ignore-missing-imports = ["*"]
++ignore-errors-in-generated-code = true
\ No newline at end of file

tests/__init__.pyc

@@ -0,0 +1,3 @@
+diff --git a/tests/__init__.pyc b/tests/__init__.pyc
+new file mode 100644
+Binary files /dev/null and b/tests/__init__.pyc differ
\ No newline at end of file

tests/test_credentials_provider.py

@@ -41,9 +41,9 @@
 +        expiry = (datetime.now() + timedelta(hours=1)).strftime("%Y-%m-%dT%H:%M:%S")
 +        return json.dumps({"access_token": access_token, "token_type": "Bearer", "expiry": expiry})
 +
-+    def test_force_refresh_always_tried_first(self, mocker):
-+        """refresh() always tries --force-refresh first."""
-+        ts = self._make_token_source()
++    def test_force_refresh_tried_first_with_profile(self, mocker):
++        """When profile is configured, refresh() tries --force-refresh first."""
++        ts = self._make_token_source(profile="my-profile")
 +
 +        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
 +        mock_run.return_value = Mock(stdout=self._valid_response_json("refreshed").encode())
@@ -54,10 +54,26 @@
 +
 +        cmd = mock_run.call_args[0][0]
 +        assert "--force-refresh" in cmd
++        assert "--profile" in cmd
++
++    def test_host_only_skips_force_refresh(self, mocker):
++        """When only host is configured, --force-refresh is not used."""
++        ts = self._make_token_source()
++
++        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
++        mock_run.return_value = Mock(stdout=self._valid_response_json("token").encode())
++
++        token = ts.refresh()
++        assert token.access_token == "token"
++        assert mock_run.call_count == 1
 +
++        cmd = mock_run.call_args[0][0]
++        assert "--force-refresh" not in cmd
++        assert "--host" in cmd
++
 +    def test_force_refresh_fallback_when_unsupported(self, mocker):
 +        """Old CLI without --force-refresh: falls back to cmd without --force-refresh."""
-+        ts = self._make_token_source()
++        ts = self._make_token_source(profile="my-profile")
 +
 +        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
 +        mock_run.side_effect = [

tests/test_config.py

@@ -1,11 +0,0 @@
-diff --git a/tests/test_config.py b/tests/test_config.py
---- a/tests/test_config.py
-+++ b/tests/test_config.py
- 
- def test_config_deep_copy(monkeypatch, mocker, tmp_path):
-     mocker.patch(
--        "databricks.sdk.credentials_provider.CliTokenSource.refresh",
-+        "databricks.sdk.credentials_provider.CliTokenSource._exec_cli_command",
-         return_value=oauth.Token(
-             access_token="token",
-             token_type="Bearer",
\ No newline at end of file

tests/test_core.py

@@ -1,27 +0,0 @@
-diff --git a/tests/test_core.py b/tests/test_core.py
---- a/tests/test_core.py
-+++ b/tests/test_core.py
- 
- def test_databricks_cli_credential_provider_installed_new(config, monkeypatch, tmp_path, mocker):
-     get_mock = mocker.patch(
--        "databricks.sdk.credentials_provider.CliTokenSource.refresh",
-+        "databricks.sdk.credentials_provider.CliTokenSource._exec_cli_command",
-         return_value=Token(
-             access_token="token",
-             token_type="Bearer",
-     config, monkeypatch, tmp_path, mocker, token_claims, configured_scopes, auth_type, expect
- ):
-     mocker.patch(
--        "databricks.sdk.credentials_provider.CliTokenSource.refresh",
-+        "databricks.sdk.credentials_provider.CliTokenSource._exec_cli_command",
-         return_value=Token(access_token=_make_jwt(token_claims), token_type="Bearer", expiry=datetime(2023, 5, 22)),
-     )
-     write_large_dummy_executable(tmp_path)
- 
- def test_databricks_cli_scope_validation_error_message(config, monkeypatch, tmp_path, mocker):
-     mocker.patch(
--        "databricks.sdk.credentials_provider.CliTokenSource.refresh",
-+        "databricks.sdk.credentials_provider.CliTokenSource._exec_cli_command",
-         return_value=Token(
-             access_token=_make_jwt({"scope": "all-apis"}), token_type="Bearer", expiry=datetime(2023, 5, 22)
-         ),
\ No newline at end of file

_{Reproduce locally: git range-diff 78914fa..469cb44 78914fa..32c2cbd | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (32c2cbd -> 194256c)

pyrefly.toml

@@ -1,15 +0,0 @@
-diff --git a/pyrefly.toml b/pyrefly.toml
-new file mode 100644
---- /dev/null
-+++ b/pyrefly.toml
-+project-includes = [
-+   "**/*.py"
-+]
-+
-+project-excludes = []
-+
-+search-path = []
-+
-+disable-search-path-heuristics = true
-+ignore-missing-imports = ["*"]
-+ignore-errors-in-generated-code = true
\ No newline at end of file

tests/__init__.pyc

@@ -1,3 +0,0 @@
-diff --git a/tests/__init__.pyc b/tests/__init__.pyc
-new file mode 100644
-Binary files /dev/null and b/tests/__init__.pyc differ
\ No newline at end of file

_{Reproduce locally: git range-diff 78914fa..32c2cbd 78914fa..194256c | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: main (194256c -> cd6c876)

NEXT_CHANGELOG.md

@@ -4,7 +4,7 @@
  ### New Features and Improvements
  * Add support for unified hosts. A single configuration profile can now be used for both account-level and workspace-level operations when the host supports it and both `account_id` and `workspace_id` are available. The `experimental_is_unified_host` flag has been removed; unified host detection is now automatic.
  * Accept `DATABRICKS_OIDC_TOKEN_FILEPATH` environment variable for consistency with other Databricks SDKs (Go, CLI, Terraform). The previous `DATABRICKS_OIDC_TOKEN_FILE` is still supported as an alias.
-+* Pass `--force-refresh` to the Databricks CLI `auth token` command so the SDK always receives a freshly minted token instead of a potentially stale cached one. Falls back gracefully on older CLIs that do not support the flag.
++* Pass `--force-refresh` to the Databricks CLI `auth token` command so the SDK always receives a freshly minted token instead of a potentially stale cached one.
  
  ### Security
  
\ No newline at end of file

_{Reproduce locally: git range-diff 78914fa..194256c 78914fa..cd6c876 | Disable: git config gitstack.push-range-diff false}

[simonfaltum]

Review Swarm: Isaac + Cursor (2-round deep review)

Verdict: Not ready yet. 0 Critical | 3 Major | 2 Gap | 5 Nit | 3 Suggestion

Two independent reviewers (Isaac, Cursor) reviewed this PR, cross-checked each other's findings, and a synthesis follows. Inline comments are attached to specific lines below; this top-level body highlights the priorities and disputes.

Top priorities (block merge)

1. Unnecessary subprocess on host-only path (Major, both reviewers) — _resolve_cli_command unconditionally calls _get_cli_version, but the version is ignored when cfg.profile is unset. The common host-only config pays a fork/exec + possible WARNING for a value it never uses. Short-circuit the host-only case before version detection.
2. No subprocess timeout on databricks version (Major, Isaac + confirmed by Cursor) — _get_cli_version calls _run_subprocess(..., check=True) with the default timeout=None. If the CLI hangs (slow disk, AV scan, stdin wait, network in a subcommand), SDK init blocks forever. Pass timeout=5.
3. Happy-path test missing (Gap/Major, Isaac + confirmed by Cursor) — tests exist for dev-build, detection failure, old-CLI fallback, and error cases, but none stubs a successful post-v0.207.1 version probe and asserts the resulting command includes --profile. Easy to regress silently.

Other Majors

4. Per-instance subprocess; no caching — _get_cli_version re-spawns on every DatabricksCliTokenSource construction. Immutable for a given cli_path; wrap with functools.lru_cache. (Severity disagreement: Isaac Major, Cursor Suggestion; taking higher.)
5. Disputed: init-time exception type change — __init__ now raises ValueError where it previously surfaced as IOError at refresh(). The databricks_cli() credential chain catches IOError and a narrow ValueError; the new ValueError from _resolve_cli_command could skip graceful-skip. Cursor noted @oauth_credentials_strategy("databricks-cli", ["host"]) limits reachability. At minimum, add a NEXT_CHANGELOG.md note.

Nits

• Fallback WARNING overstates certainty for unknown / v0.0.0 states — both reviewers.
• CliVersion.__str__ renders the dev-build sentinel as v0.0.0, ambiguous in logs.
• Adjacent string-literal concatenation in the version-detect WARNING reads like a typo (disputed).
• Forward reference to module-level _DEFAULT_DEV_BUILD from inside CliVersion (disputed, stylistic).
• No caplog assertion on the old-CLI fallback WARNING, though the parallel dev-build/detection-failure tests have one.

Suggestions

• Extensibility: declare a list of (min_version, args) tuples instead of nested-if for capability gates. Pairs naturally with stacked PR #1378.
• Broaden test_cli_version_ge into a total-order test covering <, <=, == to lock in the dataclass(order=True) contract.
• Consider lazy-at-refresh detection with caching (disputed; current init-time approach mirrors the Go SDK and is defensible).

Overall

The design is sound and mirrors databricks/databricks-sdk-go#1605. The test restructuring from class-based to parametrized is a clear improvement. No correctness bugs. The Majors are hot-path regressions and init-time robustness; fixing those (plus the happy-path test) makes this ready to merge.

[simonfaltum]

Take a look at my comment, otherwise looks good to me

[simonfaltum]

_get_cli_version caches transient failures for the full process lifetime. If databricks version times out or hits a temporary OSError once, the unknown CliVersion() sentinel is cached for that CLI path. Later DatabricksCliTokenSource instances never retry, so a supported CLI can keep falling back to --host, or fail with “no host fallback is configured,” until process restart.

Suggested fix: cache only successful detections, or keep the failure handling outside the cached helper. Add a regression test where the first probe fails and the second probe returns v0.207.1+, then assert the second source uses --profile.

[github-actions]

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-py

Inputs:

• PR number: 1377
• Commit SHA: c714a135fb7e9a105b78baa73f203a3875903c78

Checks will be approved automatically on success.