Generalize CLI token source into progressive command list

Replace the three explicit command fields (force_cmd, cmd, fallback_cmd)
and manual fallback methods with a CliCommand dataclass and an optional
commands list on CliTokenSource.

When commands is provided, refresh() delegates to _refresh_progressive()
which walks the list from activeCommandIndex, falling back on
unsupported-flag errors. When commands is None, refresh() delegates to
_refresh_single() which preserves the original single-command behavior
with zero changes for AzureCliTokenSource.

DatabricksCliTokenSource._build_commands() produces the progressive
list: --profile + --force-refresh first, plain --profile second, and
--host as a terminal fallback. --force-refresh is only paired with
--profile, never with --host. Adding future flags (e.g. --scopes)
requires only adding entries to _build_commands().

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -12,6 +12,7 @@
 ### Documentation
 
 ### Internal Changes
+* Generalize CLI token source into a progressive command list for forward-compatible flag support.
 
 ### API Changes
 * Add `disable_gov_tag_creation` field for `databricks.sdk.service.settings.RestrictWorkspaceAdminsMessage`.

databricks/sdk/credentials_provider.py

@@ -1,5 +1,6 @@
 import abc
 import base64
+import dataclasses
 import functools
 import io
 import json
@@ -650,6 +651,15 @@ def refreshed_headers() -> Dict[str, str]:
     return OAuthCredentialsProvider(refreshed_headers, token)
 
 
+@dataclasses.dataclass
+class CliCommand:
+    """A single CLI command variant with metadata for progressive fallback."""
+
+    args: List[str]
+    flags: List[str]
+    warning: str
+
+
 class CliTokenSource(oauth.Refreshable):
     _UNKNOWN_FLAG_RE = re.compile(r"unknown flag: (--[a-z-]+)")
 
@@ -660,18 +670,15 @@ def __init__(
         access_token_field: str,
         expiry_field: str,
         disable_async: bool = True,
-        fallback_cmd: Optional[List[str]] = None,
+        commands: Optional[List[CliCommand]] = 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
+        self._commands = commands
+        self._active_command_index = 0
 
     @staticmethod
     def _parse_expiry(expiry: str) -> datetime:
@@ -703,22 +710,34 @@ def _exec_cli_command(self, cmd: List[str]) -> oauth.Token:
             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 _is_unknown_flag_error(error: IOError, flags: List[str]) -> bool:
+        """Check if the error indicates the CLI rejected one of the given flags."""
+        msg = str(error)
+        return any(f"unknown flag: {flag}" in msg for flag in flags)
 
     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
+        if self._commands is not None:
+            return self._refresh_progressive()
+        return self._refresh_single()
+
+    def _refresh_single(self) -> oauth.Token:
+        return self._exec_cli_command(self._cmd)
+
+    def _refresh_progressive(self) -> oauth.Token:
+        last_err: Optional[IOError] = None
+        for i in range(self._active_command_index, len(self._commands)):
+            cmd = self._commands[i]
+            try:
+                token = self._exec_cli_command(cmd.args)
+                self._active_command_index = i
+                return token
+            except IOError as e:
+                is_last = i == len(self._commands) - 1
+                if is_last or not self._is_unknown_flag_error(e, cmd.flags):
+                    raise
+                logger.warning(cmd.warning)
+                last_err = e
+        raise last_err
 
 
 def _run_subprocess(
@@ -907,15 +926,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:
-            args = ["auth", "token", "--profile", cfg.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"]
+        commands = self.__class__._build_commands(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
@@ -925,29 +936,57 @@ def __init__(self, cfg: "Config"):
         self._host = cfg.host
 
         super().__init__(
-            cmd=[cli_path, *args],
+            cmd=commands[-1].args,
             token_type_field="token_type",
             access_token_field="access_token",
             expiry_field="expiry",
             disable_async=cfg.disable_async_token_refresh,
-            fallback_cmd=fallback_cmd,
+            commands=commands,
         )
 
-    _KNOWN_CLI_FLAGS = {"--force-refresh", "--profile"}
+    @staticmethod
+    def _build_commands(cli_path: str, cfg: "Config") -> List[CliCommand]:
+        commands: List[CliCommand] = []
+        if cfg.profile:
+            profile_args = [cli_path, "auth", "token", "--profile", cfg.profile]
+            commands.append(
+                CliCommand(
+                    args=profile_args + ["--force-refresh"],
+                    flags=["--force-refresh", "--profile"],
+                    warning="Databricks CLI does not support --force-refresh. "
+                    "Please upgrade your CLI to the latest version.",
+                )
+            )
+            commands.append(
+                CliCommand(
+                    args=profile_args,
+                    flags=["--profile"],
+                    warning="Databricks CLI does not support --profile flag. "
+                    "Falling back to --host. "
+                    "Please upgrade your CLI to the latest version.",
+                )
+            )
+            if cfg.host:
+                commands.append(
+                    CliCommand(
+                        args=[cli_path, *DatabricksCliTokenSource._build_host_args(cfg)],
+                        flags=[],
+                        warning="",
+                    )
+                )
+        else:
+            host_args = [cli_path, *DatabricksCliTokenSource._build_host_args(cfg)]
+            commands.append(
+                CliCommand(
+                    args=host_args,
+                    flags=[],
+                    warning="",
+                )
+            )
+        return commands
 
     def refresh(self) -> oauth.Token:
-        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
+        token = super().refresh()
         if self._requested_scopes:
             self._validate_token_scopes(token)
         return token