Add best-effort --force-refresh support for databricks-cli auth

When the SDK's cached CLI token is stale, try `databricks auth token
--force-refresh` to get a freshly minted token from the IdP. If the
installed CLI is too old to recognise the flag, fall back to regular
`auth token` and remember the capability for future refreshes.

Centralise unknown-flag detection in CliTokenSource._exec_cli_command()
via UnsupportedCliFlagError so the same classifier is reused by both the
legacy --profile fallback and the new --force-refresh downgrade path in
DatabricksCliTokenSource.

See: databricks/cli#4767

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 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
 

databricks/sdk/credentials_provider.py

@@ -7,6 +7,7 @@
 import os
 import pathlib
 import platform
+import re
 import subprocess
 import sys
 import threading
@@ -650,6 +651,8 @@ def refreshed_headers() -> Dict[str, str]:
 
 
 class CliTokenSource(oauth.Refreshable):
+    _UNKNOWN_FLAG_RE = re.compile(r"unknown flag: (--[a-z-]+)")
+
     def __init__(
         self,
         cmd: List[str],
@@ -699,6 +702,12 @@ def _exec_cli_command(self, cmd: List[str]) -> oauth.Token:
             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)
@@ -900,15 +909,14 @@ def __init__(self, cfg: "Config"):
 
         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)
 
+        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
@@ -925,13 +933,21 @@ def __init__(self, cfg: "Config"):
             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,
-        # every token is validated — both at initial auth and on every subsequent refresh
-        # 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:
+            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.",
+                    flag,
+                )
+                token = super().refresh()
+            else:
+                raise
         if self._requested_scopes:
             self._validate_token_scopes(token)
         return token

pyrefly.toml

@@ -0,0 +1,11 @@
+project-includes = [
+   "**/*.py"
+]
+
+project-excludes = []
+
+search-path = []
+
+disable-search-path-heuristics = true
+ignore-missing-imports = ["*"]
+ignore-errors-in-generated-code = true

tests/__init__.pyc

@@ -117,7 +117,7 @@ def test_config_copy_deep_copies_user_agent_other_info(config):
 
 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",

tests/test_config.py

@@ -166,7 +166,7 @@ def test_databricks_cli_credential_provider_installed_legacy(config, monkeypatch
 
 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",
@@ -222,7 +222,7 @@ def test_databricks_cli_scope_validation(
     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)
@@ -244,7 +244,7 @@ def test_databricks_cli_scope_validation(
 
 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)
         ),

tests/test_core.py

@@ -471,6 +471,138 @@ def test_no_fallback_when_fallback_cmd_not_set(self, mocker):
         assert mock_run.call_count == 1
 
 
+class TestDatabricksCliForceRefresh:
+    """Tests for --force-refresh support in DatabricksCliTokenSource."""
+
+    @staticmethod
+    def _make_process_error(stderr: str, stdout: str = ""):
+        import subprocess
+
+        err = subprocess.CalledProcessError(1, ["databricks"])
+        err.stdout = stdout.encode()
+        err.stderr = stderr.encode()
+        return err
+
+    @staticmethod
+    def _make_token_source(
+        *,
+        profile=None,
+        host="https://workspace.databricks.com",
+        cli_path="/path/to/databricks",
+    ):
+        """Build a DatabricksCliTokenSource by mocking only the executable check."""
+        mock_cfg = Mock()
+        mock_cfg.profile = profile
+        mock_cfg.host = host
+        mock_cfg.databricks_cli_path = cli_path
+        mock_cfg.disable_async_token_refresh = True
+        mock_cfg.scopes = None
+        mock_cfg.get_scopes = Mock(return_value=["all-apis"])
+        mock_cfg.client_type = ClientType.WORKSPACE
+        mock_cfg.account_id = None
+        return credentials_provider.DatabricksCliTokenSource(mock_cfg)
+
+    def _valid_response_json(self, access_token="fresh-token"):
+        import json
+
+        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()
+
+        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
+        mock_run.return_value = Mock(stdout=self._valid_response_json("refreshed").encode())
+
+        token = ts.refresh()
+        assert token.access_token == "refreshed"
+        assert mock_run.call_count == 1
+
+        cmd = mock_run.call_args[0][0]
+        assert "--force-refresh" 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()
+
+        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
+        mock_run.side_effect = [
+            self._make_process_error("Error: unknown flag: --force-refresh"),
+            Mock(stdout=self._valid_response_json("fallback").encode()),
+        ]
+
+        token = ts.refresh()
+        assert token.access_token == "fallback"
+        assert mock_run.call_count == 2
+
+        first_cmd = mock_run.call_args_list[0][0][0]
+        second_cmd = mock_run.call_args_list[1][0][0]
+        assert "--force-refresh" in first_cmd
+        assert "--force-refresh" not in second_cmd
+
+    def test_profile_fallback_when_unsupported(self, mocker):
+        """Old CLI without --profile: force_cmd fails, fallback retries with --host."""
+        ts = self._make_token_source(profile="my-profile")
+
+        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
+        mock_run.side_effect = [
+            # force_cmd: --profile + --force-refresh → unknown --profile
+            self._make_process_error("Error: unknown flag: --profile"),
+            # _refresh_without_force cmd: --profile → unknown --profile
+            self._make_process_error("Error: unknown flag: --profile"),
+            # _refresh_without_force fallback_cmd: --host → success
+            Mock(stdout=self._valid_response_json("host-token").encode()),
+        ]
+
+        token = ts.refresh()
+        assert token.access_token == "host-token"
+        assert mock_run.call_count == 3
+        assert "--host" in mock_run.call_args_list[2][0][0]
+
+    def test_two_step_downgrade_both_flags_unsupported(self, mocker):
+        """CLI supports neither --force-refresh nor --profile: force_cmd fails, then full fallback."""
+        ts = self._make_token_source(profile="my-profile")
+
+        mock_run = mocker.patch("databricks.sdk.credentials_provider._run_subprocess")
+        mock_run.side_effect = [
+            # 1st: force_cmd (--profile + --force-refresh) → unknown --force-refresh
+            self._make_process_error("Error: unknown flag: --force-refresh"),
+            # 2nd: _refresh_without_force cmd (--profile) → unknown --profile
+            self._make_process_error("Error: unknown flag: --profile"),
+            # 3rd: _refresh_without_force fallback_cmd (--host) → success
+            Mock(stdout=self._valid_response_json("plain").encode()),
+        ]
+
+        token = ts.refresh()
+        assert token.access_token == "plain"
+        assert mock_run.call_count == 3
+
+        cmds = [call[0][0] for call in mock_run.call_args_list]
+        assert "--force-refresh" in cmds[0] and "--profile" in cmds[0]
+        assert "--force-refresh" not in cmds[1] and "--profile" in cmds[1]
+        assert "--host" in cmds[2]
+
+    def test_real_auth_error_does_not_trigger_fallback(self, mocker):
+        """Real auth failures (not unknown-flag) surface immediately."""
+        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")
+
+        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_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:
     """Tests that credential providers work with cloud-agnostic hosts after removing is_azure/is_gcp checks."""