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: databricks/cli#855

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -7,11 +7,13 @@
 ### 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 (e.g. `--force-refresh`) without additional subprocess calls.
 
 ### API Changes

databricks/sdk/credentials_provider.py

@@ -1,12 +1,14 @@
 import abc
 import base64
+import dataclasses
 import functools
 import io
 import json
 import logging
 import os
 import pathlib
 import platform
+import re
 import subprocess
 import sys
 import threading
@@ -662,23 +664,103 @@ def refreshed_headers() -> Dict[str, str]:
     return OAuthCredentialsProvider(refreshed_headers, token)
 
 
+@dataclasses.dataclass
+class CliVersion:
+    """Semver-like version of the Databricks CLI.
+
+    A zero value (0, 0, 0) with empty `suffix` means version detection failed
+    and the SDK should fall back to the most conservative command. The `suffix`
+    field preserves any prerelease tag and build metadata (e.g. "-dev+sha") so
+    dev builds can be identified and compared correctly against releases.
+    """
+
+    major: int = 0
+    minor: int = 0
+    patch: int = 0
+    # e.g. "-dev+abcdef", "-rc.1", "+build.42", or "" for a stable release.
+    suffix: str = ""
+
+    @property
+    def is_default_dev_build(self) -> bool:
+        """True when this is the default Databricks CLI dev build (v0.0.0-dev...).
+
+        Narrowly matches the CLI's "no version injected" marker: base triple
+        equal to the default dev build's and suffix starting with the default's
+        "-dev" prerelease tag (build metadata like "+commit" is tolerated).
+
+        A version the user explicitly set (e.g. v1.0.0-dev) is intentional and
+        is not flagged here; feature gates still apply via at_least.
+        """
+        return (
+            (self.major, self.minor, self.patch)
+            == (_DEFAULT_DEV_BUILD.major, _DEFAULT_DEV_BUILD.minor, _DEFAULT_DEV_BUILD.patch)
+            and self.suffix.startswith(_DEFAULT_DEV_BUILD.suffix)
+        )
+
+    @property
+    def _has_prerelease(self) -> bool:
+        # In semver, a leading "-" introduces a prerelease tag that compares
+        # less than the same base version without one. A leading "+" is build
+        # metadata and does not affect ordering.
+        return self.suffix.startswith("-")
+
+    def _cmp_key(self) -> Tuple[int, int, int, int]:
+        """Tuple-orderable key for semver-aware comparison.
+
+        Semver puts "X.Y.Z-anything" before "X.Y.Z" (prerelease sorts before
+        release). We encode that by ranking prereleases as 0 and the absence
+        of a prerelease (release or build-metadata-only) as 1 in the last
+        slot, so normal tuple ordering gives semver ordering. Two prereleases
+        of the same base compare as equal under this key even if their tags
+        differ, which is a simplification suited to our use (release-valued
+        feature-gate thresholds).
+        """
+        return (self.major, self.minor, self.patch, 0 if self._has_prerelease else 1)
+
+    def __lt__(self, other: "CliVersion") -> bool:
+        if not isinstance(other, CliVersion):
+            return NotImplemented
+        return self._cmp_key() < other._cmp_key()
+
+    def __le__(self, other: "CliVersion") -> bool:
+        if not isinstance(other, CliVersion):
+            return NotImplemented
+        return self._cmp_key() <= other._cmp_key()
+
+    def __gt__(self, other: "CliVersion") -> bool:
+        if not isinstance(other, CliVersion):
+            return NotImplemented
+        return self._cmp_key() > other._cmp_key()
+
+    def __ge__(self, other: "CliVersion") -> bool:
+        if not isinstance(other, CliVersion):
+            return NotImplemented
+        return self._cmp_key() >= other._cmp_key()
+
+    def __str__(self) -> str:
+        if (self.major, self.minor, self.patch) == (0, 0, 0) and not self.suffix:
+            return "unknown"
+        return f"v{self.major}.{self.minor}.{self.patch}{self.suffix}"
+
+
+# The Databricks CLI emits exactly "v0.0.0-dev[+commit]" when the binary was
+# built without injecting a version via ldflags (e.g. a plain `go build`). See
+# https://github.com/databricks/cli/blob/main/internal/build/info.go.
+_DEFAULT_DEV_BUILD = CliVersion(0, 0, 0, suffix="-dev")
+
+
 class CliTokenSource(oauth.Refreshable):
+
     def __init__(
         self,
         cmd: List[str],
         token_type_field: str,
         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 +795,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 +984,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 +994,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 +1056,116 @@ 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)
+
+    # Captures "vX.Y.Z" plus any prerelease/build-metadata suffix so dev
+    # builds like "v0.0.0-dev+abcdef" can be distinguished from detection
+    # failure and compared against release versions.
+    _CLI_VERSION_RE = re.compile(r"Databricks CLI v(\d+)\.(\d+)\.(\d+)(\S*)")
+
+    @staticmethod
+    def _parse_cli_version(output: str) -> CliVersion:
+        """Parse 'Databricks CLI vX.Y.Z[-suffix]' output into a CliVersion.
+
+        Returns CliVersion() on failure, matching Go's behavior where an
+        unparseable version disables every feature gate.
+        """
+        m = DatabricksCliTokenSource._CLI_VERSION_RE.search(output)
+        if not m:
+            logger.debug(f"Failed to parse Databricks CLI version from output: {output!r}")
+            return CliVersion()
+        return CliVersion(
+            int(m.group(1)),
+            int(m.group(2)),
+            int(m.group(3)),
+            suffix=m.group(4),
+        )
+
+    @staticmethod
+    def _get_cli_version(cli_path: str) -> CliVersion:
+        """Run 'databricks version' and return the parsed CliVersion.
+
+        Returns CliVersion() on failure.
+        """
+        try:
+            out = _run_subprocess([cli_path, "version"], capture_output=True, check=True)
+            return DatabricksCliTokenSource._parse_cli_version(out.stdout.decode())
+        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 ValueError with a precise message when no usable command can be
+        built (missing profile/host, or --profile-unsupported CLI with no host
+        fallback configured).
+        """
+        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 ValueError as e:
+            raise ValueError(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 ValueError 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 ValueError as e:
+                raise ValueError(f"neither profile nor host is configured: {e}") from e
+
+        # Flag --profile is a global CLI 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:
+            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 ValueError as e:
+                raise ValueError(
+                    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 ValueError("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: