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,104 @@ 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 `prerelease` means version detection
+    failed and the SDK should fall back to the most conservative command.
+    `prerelease` matches the CLI's own `Prerelease` field in the JSON version
+    output (e.g. "dev", "rc.1", or "" for a stable release).
+    """
+
+    major: int = 0
+    minor: int = 0
+    patch: int = 0
+    # e.g. "dev", "rc.1", "alpha", or "" for a stable release. No leading "-".
+    prerelease: 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
+        and prerelease equal to the default dev build's.
+
+        A version the user explicitly set (e.g. v1.0.0-dev) is intentional and
+        is not flagged here; feature gates still apply via the comparison
+        operators.
+        """
+        return (
+            self.major == _DEFAULT_DEV_BUILD.major
+            and self.minor == _DEFAULT_DEV_BUILD.minor
+            and self.patch == _DEFAULT_DEV_BUILD.patch
+            and self.prerelease == _DEFAULT_DEV_BUILD.prerelease
+        )
+
+    @property
+    def _has_prerelease(self) -> bool:
+        # In semver, a prerelease tag (non-empty) compares less than the same
+        # base version without one.
+        return bool(self.prerelease)
+
+    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 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.prerelease:
+            return "unknown"
+        base = f"v{self.major}.{self.minor}.{self.patch}"
+        return f"{base}-{self.prerelease}" if self.prerelease else base
+
+
+# The Databricks CLI emits exactly "v0.0.0-dev" (Major/Minor/Patch=0 and
+# Prerelease="dev" in the JSON output) when the binary was built without
+# injecting a version via ldflags (a plain `go build`). See
+# https://github.com/databricks/cli/blob/main/internal/build/info.go.
+_DEFAULT_DEV_BUILD = CliVersion(0, 0, 0, prerelease="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 +796,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 +985,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 +995,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 +1057,126 @@ 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, and
+        derives the prerelease tag from the `Version` string. The JSON's own
+        `Prerelease` field is only populated by goreleaser ldflags; a plain
+        `go build` leaves it empty even when `Version` reports
+        "0.0.0-dev+<commit>". Splitting `Version` ourselves catches that case.
+
+        Returns CliVersion() on failure, matching Go's behavior where an
+        unparseable version disables every feature gate.
+        """
+        try:
+            data = json.loads(output)
+            version = data["Version"]
+            # Drop semver build metadata ("+abcdef") — it does not affect ordering.
+            # Then split base from prerelease ("0.295.0-rc.1" → prerelease "rc.1").
+            _, _, tail = version.partition("-")
+            prerelease = tail.partition("+")[0]
+            return CliVersion(
+                int(data["Major"]),
+                int(data["Minor"]),
+                int(data["Patch"]),
+                prerelease=prerelease,
+            )
+        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
+    def _get_cli_version(cli_path: str) -> CliVersion:
+        """Run `databricks version --output json` and return the parsed CliVersion.
+
+        Returns CliVersion() on failure. `--output json` is a global CLI flag
+        available since CLI v0.100.3; older CLIs that ignore or error on it
+        simply fail the JSON parse and fall back to the zero version.
+        """
+        try:
+            out = _run_subprocess(
+                [cli_path, "version", "--output", "json"],
+                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: