66import com .databricks .sdk .core .oauth .OAuthHeaderFactory ;
77import com .databricks .sdk .core .oauth .Token ;
88import com .databricks .sdk .core .oauth .TokenSource ;
9+ import com .databricks .sdk .core .utils .Environment ;
910import com .databricks .sdk .core .utils .OSUtils ;
1011import com .databricks .sdk .support .InternalApi ;
1112import com .fasterxml .jackson .core .JsonProcessingException ;
13+ import com .fasterxml .jackson .databind .JsonNode ;
1214import com .fasterxml .jackson .databind .ObjectMapper ;
15+ import java .io .IOException ;
16+ import java .io .InputStream ;
1317import java .nio .charset .StandardCharsets ;
1418import java .util .*;
19+ import java .util .concurrent .ConcurrentHashMap ;
20+ import java .util .concurrent .TimeUnit ;
21+ import org .apache .commons .io .IOUtils ;
1522
1623@ InternalApi
1724public class DatabricksCliCredentialsProvider implements CredentialsProvider {
@@ -22,6 +29,18 @@ public class DatabricksCliCredentialsProvider implements CredentialsProvider {
2229
2330 private static final ObjectMapper MAPPER = new ObjectMapper ();
2431
32+ // --profile support added in CLI v0.207.1: https://github.com/databricks/cli/pull/855
33+ static final DatabricksCliVersion CLI_VERSION_FOR_PROFILE = new DatabricksCliVersion (0 , 207 , 1 );
34+
35+ // 5-second cap on `databricks version` so a hung CLI (slow first-run scan, antivirus, blocked
36+ // stdin) does not wedge SDK init indefinitely.
37+ private static final long VERSION_PROBE_TIMEOUT_SECONDS = 5 ;
38+
39+ // Successful version probes keyed by cliPath. Failures are deliberately not cached, so a
40+ // transient error (timeout, AV scan) does not pin every later token source to the conservative
41+ // fallback for the rest of the process lifetime.
42+ private static final Map <String , DatabricksCliVersion > VERSION_CACHE = new ConcurrentHashMap <>();
43+
2544 /** Thrown when the cached CLI token's scopes don't match the SDK's configured scopes. */
2645 static class ScopeMismatchException extends DatabricksException {
2746 ScopeMismatchException (String message ) {
@@ -58,6 +77,150 @@ List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
5877 return cmd ;
5978 }
6079
80+ /**
81+ * Builds the {@code auth token} command for the given CLI version.
82+ *
83+ * <p>Falls back to {@code --host} when {@code --profile} is either not configured or not
84+ * supported by the installed CLI.
85+ */
86+ List <String > buildCliCommand (
87+ String cliPath , DatabricksConfig config , DatabricksCliVersion version ) {
88+ if (config .getProfile () == null ) {
89+ return buildHostArgs (cliPath , config );
90+ }
91+
92+ // Flag --profile is a global CLI flag and is recognized for all commands even the ones that
93+ // do not support it. Only use --profile in CLI versions known to support it in `auth token`.
94+ if (!version .atLeast (CLI_VERSION_FOR_PROFILE )) {
95+ if (version .equals (DatabricksCliVersion .UNKNOWN ) || version .isDefaultDevBuild ()) {
96+ // We didn't actually prove the CLI lacks --profile; we just failed to confirm it.
97+ LOG .warn (
98+ "Could not confirm --profile support for Databricks CLI {} (requires >= {}). "
99+ + "Falling back to --host." ,
100+ version ,
101+ CLI_VERSION_FOR_PROFILE );
102+ } else {
103+ LOG .warn (
104+ "Databricks CLI {} does not support --profile (requires >= {}). Falling back to --host." ,
105+ version ,
106+ CLI_VERSION_FOR_PROFILE );
107+ }
108+ return buildHostArgs (cliPath , config );
109+ }
110+
111+ return new ArrayList <>(
112+ Arrays .asList (cliPath , "auth" , "token" , "--profile" , config .getProfile ()));
113+ }
114+
115+ /**
116+ * Detects the installed CLI version and builds the {@code auth token} command. Falls back to the
117+ * most conservative command when version detection fails.
118+ */
119+ List <String > resolveCliCommand (String cliPath , DatabricksConfig config ) {
120+ DatabricksCliVersion version = getCliVersion (cliPath , config .getEnv ());
121+ if (version .isDefaultDevBuild ()) {
122+ // A default-marker dev build has no injected version, so every feature gate fails.
123+ // Surface an informational hint so users know why their feature flags aren't taking effect.
124+ LOG .info (
125+ "Databricks CLI {} is a development build; feature detection will use conservative "
126+ + "fallbacks. Rebuild the CLI with an explicit version to enable capability-based "
127+ + "flag selection." ,
128+ version );
129+ }
130+ return buildCliCommand (cliPath , config , version );
131+ }
132+
133+ /**
134+ * Returns the CLI version, catching subprocess failures so the caller can proceed with the
135+ * conservative fallback. Successful results are cached per {@code cliPath} for the process
136+ * lifetime; failures are not cached and will be retried on the next call.
137+ */
138+ DatabricksCliVersion getCliVersion (String cliPath , Environment env ) {
139+ DatabricksCliVersion cached = VERSION_CACHE .get (cliPath );
140+ if (cached != null ) {
141+ return cached ;
142+ }
143+
144+ try {
145+ DatabricksCliVersion version = probeCliVersion (cliPath , env );
146+ VERSION_CACHE .put (cliPath , version );
147+ return version ;
148+ } catch (Exception e ) {
149+ LOG .warn (
150+ "Failed to detect Databricks CLI version: {}. Falling back to conservative flag set." ,
151+ e .getMessage ());
152+ return DatabricksCliVersion .UNKNOWN ;
153+ }
154+ }
155+
156+ /**
157+ * Runs {@code databricks version --output json} and returns the parsed {@link
158+ * DatabricksCliVersion}.
159+ */
160+ DatabricksCliVersion probeCliVersion (String cliPath , Environment env ) throws IOException {
161+ List <String > versionArgs = Arrays .asList (cliPath , "version" , "--output" , "json" );
162+ List <String > cmd = OSUtils .get (env ).getCliExecutableCommand (versionArgs );
163+
164+ ProcessBuilder pb = new ProcessBuilder (cmd );
165+ pb .environment ().putAll (env .getEnv ());
166+ Process process = pb .start ();
167+
168+ try {
169+ if (!process .waitFor (VERSION_PROBE_TIMEOUT_SECONDS , TimeUnit .SECONDS )) {
170+ process .destroyForcibly ();
171+ throw new IOException (
172+ "timed out after "
173+ + VERSION_PROBE_TIMEOUT_SECONDS
174+ + "s waiting for `databricks version`" );
175+ }
176+ } catch (InterruptedException e ) {
177+ Thread .currentThread ().interrupt ();
178+ throw new IOException ("interrupted waiting for `databricks version`" , e );
179+ }
180+
181+ if (process .exitValue () != 0 ) {
182+ String stderr = readStream (process .getErrorStream ());
183+ throw new IOException (
184+ "`databricks version` exited with code " + process .exitValue () + ": " + stderr );
185+ }
186+
187+ return parseCliVersion (readStream (process .getInputStream ()));
188+ }
189+
190+ private static String readStream (InputStream stream ) throws IOException {
191+ return new String (IOUtils .toByteArray (stream ), StandardCharsets .UTF_8 );
192+ }
193+
194+ /**
195+ * Parses the JSON output of {@code databricks version --output json}.
196+ *
197+ * <p>Takes Major/Minor/Patch from the JSON's pre-parsed numeric fields. The Prerelease field and
198+ * the Version string are intentionally ignored: for our feature-gate purposes the base triple is
199+ * sufficient, and the (0, 0, 0) case already identifies the default dev build (a CLI built
200+ * without version metadata leaves these fields at their zero defaults).
201+ *
202+ * <p>Returns {@link DatabricksCliVersion#UNKNOWN} on failure so that an unparseable version
203+ * disables every feature gate.
204+ */
205+ static DatabricksCliVersion parseCliVersion (String output ) {
206+ try {
207+ JsonNode node = MAPPER .readTree (output );
208+ JsonNode major = node .get ("Major" );
209+ JsonNode minor = node .get ("Minor" );
210+ JsonNode patch = node .get ("Patch" );
211+ if (major == null || minor == null || patch == null ) {
212+ LOG .debug (
213+ "Failed to parse Databricks CLI version: missing Major/Minor/Patch in {}" , output );
214+ return DatabricksCliVersion .UNKNOWN ;
215+ }
216+ return new DatabricksCliVersion (major .asInt (), minor .asInt (), patch .asInt ());
217+ } catch (JsonProcessingException e ) {
218+ LOG .debug (
219+ "Failed to parse Databricks CLI version from output: {} ({})" , output , e .getMessage ());
220+ return DatabricksCliVersion .UNKNOWN ;
221+ }
222+ }
223+
61224 private CliTokenSource getDatabricksCliTokenSource (DatabricksConfig config ) {
62225 String cliPath = config .getDatabricksCliPath ();
63226 if (cliPath == null ) {
@@ -68,25 +231,8 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
68231 return null ;
69232 }
70233
71- List <String > cmd ;
72- List <String > fallbackCmd = null ;
73-
74- if (config .getProfile () != null ) {
75- // When profile is set, use --profile as the primary command.
76- // The profile contains the full config (host, account_id, etc.).
77- cmd =
78- new ArrayList <>(
79- Arrays .asList (cliPath , "auth" , "token" , "--profile" , config .getProfile ()));
80- // Build a --host fallback for older CLIs that don't support --profile.
81- if (config .getHost () != null ) {
82- fallbackCmd = buildHostArgs (cliPath , config );
83- }
84- } else {
85- cmd = buildHostArgs (cliPath , config );
86- }
87-
88- return new CliTokenSource (
89- cmd , "token_type" , "access_token" , "expiry" , config .getEnv (), fallbackCmd );
234+ List <String > cmd = resolveCliCommand (cliPath , config );
235+ return new CliTokenSource (cmd , "token_type" , "access_token" , "expiry" , config .getEnv ());
90236 }
91237
92238 @ Override
0 commit comments