11import abc
22import base64
3+ import dataclasses
34import functools
45import io
56import 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+
652692class 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
715741def _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+
879936class 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