Refactor CliTokenSource to use an ordered attempt chain

Replace ad-hoc forceCmd/profileCmd/fallbackCmd fields with a generic
ordered attempt model: each CliCommand describes its command, the flags
it uses (for detecting "unknown flag" errors), and an optional warning
message.

getToken() is now a simple loop over the attempt chain instead of
nested if-else branches. This makes adding future CLI flag fallbacks
straightforward without introducing more fields.

DatabricksCliCredentialsProvider builds the attempt chain explicitly:
- with profile: force+profile -> profile -> host
- without profile: force+host -> host

Old constructors delegate to the new model. Azure CLI callers use
the 5-arg constructor which creates a single-attempt chain, so their
behavior is unchanged.

Signed-off-by: Mihai Mitrea <mihai.mitrea@databricks.com>

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -14,6 +14,7 @@
 ### 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.

databricks-sdk-java/src/main/java/com/databricks/sdk/core/CliTokenSource.java

@@ -16,7 +16,9 @@
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,18 +27,25 @@
 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;
+  /**
+   * 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;
+
+    // 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) {
+      this.cmd = cmd;
+      this.usedFlags = usedFlags != null ? usedFlags : Collections.emptyList();
+      this.fallbackMessage = fallbackMessage;
+    }
+  }
 
   /**
    * Internal exception that carries the clean stderr message but exposes full output for checks.
@@ -54,42 +63,65 @@ String getFullOutput() {
     }
   }
 
+  private final List<CliCommand> attempts;
+  private final String tokenTypeField;
+  private final String accessTokenField;
+  private final String expiryField;
+  private final Environment env;
+
+  /** Constructs a single-attempt source. Used by Azure CLI and simple callers. */
   public CliTokenSource(
       List<String> cmd,
       String tokenTypeField,
       String accessTokenField,
       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);
   }
 
-  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);
+      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,
+  private CliTokenSource(
+      List<CliCommand> attempts,
       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;
     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;
   }
 
   /**
@@ -150,8 +182,6 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
         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);
@@ -167,54 +197,48 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
     }
   }
 
-  private String getErrorText(IOException e) {
+  private static String getErrorText(IOException e) {
     return e instanceof CliCommandException
         ? ((CliCommandException) e).getFullOutput()
         : 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 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 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.forceCmd);
-    } catch (IOException e) {
-      String textToCheck = getErrorText(e);
-      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() && isUnknownFlagError(getErrorText(e), attempt.usedFlags)) {
+          if (attempt.fallbackMessage != null) {
+            LOG.warn(attempt.fallbackMessage);
+          }
+          lastException = e;
+          continue;
+        }
+        throw new DatabricksException(e.getMessage(), e);
       }
-      throw new DatabricksException(e.getMessage(), e);
     }
+
+    throw new DatabricksException(lastException.getMessage(), lastException);
   }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java

@@ -75,9 +75,47 @@ List<String> buildProfileArgs(String cliPath, DatabricksConfig config) {
   }
 
   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;
+
+    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 (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));
+    }
+
+    return attempts;
   }
 
   private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
@@ -90,23 +128,8 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
       return null;
     }
 
-    List<String> profileCmd;
-    List<String> fallbackCmd = null;
-    List<String> forceCmd;
-
-    if (config.getProfile() != null) {
-      profileCmd = buildProfileArgs(cliPath, config);
-      forceCmd = withForceRefresh(profileCmd);
-      if (config.getHost() != null) {
-        fallbackCmd = buildHostArgs(cliPath, config);
-      }
-    } else {
-      profileCmd = buildHostArgs(cliPath, config);
-      forceCmd = withForceRefresh(profileCmd);
-    }
-
-    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());
   }
 
   @Override

databricks-sdk-java/src/test/java/com/databricks/sdk/core/CliTokenSourceTest.java

@@ -24,6 +24,7 @@
 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;
@@ -226,12 +227,28 @@ private CliTokenSource makeTokenSource(
 
   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);
     }
   }