Skip to content

Commit ba0f51c

Browse files
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
1 parent 78914fa commit ba0f51c

3 files changed

Lines changed: 452 additions & 195 deletions

File tree

NEXT_CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
### Security
1010

1111
### Bug Fixes
12+
* 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.
1213

1314
### Documentation
1415

@@ -18,6 +19,7 @@
1819
### Internal Changes
1920
* Replace the async-disabling mechanism on token refresh failure with a 1-minute retry backoff. Previously, a single failed async refresh would disable proactive token renewal until the token expired. Now, the SDK waits a short cooldown period and retries, improving resilience to transient errors.
2021
* Extract `_resolve_profile` to simplify config file loading and improve `__settings__` error messages.
22+
* Detect Databricks CLI version at init time via `databricks version`, enabling version-gated flag support without additional subprocess calls.
2123

2224
### API Changes
2325
* Add `create_catalog()`, `create_synced_table()`, `delete_catalog()`, `delete_synced_table()`, `get_catalog()` and `get_synced_table()` methods for [w.postgres](https://databricks-sdk-py.readthedocs.io/en/latest/workspace/postgres/postgres.html) workspace-level service.

databricks/sdk/credentials_provider.py

Lines changed: 181 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import abc
22
import base64
3+
import dataclasses
34
import functools
45
import io
56
import json
@@ -649,23 +650,57 @@ def refreshed_headers() -> Dict[str, str]:
649650
return OAuthCredentialsProvider(refreshed_headers, token)
650651

651652

653+
@dataclasses.dataclass(order=True)
654+
class CliVersion:
655+
"""Semver version triple of the Databricks CLI.
656+
657+
Three sentinel states in the (major, minor, patch) tuple:
658+
* `(-1, -1, -1)` — the default, meaning version detection failed. It
659+
compares less than every real release so every feature gate fails.
660+
* `(0, 0, 0)` — the CLI's default dev build, emitted when the binary
661+
was built without version metadata. See `is_default_dev_build`.
662+
* anything else — a real CLI version.
663+
664+
Prerelease tags (e.g. "-rc.1", "-dev+commit") are deliberately ignored:
665+
for our purposes the base triple is sufficient, and feature gates are
666+
release-based so a prerelease of a version with a flag is assumed to
667+
have the flag too.
668+
"""
669+
670+
major: int = -1
671+
minor: int = -1
672+
patch: int = -1
673+
674+
@property
675+
def is_default_dev_build(self) -> bool:
676+
"""True when the CLI reports (0, 0, 0), i.e. its default dev build.
677+
678+
Narrowly matches the CLI's "no version injected" marker: its version
679+
metadata stays at the zero defaults. A version the user explicitly
680+
set (e.g. v1.0.0-dev) is intentional and is not flagged here.
681+
"""
682+
return (self.major, self.minor, self.patch) == (0, 0, 0)
683+
684+
def __str__(self) -> str:
685+
if self == CliVersion():
686+
return "unknown"
687+
if self.is_default_dev_build:
688+
return "v0.0.0-dev"
689+
return f"v{self.major}.{self.minor}.{self.patch}"
690+
691+
652692
class CliTokenSource(oauth.Refreshable):
693+
653694
def __init__(
654695
self,
655696
cmd: List[str],
656697
token_type_field: str,
657698
access_token_field: str,
658699
expiry_field: str,
659700
disable_async: bool = True,
660-
fallback_cmd: Optional[List[str]] = None,
661701
):
662702
super().__init__(disable_async=disable_async)
663703
self._cmd = cmd
664-
# fallback_cmd is tried when the primary command fails with "unknown flag: --profile",
665-
# indicating the CLI is too old to support --profile. Can be removed once support
666-
# for CLI versions predating --profile is dropped.
667-
# See: https://github.com/databricks/databricks-sdk-go/pull/1497
668-
self._fallback_cmd = fallback_cmd
669704
self._token_type_field = token_type_field
670705
self._access_token_field = access_token_field
671706
self._expiry_field = expiry_field
@@ -700,16 +735,7 @@ def _exec_cli_command(self, cmd: List[str]) -> oauth.Token:
700735
raise IOError(f"cannot get access token: {message}") from e
701736

702737
def refresh(self) -> oauth.Token:
703-
try:
704-
return self._exec_cli_command(self._cmd)
705-
except IOError as e:
706-
if self._fallback_cmd is not None and "unknown flag: --profile" in str(e):
707-
logger.warning(
708-
"Databricks CLI does not support --profile flag. Falling back to --host. "
709-
"Please upgrade your CLI to the latest version."
710-
)
711-
return self._exec_cli_command(self._fallback_cmd)
712-
raise
738+
return self._exec_cli_command(self._cmd)
713739

714740

715741
def _run_subprocess(
@@ -876,6 +902,37 @@ def inner() -> Dict[str, str]:
876902
return inner
877903

878904

905+
@functools.lru_cache(maxsize=8)
906+
def _get_cli_version(cli_path: str) -> "CliVersion":
907+
"""Run `databricks version --output json` and return the parsed CliVersion.
908+
909+
Cached by `cli_path` for the Python process lifetime: the installed binary's
910+
version is immutable during the process, and the credential chain may
911+
instantiate DatabricksCliTokenSource many times (per WorkspaceClient, per
912+
request in some patterns). Without the cache, every construction would
913+
re-fork `databricks version`.
914+
915+
The 5-second timeout prevents a hung CLI (first-run disk scan, antivirus,
916+
stdin wait) from wedging SDK init indefinitely; TimeoutExpired is caught
917+
by the generic handler and falls through to CliVersion().
918+
919+
Returns CliVersion() on any failure.
920+
"""
921+
try:
922+
out = _run_subprocess(
923+
[cli_path, "version", "--output", "json"],
924+
capture_output=True,
925+
check=True,
926+
timeout=5,
927+
)
928+
return DatabricksCliTokenSource._parse_cli_version(out.stdout.decode())
929+
except Exception as e:
930+
logger.warning(
931+
f"Failed to detect Databricks CLI version: {e}. Falling back to conservative flag set."
932+
)
933+
return CliVersion()
934+
935+
879936
class DatabricksCliTokenSource(CliTokenSource):
880937
"""Obtain the token granted by `databricks auth login` CLI command"""
881938

@@ -898,16 +955,7 @@ def __init__(self, cfg: "Config"):
898955
elif cli_path.count("/") == 0:
899956
cli_path = self.__class__._find_executable(cli_path)
900957

901-
fallback_cmd = None
902-
if cfg.profile:
903-
# When profile is set, use --profile as the primary command.
904-
# The profile contains the full config (host, account_id, etc.).
905-
args = ["auth", "token", "--profile", cfg.profile]
906-
# Build a --host fallback for older CLIs that don't support --profile.
907-
if cfg.host:
908-
fallback_cmd = [cli_path, *self.__class__._build_host_args(cfg)]
909-
else:
910-
args = self.__class__._build_host_args(cfg)
958+
cmd = self.__class__._resolve_cli_command(cli_path, cfg)
911959

912960
# get_scopes() defaults to ["all-apis"] when nothing is configured, which would
913961
# cause false-positive mismatches against every token that wasn't issued with
@@ -917,12 +965,11 @@ def __init__(self, cfg: "Config"):
917965
self._host = cfg.host
918966

919967
super().__init__(
920-
cmd=[cli_path, *args],
968+
cmd=cmd,
921969
token_type_field="token_type",
922970
access_token_field="access_token",
923971
expiry_field="expiry",
924972
disable_async=cfg.disable_async_token_refresh,
925-
fallback_cmd=fallback_cmd,
926973
)
927974

928975
def refresh(self) -> oauth.Token:
@@ -980,15 +1027,114 @@ def _validate_token_scopes(self, token: oauth.Token):
9801027
f"Scopes default to all-apis."
9811028
)
9821029

1030+
# --profile support added in CLI v0.207.1: https://github.com/databricks/cli/pull/855
1031+
_CLI_VERSION_FOR_PROFILE = CliVersion(0, 207, 1)
1032+
1033+
@staticmethod
1034+
def _parse_cli_version(output: str) -> CliVersion:
1035+
"""Parse the JSON output of `databricks version --output json`.
1036+
1037+
Takes Major/Minor/Patch from the JSON's pre-parsed numeric fields. The
1038+
`Prerelease` JSON field and the `Version` string are intentionally
1039+
ignored: for our feature-gate purposes the base triple is sufficient,
1040+
and the (0, 0, 0) case already identifies the default dev build (a
1041+
CLI built without version metadata leaves these fields at their zero
1042+
defaults).
1043+
1044+
Returns CliVersion() (unknown, (-1, -1, -1)) on failure so that an
1045+
unparseable version disables every feature gate.
1046+
"""
1047+
try:
1048+
data = json.loads(output)
1049+
return CliVersion(
1050+
int(data["Major"]),
1051+
int(data["Minor"]),
1052+
int(data["Patch"]),
1053+
)
1054+
except (json.JSONDecodeError, KeyError, TypeError, ValueError) as e:
1055+
logger.debug(f"Failed to parse Databricks CLI version from output: {output!r} ({e})")
1056+
return CliVersion()
1057+
1058+
@staticmethod
1059+
def _resolve_cli_command(cli_path: str, cfg: "Config") -> List[str]:
1060+
"""Detect CLI version and build the auth-token command.
1061+
1062+
Raises IOError with a precise message when no usable command can be
1063+
built (missing profile/host, or --profile-unsupported CLI with no host
1064+
fallback configured). IOError matches the exception type already
1065+
handled by the `databricks_cli()` credential-chain strategy, so
1066+
graceful-skip behaviour is preserved.
1067+
"""
1068+
version = _get_cli_version(cli_path)
1069+
if version.is_default_dev_build:
1070+
# A default-marker dev build has no injected version, so every
1071+
# feature gate fails via at_least. Surface an informational hint so
1072+
# users know why their feature flags aren't taking effect.
1073+
logger.info(
1074+
f"Databricks CLI {version} is a development build; feature detection will use "
1075+
"conservative fallbacks. Rebuild the CLI with an explicit version to enable "
1076+
"capability-based flag selection."
1077+
)
1078+
try:
1079+
return DatabricksCliTokenSource._build_cli_command(cli_path, cfg, version)
1080+
except IOError as e:
1081+
raise IOError(f"cannot configure CLI token source: {e}") from e
1082+
1083+
@staticmethod
1084+
def _build_cli_command(cli_path: str, cfg: "Config", version: CliVersion) -> List[str]:
1085+
"""Build the `auth token` command.
1086+
1087+
Falls back to --host when --profile is either not configured or not
1088+
supported by the installed CLI. Raises IOError describing which
1089+
piece of configuration is missing when no command can be built.
1090+
"""
1091+
if not cfg.profile:
1092+
try:
1093+
return DatabricksCliTokenSource._build_host_command(cli_path, cfg)
1094+
except IOError as e:
1095+
raise IOError(f"neither profile nor host is configured: {e}") from e
1096+
1097+
# Flag --profile is a global flag and is recognized for all
1098+
# commands even the ones that do not support it. Only use --profile
1099+
# in CLI versions that are known to support it in `auth token`.
1100+
if version < DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE:
1101+
# For the unknown (detection failed) and dev-build (no version
1102+
# metadata) cases we can't actually prove the CLI lacks --profile;
1103+
# we just failed to confirm it. Log accordingly.
1104+
if version == CliVersion() or version.is_default_dev_build:
1105+
logger.warning(
1106+
f"Could not confirm --profile support for Databricks CLI {version} "
1107+
f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}). "
1108+
"Falling back to --host."
1109+
)
1110+
else:
1111+
logger.warning(
1112+
f"Databricks CLI {version} does not support --profile "
1113+
f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}). "
1114+
"Falling back to --host."
1115+
)
1116+
try:
1117+
return DatabricksCliTokenSource._build_host_command(cli_path, cfg)
1118+
except IOError as e:
1119+
raise IOError(
1120+
f"Databricks CLI {version} does not support --profile "
1121+
f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_PROFILE}) "
1122+
f"and no host fallback is configured: {e}"
1123+
) from e
1124+
1125+
return [cli_path, "auth", "token", "--profile", cfg.profile]
1126+
9831127
@staticmethod
984-
def _build_host_args(cfg: "Config") -> List[str]:
985-
"""Build CLI arguments using --host (legacy path)."""
986-
args = ["auth", "token", "--host", cfg.host]
987-
# This is here to support older versions of the Databricks CLI, so we need to keep the client type check.
988-
# This won't work for unified hosts, but it is not supposed to.
1128+
def _build_host_command(cli_path: str, cfg: "Config") -> List[str]:
1129+
"""Build the --host based auth token command. Raises when cfg.host is unset."""
1130+
if not cfg.host:
1131+
raise IOError("host is not set")
1132+
cmd = [cli_path, "auth", "token", "--host", cfg.host]
1133+
# Older Databricks CLIs need --account-id for account-scoped calls;
1134+
# unified hosts don't use this path.
9891135
if cfg.client_type == ClientType.ACCOUNT:
990-
args += ["--account-id", cfg.account_id]
991-
return args
1136+
cmd += ["--account-id", cfg.account_id]
1137+
return cmd
9921138

9931139
@staticmethod
9941140
def _find_executable(name) -> str:

0 commit comments

Comments
 (0)