Skip to content

Refactor CliTokenSource to use an ordered attempt chain#752

Draft
mihaimitrea-db wants to merge 2 commits intomainfrom
mihaimitrea-db/stack/cli-attempt-chain
Draft

Refactor CliTokenSource to use an ordered attempt chain#752
mihaimitrea-db wants to merge 2 commits intomainfrom
mihaimitrea-db/stack/cli-attempt-chain

Conversation

@mihaimitrea-db
Copy link
Copy Markdown
Contributor

@mihaimitrea-db mihaimitrea-db commented Mar 31, 2026

Stacked PR

Use this link to review incremental changes.


Summary

Generalize CliTokenSource from three explicit command fields (cmd, fallbackCmd, secondFallbackCmd) into a List<CliCommand> attempt chain, so that adding future CLI flags is straightforward.

Why

The parent PR (#751) introduced --force-refresh support by adding a third command field and hand-writing each fallback block in getToken(). This works, but every new flag would require adding another field, another if block, another error check, and another test — the pattern doesn't scale.

We expect future flags like --scopes (forwarding custom OAuth scopes to the CLI). Rather than growing the class linearly with each flag, this PR extracts the repeating pattern into a loop over a command list.

Why try-and-retry over version detection or --help parsing

Three approaches were evaluated for resolving which flags the installed CLI supports:

  • Version detection (databricks version + static version table) was rejected because it creates a maintenance burden and a second source of truth. Every SDK (Go, Python, Java) would need to independently maintain a table mapping flags to the CLI version that introduced them. If any SDK's table falls out of sync with the CLI's actual releases, users silently get degraded commands.
  • --help flag parsing (databricks auth token --help + contains) was rejected because it depends on the output format of --help — which is not a stable API. Cobra format changes could break detection, and naive substring matching is fragile.
  • Feature probing with try-and-retry (the approach taken here) uses the CLI itself as the authority on what it supports. Commands are built at init time from most-featured to simplest. On the first getToken() call, each command is tried in order; when the CLI responds with "unknown flag:", the next simpler command is tried. This approach has zero maintenance burden (no version numbers or flag registries to keep in sync), zero overhead on the happy path (newest CLI succeeds on the first command), and requires no signature changes.

What changed

Interface changes

None. CliTokenSource is not part of the public API surface.

Behavioral changes

None. The set of commands tried is identical to the parent PR. The only observable difference is that isUnknownFlagError now matches against specific usedFlags per command rather than a blanket "unknown flag:" check, preventing false-positive fallbacks from unexpected unknown-flag errors.

Internal changes

  • CliCommand inner class: replaces the three separate List<String> fields. Each entry holds cmd (the full CLI command), usedFlags (flags in this command, used for error matching), and fallbackMessage (logged when falling back from this command).
  • CliTokenSource: now holds List<CliCommand> attempts instead of cmd, fallbackCmd, secondFallbackCmd. The 6-arg and 7-arg public constructors are removed; only the 5-arg constructor (Azure CLI) and the fromAttempts factory remain.
  • DatabricksCliCredentialsProvider.buildAttempts: constructs command variants inline — the most-featured command (--profile + --force-refresh) first, then the plain --profile command, then --host if available. Profile and host commands are built independently. Adding a future flag means adding one more CliCommand literal here.
  • getToken(): a single for loop over attempts. Non-terminal commands fall through when the error matches one of the command's usedFlags; the terminal command returns errors directly.
  • Empty-attempts guard: construction-time validation throws immediately if the attempt list is empty, failing fast instead of at token-fetch time.

How is this tested?

Unit tests in DatabricksCliCredentialsProviderTest:

  • testBuildAttempts_WithProfileAndHost — verifies 3 commands with correct cmd, usedFlags, and fallbackMessage.
  • testBuildAttempts_WithProfileOnly — verifies 2 commands (no host fallback).
  • testBuildAttempts_WithHostOnly — verifies single --host command, no fallback chain.
  • testBuildAttempts_WithAccountHost — verifies --host + --account-id for account-level hosts.

Existing CliTokenSourceTest fallback behavior tests continue to pass against the refactored getToken() loop.

@mihaimitrea-db
Copy link
Copy Markdown
Contributor Author

Range-diff: stack/cli-force-refresh (628f509 -> 09cdbc4)
NEXT_CHANGELOG.md
@@ -0,0 +1,10 @@
+diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
+--- a/NEXT_CHANGELOG.md
++++ b/NEXT_CHANGELOG.md
+ ### Documentation
+ 
+ ### Internal Changes
++* Generalize CLI token source into a progressive command list for forward-compatible flag support.
+ 
+ ### API Changes
+ * Add `createCatalog()`, `createSyncedTable()`, `deleteCatalog()`, `deleteSyncedTable()`, `getCatalog()` and `getSyncedTable()` methods for `workspaceClient.postgres()` service.
\ No newline at end of file
databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
@@ -132,8 +132,7 @@
      this.env = env;
 -    this.fallbackCmd =
 -        fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
--    this.forceCmd =
--        forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
+-    this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
 +  }
 +
 +  private static final String UNKNOWN_PROFILE_FLAG = "unknown flag: --profile";
databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
@@ -21,8 +21,7 @@
 +    return result;
 +  }
 +
-+  List<CliTokenSource.CliCommand> buildAttempts(
-+      String cliPath, DatabricksConfig config) {
++  List<CliTokenSource.CliCommand> buildAttempts(String cliPath, DatabricksConfig config) {
 +    List<CliTokenSource.CliCommand> attempts = new ArrayList<>();
 +
 +    List<String> profileCmd;
@@ -54,9 +53,7 @@
 +          new CliTokenSource.CliCommand(
 +              buildHostArgs(cliPath, config), Collections.emptyList(), null));
 +    } else {
-+      attempts.add(
-+          new CliTokenSource.CliCommand(
-+              profileCmd, Collections.emptyList(), null));
++      attempts.add(new CliTokenSource.CliCommand(profileCmd, Collections.emptyList(), null));
 +    }
 +
 +    return attempts;
@@ -84,11 +81,7 @@
 -    return new CliTokenSource(
 -        profileCmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd, forceCmd);
 +    return CliTokenSource.fromAttempts(
-+        buildAttempts(cliPath, config),
-+        "token_type",
-+        "access_token",
-+        "expiry",
-+        config.getEnv());
++        buildAttempts(cliPath, config), "token_type", "access_token", "expiry", config.getEnv());
    }
  
    @Override
\ No newline at end of file
databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
@@ -9,8 +9,7 @@
 +
 +  @Test
 +  void testBuildAttempts_WithProfileAndHost() {
-+    DatabricksConfig config =
-+        new DatabricksConfig().setHost(HOST).setProfile("my-profile");
++    DatabricksConfig config = new DatabricksConfig().setHost(HOST).setProfile("my-profile");
 +
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
@@ -18,12 +17,9 @@
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile", "--force-refresh"),
 +        attempts.get(0).cmd);
-+    assertEquals(
-+        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"),
-+        attempts.get(1).cmd);
 +    assertEquals(
-+        Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST),
-+        attempts.get(2).cmd);
++        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), attempts.get(1).cmd);
++    assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), attempts.get(2).cmd);
 +  }
 +
 +  @Test
@@ -37,8 +33,7 @@
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile", "--force-refresh"),
 +        attempts.get(0).cmd);
 +    assertEquals(
-+        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"),
-+        attempts.get(1).cmd);
++        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), attempts.get(1).cmd);
 +  }
 +
 +  @Test
@@ -51,23 +46,26 @@
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST, "--force-refresh"),
 +        attempts.get(0).cmd);
-+    assertEquals(
-+        Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST),
-+        attempts.get(1).cmd);
++    assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), attempts.get(1).cmd);
 +  }
 +
 +  @Test
 +  void testBuildAttempts_WithAccountHost() {
-+    DatabricksConfig config =
-+        new DatabricksConfig().setHost(ACCOUNT_HOST).setAccountId(ACCOUNT_ID);
++    DatabricksConfig config = new DatabricksConfig().setHost(ACCOUNT_HOST).setAccountId(ACCOUNT_ID);
 +
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
 +    assertEquals(2, attempts.size());
 +    assertEquals(
 +        Arrays.asList(
-+            CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST,
-+            "--account-id", ACCOUNT_ID, "--force-refresh"),
++            CLI_PATH,
++            "auth",
++            "token",
++            "--host",
++            ACCOUNT_HOST,
++            "--account-id",
++            ACCOUNT_ID,
++            "--force-refresh"),
 +        attempts.get(0).cmd);
 +    assertEquals(
 +        Arrays.asList(

Reproduce locally: git range-diff 048a903..628f509 5e8f476..09cdbc4 | Disable: git config gitstack.push-range-diff false

@mihaimitrea-db mihaimitrea-db force-pushed the mihaimitrea-db/stack/cli-attempt-chain branch from 09cdbc4 to f88c52a Compare March 31, 2026 13:22
@mihaimitrea-db
Copy link
Copy Markdown
Contributor Author

Range-diff: stack/cli-force-refresh (09cdbc4 -> f88c52a)
databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
@@ -26,9 +26,9 @@
 +   * command in the chain and an optional log message emitted on fallback.
 +   */
 +  static class CliCommand {
-+    private final List<String> cmd;
-+    private final List<String> fallbackTriggers;
-+    private final String fallbackMessage;
++    final List<String> cmd;
++    final List<String> fallbackTriggers;
++    final String fallbackMessage;
  
 -  private List<String> profileCmd;
 -  private String tokenTypeField;

Reproduce locally: git range-diff 5e8f476..09cdbc4 5e8f476..f88c52a | Disable: git config gitstack.push-range-diff false

@mihaimitrea-db mihaimitrea-db force-pushed the mihaimitrea-db/stack/cli-attempt-chain branch from f88c52a to 25a3779 Compare March 31, 2026 14:04
@mihaimitrea-db
Copy link
Copy Markdown
Contributor Author

Range-diff: stack/cli-force-refresh (f88c52a -> 25a3779)
databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
@@ -15,20 +15,19 @@
  public class CliTokenSource implements TokenSource {
    private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
  
--  private static final String UNKNOWN_PROFILE_FLAG = "unknown flag: --profile";
--  private static final String UNKNOWN_FORCE_REFRESH_FLAG = "unknown flag: --force-refresh";
--
 -  // forceCmd is tried before profileCmd when non-null. If the CLI rejects
 -  // --force-refresh or --profile, execution falls through to profileCmd.
 -  private List<String> forceCmd;
 +  /**
-+   * Describes a CLI command with the error substrings that allow falling through to the next
-+   * command in the chain and an optional log message emitted on fallback.
++   * Describes a CLI command with an optional warning message emitted when falling through to the
++   * next command in the chain.
 +   */
 +  static class CliCommand {
 +    final List<String> cmd;
-+    final List<String> fallbackTriggers;
-+    final String fallbackMessage;
++
++    // Flags used by this command (e.g. "--force-refresh", "--profile"). Used to distinguish
++    // "unknown flag" errors (which trigger fallback) from real auth errors (which propagate).
++    final List<String> usedFlags;
  
 -  private List<String> profileCmd;
 -  private String tokenTypeField;
@@ -38,9 +37,11 @@
 -  // fallbackCmd is tried when profileCmd fails with "unknown flag: --profile",
 -  // indicating the CLI is too old to support --profile.
 -  private List<String> fallbackCmd;
-+    CliCommand(List<String> cmd, List<String> fallbackTriggers, String fallbackMessage) {
++    final String fallbackMessage;
++
++    CliCommand(List<String> cmd, List<String> usedFlags, String fallbackMessage) {
 +      this.cmd = cmd;
-+      this.fallbackTriggers = fallbackTriggers != null ? fallbackTriggers : Collections.emptyList();
++      this.usedFlags = usedFlags != null ? usedFlags : Collections.emptyList();
 +      this.fallbackMessage = fallbackMessage;
 +    }
 +  }
@@ -85,7 +86,7 @@
 +                a ->
 +                    new CliCommand(
 +                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
-+                        a.fallbackTriggers,
++                        a.usedFlags,
 +                        a.fallbackMessage))
 +            .collect(Collectors.toList()),
 +        tokenTypeField,
@@ -108,7 +109,7 @@
 +                a ->
 +                    new CliCommand(
 +                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
-+                        a.fallbackTriggers,
++                        a.usedFlags,
 +                        a.fallbackMessage))
 +            .collect(Collectors.toList()),
 +        tokenTypeField,
@@ -134,9 +135,6 @@
 -        fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
 -    this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
 +  }
-+
-+  private static final String UNKNOWN_PROFILE_FLAG = "unknown flag: --profile";
-+  private static final String UNKNOWN_FORCE_REFRESH_FLAG = "unknown flag: --force-refresh";
 +
 +  private static List<CliCommand> buildAttempts(
 +      List<String> forceCmd, List<String> profileCmd, List<String> fallbackCmd) {
@@ -146,7 +144,7 @@
 +      attempts.add(
 +          new CliCommand(
 +              forceCmd,
-+              Arrays.asList(UNKNOWN_FORCE_REFRESH_FLAG, UNKNOWN_PROFILE_FLAG),
++              Arrays.asList("--force-refresh", "--profile"),
 +              "Databricks CLI does not support --force-refresh flag. "
 +                  + "Falling back to regular token fetch. "
 +                  + "Please upgrade your CLI to the latest version."));
@@ -156,7 +154,7 @@
 +      attempts.add(
 +          new CliCommand(
 +              profileCmd,
-+              Collections.singletonList(UNKNOWN_PROFILE_FLAG),
++              Collections.singletonList("--profile"),
 +              "Databricks CLI does not support --profile flag. Falling back to --host. "
 +                  + "Please upgrade your CLI to the latest version."));
 +      attempts.add(new CliCommand(fallbackCmd, Collections.emptyList(), null));
@@ -168,6 +166,14 @@
    }
  
    /**
+         if (stderr.contains("not found")) {
+           throw new DatabricksException(stderr);
+         }
+-        // getMessage() returns the clean stderr-based message; getFullOutput() exposes
+-        // both streams so the caller can check for "unknown flag: --profile" in either.
+         throw new CliCommandException("cannot get access token: " + stderr, stdout + "\n" + stderr);
+       }
+       JsonNode jsonNode = new ObjectMapper().readTree(stdout);
      }
    }
  
@@ -179,7 +185,7 @@
    }
  
 -  private boolean isUnknownFlagError(String errorText, String flag) {
--    return errorText != null && errorText.contains(flag);
+-    return errorText != null && errorText.contains("unknown flag: " + flag);
 -  }
 -
 -  private Token execProfileCmdWithFallback() {
@@ -187,7 +193,7 @@
 -      return execCliCommand(this.profileCmd);
 -    } catch (IOException e) {
 -      String textToCheck = getErrorText(e);
--      if (fallbackCmd != null && isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
+-      if (fallbackCmd != null && isUnknownFlagError(textToCheck, "--profile")) {
 -        LOG.warn(
 -            "Databricks CLI does not support --profile flag. Falling back to --host. "
 -                + "Please upgrade your CLI to the latest version.");
@@ -196,39 +202,46 @@
 -        } catch (IOException fallbackException) {
 -          throw new DatabricksException(fallbackException.getMessage(), fallbackException);
 -        }
--      }
--      throw new DatabricksException(e.getMessage(), e);
-+  private static boolean shouldFallback(CliCommand attempt, String errorText) {
++  private static boolean isUnknownFlagError(String errorText, List<String> flags) {
 +    if (errorText == null) {
 +      return false;
++    }
++    for (String flag : flags) {
++      if (errorText.contains("unknown flag: " + flag)) {
++        return true;
+       }
+-      throw new DatabricksException(e.getMessage(), e);
      }
-+    return attempt.fallbackTriggers.stream().anyMatch(errorText::contains);
++    return false;
    }
  
    @Override
    public Token getToken() {
 -    if (forceCmd == null) {
 -      return execProfileCmdWithFallback();
--    }
-+    IOException lastException = null;
++    if (attempts.isEmpty()) {
++      throw new DatabricksException("cannot get access token: no CLI commands configured");
+     }
  
 -    try {
 -      return execCliCommand(this.forceCmd);
 -    } catch (IOException e) {
 -      String textToCheck = getErrorText(e);
--      if (isUnknownFlagError(textToCheck, UNKNOWN_FORCE_REFRESH_FLAG)
--          || isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
+-      if (isUnknownFlagError(textToCheck, "--force-refresh")
+-          || isUnknownFlagError(textToCheck, "--profile")) {
 -        LOG.warn(
 -            "Databricks CLI does not support --force-refresh flag. "
 -                + "Falling back to regular token fetch. "
 -                + "Please upgrade your CLI to the latest version.");
 -        return execProfileCmdWithFallback();
++    IOException lastException = null;
++
 +    for (int i = 0; i < attempts.size(); i++) {
 +      CliCommand attempt = attempts.get(i);
 +      try {
 +        return execCliCommand(attempt.cmd);
 +      } catch (IOException e) {
-+        if (i + 1 < attempts.size() && shouldFallback(attempt, getErrorText(e))) {
++        if (i + 1 < attempts.size() && isUnknownFlagError(getErrorText(e), attempt.usedFlags)) {
 +          if (attempt.fallbackMessage != null) {
 +            LOG.warn(attempt.fallbackMessage);
 +          }
databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
@@ -1,17 +1,8 @@
 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
 +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
-     return cmd;
    }
  
-+  private static final String UNKNOWN_PROFILE_FLAG = "unknown flag: --profile";
-+  private static final String UNKNOWN_FORCE_REFRESH_FLAG = "unknown flag: --force-refresh";
-+
-   List<String> buildProfileArgs(String cliPath, DatabricksConfig config) {
-     return new ArrayList<>(
-         Arrays.asList(cliPath, "auth", "token", "--profile", config.getProfile()));
-   }
- 
    private static List<String> withForceRefresh(List<String> cmd) {
 -    List<String> forceCmd = new ArrayList<>(cmd);
 -    forceCmd.add("--force-refresh");
@@ -37,7 +28,7 @@
 +    attempts.add(
 +        new CliTokenSource.CliCommand(
 +            withForceRefresh(profileCmd),
-+            Arrays.asList(UNKNOWN_FORCE_REFRESH_FLAG, UNKNOWN_PROFILE_FLAG),
++            Arrays.asList("--force-refresh", "--profile"),
 +            "Databricks CLI does not support --force-refresh flag. "
 +                + "Falling back to regular token fetch. "
 +                + "Please upgrade your CLI to the latest version."));
@@ -46,7 +37,7 @@
 +      attempts.add(
 +          new CliTokenSource.CliCommand(
 +              profileCmd,
-+              Collections.singletonList(UNKNOWN_PROFILE_FLAG),
++              Collections.singletonList("--profile"),
 +              "Databricks CLI does not support --profile flag. Falling back to --host. "
 +                  + "Please upgrade your CLI to the latest version."));
 +      attempts.add(

Reproduce locally: git range-diff 5e8f476..f88c52a 6b8a57f..25a3779 | Disable: git config gitstack.push-range-diff false

@mihaimitrea-db mihaimitrea-db force-pushed the mihaimitrea-db/stack/cli-attempt-chain branch from 25a3779 to 694521d Compare March 31, 2026 15:01
@mihaimitrea-db
Copy link
Copy Markdown
Contributor Author

Range-diff: stack/cli-force-refresh (25a3779 -> 694521d)
databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
@@ -1,10 +1,8 @@
 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
 +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
- import java.time.ZoneId;
  import java.time.format.DateTimeFormatter;
  import java.time.format.DateTimeParseException;
-+import java.util.ArrayList;
  import java.util.Arrays;
 +import java.util.Collections;
  import java.util.List;
@@ -61,47 +59,32 @@
    public CliTokenSource(
        List<String> cmd,
        String tokenTypeField,
-     this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
-   }
- 
-+  /** Constructs a two-attempt source with --profile to --host fallback. */
-   public CliTokenSource(
-       List<String> cmd,
-       String tokenTypeField,
-     this(cmd, tokenTypeField, accessTokenField, expiryField, env, fallbackCmd, null);
-   }
- 
-+  /** Constructs a source with optional force-refresh, profile, and host fallback chain. */
-   public CliTokenSource(
-       List<String> cmd,
-       String tokenTypeField,
-       Environment env,
-       List<String> fallbackCmd,
-       List<String> forceCmd) {
--    super();
--    this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
+       String accessTokenField,
+       String expiryField,
+       Environment env) {
+-    this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
 +    this(
-+        buildAttempts(forceCmd, cmd, fallbackCmd).stream()
-+            .map(
-+                a ->
-+                    new CliCommand(
-+                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
-+                        a.usedFlags,
-+                        a.fallbackMessage))
-+            .collect(Collectors.toList()),
++        Collections.singletonList(
++            new CliCommand(
++                OSUtils.get(env).getCliExecutableCommand(cmd), Collections.emptyList(), null)),
 +        tokenTypeField,
 +        accessTokenField,
 +        expiryField,
 +        env,
 +        true);
-+  }
-+
+   }
+ 
+-  public CliTokenSource(
+-      List<String> cmd,
 +  /** Creates a CliTokenSource from a pre-built attempt chain. */
 +  static CliTokenSource fromAttempts(
 +      List<CliCommand> attempts,
-+      String tokenTypeField,
-+      String accessTokenField,
-+      String expiryField,
+       String tokenTypeField,
+       String accessTokenField,
+       String expiryField,
+-      Environment env,
+-      List<String> fallbackCmd) {
+-    this(cmd, tokenTypeField, accessTokenField, expiryField, env, fallbackCmd, null);
 +      Environment env) {
 +    return new CliTokenSource(
 +        attempts.stream()
@@ -117,14 +100,20 @@
 +        expiryField,
 +        env,
 +        true);
-+  }
-+
+   }
+ 
+-  public CliTokenSource(
+-      List<String> cmd,
 +  private CliTokenSource(
 +      List<CliCommand> attempts,
-+      String tokenTypeField,
-+      String accessTokenField,
-+      String expiryField,
-+      Environment env,
+       String tokenTypeField,
+       String accessTokenField,
+       String expiryField,
+       Environment env,
+-      List<String> fallbackCmd,
+-      List<String> forceCmd) {
+-    super();
+-    this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
 +      boolean alreadyResolved) {
 +    this.attempts = attempts;
      this.tokenTypeField = tokenTypeField;
@@ -134,35 +123,6 @@
 -    this.fallbackCmd =
 -        fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
 -    this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
-+  }
-+
-+  private static List<CliCommand> buildAttempts(
-+      List<String> forceCmd, List<String> profileCmd, List<String> fallbackCmd) {
-+    List<CliCommand> attempts = new ArrayList<>();
-+
-+    if (forceCmd != null) {
-+      attempts.add(
-+          new CliCommand(
-+              forceCmd,
-+              Arrays.asList("--force-refresh", "--profile"),
-+              "Databricks CLI does not support --force-refresh flag. "
-+                  + "Falling back to regular token fetch. "
-+                  + "Please upgrade your CLI to the latest version."));
-+    }
-+
-+    if (fallbackCmd != null) {
-+      attempts.add(
-+          new CliCommand(
-+              profileCmd,
-+              Collections.singletonList("--profile"),
-+              "Databricks CLI does not support --profile flag. Falling back to --host. "
-+                  + "Please upgrade your CLI to the latest version."));
-+      attempts.add(new CliCommand(fallbackCmd, Collections.emptyList(), null));
-+    } else {
-+      attempts.add(new CliCommand(profileCmd, Collections.emptyList(), null));
-+    }
-+
-+    return attempts;
    }
  
    /**
databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
@@ -15,36 +15,36 @@
 +  List<CliTokenSource.CliCommand> buildAttempts(String cliPath, DatabricksConfig config) {
 +    List<CliTokenSource.CliCommand> attempts = new ArrayList<>();
 +
-+    List<String> profileCmd;
-+    boolean hasHostFallback = false;
++    boolean hasProfile = config.getProfile() != null;
 +
-+    if (config.getProfile() != null) {
-+      profileCmd = buildProfileArgs(cliPath, config);
-+      hasHostFallback = config.getHost() != null;
-+    } else {
-+      profileCmd = buildHostArgs(cliPath, config);
-+    }
++    if (hasProfile) {
++      List<String> profileCmd = buildProfileArgs(cliPath, config);
 +
-+    attempts.add(
-+        new CliTokenSource.CliCommand(
-+            withForceRefresh(profileCmd),
-+            Arrays.asList("--force-refresh", "--profile"),
-+            "Databricks CLI does not support --force-refresh flag. "
-+                + "Falling back to regular token fetch. "
-+                + "Please upgrade your CLI to the latest version."));
-+
-+    if (hasHostFallback) {
 +      attempts.add(
 +          new CliTokenSource.CliCommand(
-+              profileCmd,
-+              Collections.singletonList("--profile"),
-+              "Databricks CLI does not support --profile flag. Falling back to --host. "
++              withForceRefresh(profileCmd),
++              Arrays.asList("--force-refresh", "--profile"),
++              "Databricks CLI does not support --force-refresh flag. "
++                  + "Falling back to regular token fetch. "
 +                  + "Please upgrade your CLI to the latest version."));
++
++      if (config.getHost() != null) {
++        attempts.add(
++            new CliTokenSource.CliCommand(
++                profileCmd,
++                Collections.singletonList("--profile"),
++                "Databricks CLI does not support --profile flag. Falling back to --host. "
++                    + "Please upgrade your CLI to the latest version."));
++        attempts.add(
++            new CliTokenSource.CliCommand(
++                buildHostArgs(cliPath, config), Collections.emptyList(), null));
++      } else {
++        attempts.add(new CliTokenSource.CliCommand(profileCmd, Collections.emptyList(), null));
++      }
++    } else {
 +      attempts.add(
 +          new CliTokenSource.CliCommand(
 +              buildHostArgs(cliPath, config), Collections.emptyList(), null));
-+    } else {
-+      attempts.add(new CliTokenSource.CliCommand(profileCmd, Collections.emptyList(), null));
 +    }
 +
 +    return attempts;
databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java
@@ -0,0 +1,40 @@
+diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java
+--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java
++++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java
+ import java.time.format.DateTimeParseException;
+ import java.util.ArrayList;
+ import java.util.Arrays;
++import java.util.Collections;
+ import java.util.HashMap;
+ import java.util.List;
+ import java.util.Map;
+ 
+   private CliTokenSource makeTokenSource(
+       Environment env, List<String> primaryCmd, List<String> fallbackCmd, List<String> forceCmd) {
++    List<CliTokenSource.CliCommand> attempts = new ArrayList<>();
++
++    if (forceCmd != null) {
++      attempts.add(
++          new CliTokenSource.CliCommand(
++              forceCmd, Arrays.asList("--force-refresh", "--profile"), "force-refresh fallback"));
++    }
++
++    if (fallbackCmd != null) {
++      attempts.add(
++          new CliTokenSource.CliCommand(
++              primaryCmd, Collections.singletonList("--profile"), "profile fallback"));
++      attempts.add(new CliTokenSource.CliCommand(fallbackCmd, Collections.emptyList(), null));
++    } else {
++      attempts.add(new CliTokenSource.CliCommand(primaryCmd, Collections.emptyList(), null));
++    }
++
+     OSUtilities osUtils = mock(OSUtilities.class);
+     when(osUtils.getCliExecutableCommand(any())).thenAnswer(inv -> inv.getArgument(0));
+     try (MockedStatic<OSUtils> mockedOSUtils = mockStatic(OSUtils.class)) {
+       mockedOSUtils.when(() -> OSUtils.get(any())).thenReturn(osUtils);
+-      return new CliTokenSource(
+-          primaryCmd, "token_type", "access_token", "expiry", env, fallbackCmd, forceCmd);
++      return CliTokenSource.fromAttempts(attempts, "token_type", "access_token", "expiry", env);
+     }
+   }
+ 
\ No newline at end of file
databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
@@ -1,6 +1,27 @@
 diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
 +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
+ import static org.junit.jupiter.api.Assertions.*;
+ 
+ import java.util.Arrays;
++import java.util.Collections;
+ import java.util.List;
+ import org.junit.jupiter.api.Test;
+ 
+   private static final String ACCOUNT_ID = "test-account-123";
+   private static final String WORKSPACE_ID = "987654321";
+ 
++  private static final String FORCE_REFRESH_FALLBACK_MSG =
++      "Databricks CLI does not support --force-refresh flag. "
++          + "Falling back to regular token fetch. "
++          + "Please upgrade your CLI to the latest version.";
++  private static final String PROFILE_FALLBACK_MSG =
++      "Databricks CLI does not support --profile flag. Falling back to --host. "
++          + "Please upgrade your CLI to the latest version.";
++
+   private final DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider();
+ 
+   @Test
  
      assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), cmd);
    }
@@ -14,12 +35,21 @@
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
 +    assertEquals(3, attempts.size());
++
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile", "--force-refresh"),
 +        attempts.get(0).cmd);
++    assertEquals(Arrays.asList("--force-refresh", "--profile"), attempts.get(0).usedFlags);
++    assertEquals(FORCE_REFRESH_FALLBACK_MSG, attempts.get(0).fallbackMessage);
++
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), attempts.get(1).cmd);
++    assertEquals(Collections.singletonList("--profile"), attempts.get(1).usedFlags);
++    assertEquals(PROFILE_FALLBACK_MSG, attempts.get(1).fallbackMessage);
++
 +    assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), attempts.get(2).cmd);
++    assertEquals(Collections.emptyList(), attempts.get(2).usedFlags);
++    assertNull(attempts.get(2).fallbackMessage);
 +  }
 +
 +  @Test
@@ -29,11 +59,17 @@
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
 +    assertEquals(2, attempts.size());
++
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile", "--force-refresh"),
 +        attempts.get(0).cmd);
++    assertEquals(Arrays.asList("--force-refresh", "--profile"), attempts.get(0).usedFlags);
++    assertEquals(FORCE_REFRESH_FALLBACK_MSG, attempts.get(0).fallbackMessage);
++
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), attempts.get(1).cmd);
++    assertEquals(Collections.emptyList(), attempts.get(1).usedFlags);
++    assertNull(attempts.get(1).fallbackMessage);
 +  }
 +
 +  @Test
@@ -42,11 +78,11 @@
 +
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
-+    assertEquals(2, attempts.size());
-+    assertEquals(
-+        Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST, "--force-refresh"),
-+        attempts.get(0).cmd);
-+    assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), attempts.get(1).cmd);
++    assertEquals(1, attempts.size());
++
++    assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), attempts.get(0).cmd);
++    assertEquals(Collections.emptyList(), attempts.get(0).usedFlags);
++    assertNull(attempts.get(0).fallbackMessage);
 +  }
 +
 +  @Test
@@ -55,21 +91,13 @@
 +
 +    List<CliTokenSource.CliCommand> attempts = provider.buildAttempts(CLI_PATH, config);
 +
-+    assertEquals(2, attempts.size());
++    assertEquals(1, attempts.size());
++
 +    assertEquals(
 +        Arrays.asList(
-+            CLI_PATH,
-+            "auth",
-+            "token",
-+            "--host",
-+            ACCOUNT_HOST,
-+            "--account-id",
-+            ACCOUNT_ID,
-+            "--force-refresh"),
++            CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID),
 +        attempts.get(0).cmd);
-+    assertEquals(
-+        Arrays.asList(
-+            CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID),
-+        attempts.get(1).cmd);
++    assertEquals(Collections.emptyList(), attempts.get(0).usedFlags);
++    assertNull(attempts.get(0).fallbackMessage);
 +  }
  }
\ No newline at end of file

Reproduce locally: git range-diff 6b8a57f..25a3779 6b8a57f..694521d | Disable: git config gitstack.push-range-diff false

Try `--force-refresh` before the regular CLI command so the SDK can
bypass the CLI's own token cache when the SDK considers its token stale.
If the CLI is too old to recognise `--force-refresh` (or `--profile`),
gracefully fall back to the next command in the chain.

Chain order:
- with profile: forceCmd (--profile --force-refresh) -> profileCmd (--profile) -> fallbackCmd (--host)
- without profile: forceCmd (--host --force-refresh) -> profileCmd (--host)

Azure CLI callers are unchanged; they use constructors that leave
forceCmd null, preserving existing behavior.

Signed-off-by: Mihai Mitrea <mihai.mitrea@databricks.com>
@mihaimitrea-db mihaimitrea-db force-pushed the mihaimitrea-db/stack/cli-attempt-chain branch from 694521d to 0f6baf6 Compare March 31, 2026 16:03
@mihaimitrea-db
Copy link
Copy Markdown
Contributor Author

Range-diff: stack/cli-force-refresh (694521d -> 0f6baf6)
NEXT_CHANGELOG.md
@@ -4,7 +4,7 @@
  ### Documentation
  
  ### Internal Changes
-+* Generalize CLI token source into a progressive command list for forward-compatible flag support.
++* Generalized CLI token source into a progressive command attempt list, replacing the fixed three-field approach with an extensible chain.
  
  ### API Changes
  * Add `createCatalog()`, `createSyncedTable()`, `deleteCatalog()`, `deleteSyncedTable()`, `getCatalog()` and `getSyncedTable()` methods for `workspaceClient.postgres()` service.
\ No newline at end of file
databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java
@@ -13,9 +13,13 @@
  public class CliTokenSource implements TokenSource {
    private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
  
--  // forceCmd is tried before profileCmd when non-null. If the CLI rejects
--  // --force-refresh or --profile, execution falls through to profileCmd.
--  private List<String> forceCmd;
+-  private List<String> cmd;
+-  private List<String> fallbackCmd;
+-  private List<String> secondFallbackCmd;
+-  private String tokenTypeField;
+-  private String accessTokenField;
+-  private String expiryField;
+-  private Environment env;
 +  /**
 +   * Describes a CLI command with an optional warning message emitted when falling through to the
 +   * next command in the chain.
@@ -26,15 +30,7 @@
 +    // Flags used by this command (e.g. "--force-refresh", "--profile"). Used to distinguish
 +    // "unknown flag" errors (which trigger fallback) from real auth errors (which propagate).
 +    final List<String> usedFlags;
- 
--  private List<String> profileCmd;
--  private String tokenTypeField;
--  private String accessTokenField;
--  private String expiryField;
--  private Environment env;
--  // fallbackCmd is tried when profileCmd fails with "unknown flag: --profile",
--  // indicating the CLI is too old to support --profile.
--  private List<String> fallbackCmd;
++
 +    final String fallbackMessage;
 +
 +    CliCommand(List<String> cmd, List<String> usedFlags, String fallbackMessage) {
@@ -63,66 +59,60 @@
        String expiryField,
        Environment env) {
 -    this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
-+    this(
-+        Collections.singletonList(
-+            new CliCommand(
-+                OSUtils.get(env).getCliExecutableCommand(cmd), Collections.emptyList(), null)),
-+        tokenTypeField,
-+        accessTokenField,
-+        expiryField,
-+        env,
-+        true);
++    this(cmd, null, tokenTypeField, accessTokenField, expiryField, env);
    }
  
 -  public CliTokenSource(
--      List<String> cmd,
 +  /** Creates a CliTokenSource from a pre-built attempt chain. */
 +  static CliTokenSource fromAttempts(
 +      List<CliCommand> attempts,
-       String tokenTypeField,
-       String accessTokenField,
-       String expiryField,
--      Environment env,
--      List<String> fallbackCmd) {
--    this(cmd, tokenTypeField, accessTokenField, expiryField, env, fallbackCmd, null);
++      String tokenTypeField,
++      String accessTokenField,
++      String expiryField,
 +      Environment env) {
-+    return new CliTokenSource(
-+        attempts.stream()
-+            .map(
-+                a ->
-+                    new CliCommand(
-+                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
-+                        a.usedFlags,
-+                        a.fallbackMessage))
-+            .collect(Collectors.toList()),
-+        tokenTypeField,
-+        accessTokenField,
-+        expiryField,
-+        env,
-+        true);
-   }
- 
--  public CliTokenSource(
--      List<String> cmd,
++    return new CliTokenSource(null, attempts, tokenTypeField, accessTokenField, expiryField, env);
++  }
++
 +  private CliTokenSource(
+       List<String> cmd,
 +      List<CliCommand> attempts,
        String tokenTypeField,
        String accessTokenField,
        String expiryField,
-       Environment env,
+-      Environment env,
 -      List<String> fallbackCmd,
--      List<String> forceCmd) {
--    super();
--    this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
-+      boolean alreadyResolved) {
-+    this.attempts = attempts;
+-      List<String> secondFallbackCmd) {
+-    this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
++      Environment env) {
++    if (attempts != null) {
++      this.attempts =
++          attempts.stream()
++              .map(
++                  a ->
++                      new CliCommand(
++                          OSUtils.get(env).getCliExecutableCommand(a.cmd),
++                          a.usedFlags,
++                          a.fallbackMessage))
++              .collect(Collectors.toList());
++    } else {
++      this.attempts =
++          Collections.singletonList(
++              new CliCommand(
++                  OSUtils.get(env).getCliExecutableCommand(cmd), Collections.emptyList(), null));
++    }
++    if (this.attempts.isEmpty()) {
++      throw new DatabricksException("cannot get access token: no CLI commands configured");
++    }
      this.tokenTypeField = tokenTypeField;
      this.accessTokenField = accessTokenField;
      this.expiryField = expiryField;
      this.env = env;
 -    this.fallbackCmd =
 -        fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
--    this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
+-    this.secondFallbackCmd =
+-        secondFallbackCmd != null
+-            ? OSUtils.get(env).getCliExecutableCommand(secondFallbackCmd)
+-            : null;
    }
  
    /**
@@ -144,24 +134,8 @@
          : e.getMessage();
    }
  
--  private boolean isUnknownFlagError(String errorText, String flag) {
--    return errorText != null && errorText.contains("unknown flag: " + flag);
--  }
--
--  private Token execProfileCmdWithFallback() {
--    try {
--      return execCliCommand(this.profileCmd);
--    } catch (IOException e) {
--      String textToCheck = getErrorText(e);
--      if (fallbackCmd != null && isUnknownFlagError(textToCheck, "--profile")) {
--        LOG.warn(
--            "Databricks CLI does not support --profile flag. Falling back to --host. "
--                + "Please upgrade your CLI to the latest version.");
--        try {
--          return execCliCommand(this.fallbackCmd);
--        } catch (IOException fallbackException) {
--          throw new DatabricksException(fallbackException.getMessage(), fallbackException);
--        }
+-  private boolean isUnknownFlagError(String errorText) {
+-    return errorText != null && errorText.contains("unknown flag:");
 +  private static boolean isUnknownFlagError(String errorText, List<String> flags) {
 +    if (errorText == null) {
 +      return false;
@@ -169,33 +143,36 @@
 +    for (String flag : flags) {
 +      if (errorText.contains("unknown flag: " + flag)) {
 +        return true;
-       }
--      throw new DatabricksException(e.getMessage(), e);
-     }
++      }
++    }
 +    return false;
    }
  
    @Override
    public Token getToken() {
--    if (forceCmd == null) {
--      return execProfileCmdWithFallback();
-+    if (attempts.isEmpty()) {
-+      throw new DatabricksException("cannot get access token: no CLI commands configured");
-     }
+-    try {
+-      return execCliCommand(this.cmd);
+-    } catch (IOException e) {
+-      if (fallbackCmd != null && isUnknownFlagError(getErrorText(e))) {
+-        LOG.warn(
+-            "CLI does not support some flags used by this SDK. "
+-                + "Falling back to a compatible command. "
+-                + "Please upgrade your CLI to the latest version.");
+-      } else {
+-        throw new DatabricksException(e.getMessage(), e);
+-      }
+-    }
++    IOException lastException = null;
  
 -    try {
--      return execCliCommand(this.forceCmd);
+-      return execCliCommand(this.fallbackCmd);
 -    } catch (IOException e) {
--      String textToCheck = getErrorText(e);
--      if (isUnknownFlagError(textToCheck, "--force-refresh")
--          || isUnknownFlagError(textToCheck, "--profile")) {
+-      if (secondFallbackCmd != null && isUnknownFlagError(getErrorText(e))) {
 -        LOG.warn(
--            "Databricks CLI does not support --force-refresh flag. "
--                + "Falling back to regular token fetch. "
+-            "CLI does not support some flags used by this SDK. "
+-                + "Falling back to a compatible command. "
 -                + "Please upgrade your CLI to the latest version.");
--        return execProfileCmdWithFallback();
-+    IOException lastException = null;
-+
+-      } else {
 +    for (int i = 0; i < attempts.size(); i++) {
 +      CliCommand attempt = attempts.get(i);
 +      try {
@@ -208,11 +185,15 @@
 +          lastException = e;
 +          continue;
 +        }
-+        throw new DatabricksException(e.getMessage(), e);
+         throw new DatabricksException(e.getMessage(), e);
        }
+     }
+ 
+-    try {
+-      return execCliCommand(this.secondFallbackCmd);
+-    } catch (IOException e) {
 -      throw new DatabricksException(e.getMessage(), e);
-     }
-+
+-    }
 +    throw new DatabricksException(lastException.getMessage(), lastException);
    }
  }
\ No newline at end of file
databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
@@ -1,21 +1,14 @@
 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
 +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java
+     return forceCmd;
    }
  
-   private static List<String> withForceRefresh(List<String> cmd) {
--    List<String> forceCmd = new ArrayList<>(cmd);
--    forceCmd.add("--force-refresh");
--    return forceCmd;
-+    List<String> result = new ArrayList<>(cmd);
-+    result.add("--force-refresh");
-+    return result;
-+  }
-+
 +  List<CliTokenSource.CliCommand> buildAttempts(String cliPath, DatabricksConfig config) {
 +    List<CliTokenSource.CliCommand> attempts = new ArrayList<>();
 +
 +    boolean hasProfile = config.getProfile() != null;
++    boolean hasHost = config.getHost() != null;
 +
 +    if (hasProfile) {
 +      List<String> profileCmd = buildProfileArgs(cliPath, config);
@@ -28,49 +21,52 @@
 +                  + "Falling back to regular token fetch. "
 +                  + "Please upgrade your CLI to the latest version."));
 +
-+      if (config.getHost() != null) {
-+        attempts.add(
-+            new CliTokenSource.CliCommand(
-+                profileCmd,
-+                Collections.singletonList("--profile"),
-+                "Databricks CLI does not support --profile flag. Falling back to --host. "
-+                    + "Please upgrade your CLI to the latest version."));
-+        attempts.add(
-+            new CliTokenSource.CliCommand(
-+                buildHostArgs(cliPath, config), Collections.emptyList(), null));
-+      } else {
-+        attempts.add(new CliTokenSource.CliCommand(profileCmd, Collections.emptyList(), null));
-+      }
-+    } else {
 +      attempts.add(
 +          new CliTokenSource.CliCommand(
++              profileCmd,
++              Collections.singletonList("--profile"),
++              "Databricks CLI does not support --profile flag. Falling back to --host. "
++                  + "Please upgrade your CLI to the latest version."));
++    }
++
++    if (hasHost) {
++      attempts.add(
++          new CliTokenSource.CliCommand(
 +              buildHostArgs(cliPath, config), Collections.emptyList(), null));
 +    }
 +
 +    return attempts;
-   }
- 
++  }
++
    private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
+     String cliPath = config.getDatabricksCliPath();
+     if (cliPath == null) {
        return null;
      }
  
--    List<String> profileCmd;
+-    List<String> cmd;
 -    List<String> fallbackCmd = null;
--    List<String> forceCmd;
+-    List<String> secondFallbackCmd = null;
 -
 -    if (config.getProfile() != null) {
--      profileCmd = buildProfileArgs(cliPath, config);
--      forceCmd = withForceRefresh(profileCmd);
+-      List<String> profileArgs = buildProfileArgs(cliPath, config);
+-      cmd = withForceRefresh(profileArgs);
+-      fallbackCmd = profileArgs;
 -      if (config.getHost() != null) {
--        fallbackCmd = buildHostArgs(cliPath, config);
+-        secondFallbackCmd = buildHostArgs(cliPath, config);
 -      }
 -    } else {
--      profileCmd = buildHostArgs(cliPath, config);
--      forceCmd = withForceRefresh(profileCmd);
+-      cmd = buildHostArgs(cliPath, config);
 -    }
 -
 -    return new CliTokenSource(
--        profileCmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd, forceCmd);
+-        cmd,
+-        "token_type",
+-        "access_token",
+-        "expiry",
+-        config.getEnv(),
+-        fallbackCmd,
+-        secondFallbackCmd);
 +    return CliTokenSource.fromAttempts(
 +        buildAttempts(cliPath, config), "token_type", "access_token", "expiry", config.getEnv());
    }
databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java
@@ -10,22 +10,29 @@
  import java.util.Map;
  
    private CliTokenSource makeTokenSource(
-       Environment env, List<String> primaryCmd, List<String> fallbackCmd, List<String> forceCmd) {
+       Environment env, List<String> cmd, List<String> fallbackCmd, List<String> secondFallbackCmd) {
 +    List<CliTokenSource.CliCommand> attempts = new ArrayList<>();
 +
-+    if (forceCmd != null) {
-+      attempts.add(
-+          new CliTokenSource.CliCommand(
-+              forceCmd, Arrays.asList("--force-refresh", "--profile"), "force-refresh fallback"));
-+    }
++    attempts.add(
++        new CliTokenSource.CliCommand(
++            cmd,
++            fallbackCmd != null
++                ? Arrays.asList("--force-refresh", "--profile")
++                : Collections.emptyList(),
++            fallbackCmd != null ? "fallback" : null));
 +
 +    if (fallbackCmd != null) {
 +      attempts.add(
 +          new CliTokenSource.CliCommand(
-+              primaryCmd, Collections.singletonList("--profile"), "profile fallback"));
-+      attempts.add(new CliTokenSource.CliCommand(fallbackCmd, Collections.emptyList(), null));
-+    } else {
-+      attempts.add(new CliTokenSource.CliCommand(primaryCmd, Collections.emptyList(), null));
++              fallbackCmd,
++              secondFallbackCmd != null
++                  ? Collections.singletonList("--profile")
++                  : Collections.emptyList(),
++              secondFallbackCmd != null ? "second fallback" : null));
++    }
++
++    if (secondFallbackCmd != null) {
++      attempts.add(new CliTokenSource.CliCommand(secondFallbackCmd, Collections.emptyList(), null));
 +    }
 +
      OSUtilities osUtils = mock(OSUtilities.class);
@@ -33,7 +40,7 @@
      try (MockedStatic<OSUtils> mockedOSUtils = mockStatic(OSUtils.class)) {
        mockedOSUtils.when(() -> OSUtils.get(any())).thenReturn(osUtils);
 -      return new CliTokenSource(
--          primaryCmd, "token_type", "access_token", "expiry", env, fallbackCmd, forceCmd);
+-          cmd, "token_type", "access_token", "expiry", env, fallbackCmd, secondFallbackCmd);
 +      return CliTokenSource.fromAttempts(attempts, "token_type", "access_token", "expiry", env);
      }
    }
databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java
@@ -68,8 +68,8 @@
 +
 +    assertEquals(
 +        Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile"), attempts.get(1).cmd);
-+    assertEquals(Collections.emptyList(), attempts.get(1).usedFlags);
-+    assertNull(attempts.get(1).fallbackMessage);
++    assertEquals(Collections.singletonList("--profile"), attempts.get(1).usedFlags);
++    assertEquals(PROFILE_FALLBACK_MSG, attempts.get(1).fallbackMessage);
 +  }
 +
 +  @Test

Reproduce locally: git range-diff 6b8a57f..694521d 61686da..0f6baf6 | Disable: git config gitstack.push-range-diff false

@github-actions
Copy link
Copy Markdown

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-java

Inputs:

  • PR number: 752
  • Commit SHA: 0f6baf6f2f1944215fc0c3c1023bab1f584197f8

Checks will be approved automatically on success.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant