Refactor CliTokenSource to use an ordered attempt chain

mihaimitrea-db · mihaimitrea-db · commit f88c52af3f5f · 2026-03-31T13:21:43.000Z
Replace ad-hoc forceCmd/profileCmd/fallbackCmd fields with a generic
ordered attempt model: each CliCommandAttempt describes its command,
the error substrings that trigger fallback to the next attempt, and
an optional log 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 -&gt; profile -&gt; host
- without profile: force+host -&gt; 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 &lt;mihai.mitrea@databricks.com&gt;
diff --git a/NEXT_CHANGELOG.md b/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.
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
@@ -15,8 +15,11 @@
 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;
+import java.util.stream.Collectors;
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,21 +28,21 @@
 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.
+   */
+  static class CliCommand {
+    final List<String> cmd;
+    final List<String> fallbackTriggers;
+    final String fallbackMessage;
 
-  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;
+    CliCommand(List<String> cmd, List<String> fallbackTriggers, String fallbackMessage) {
+      this.cmd = cmd;
+      this.fallbackTriggers = fallbackTriggers != null ? fallbackTriggers : Collections.emptyList();
+      this.fallbackMessage = fallbackMessage;
+    }
+  }
 
   /**
    * Internal exception that carries the clean stderr message but exposes full output for checks.
@@ -57,6 +60,13 @@ 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,
@@ -66,6 +76,7 @@ public CliTokenSource(
     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,
@@ -76,6 +87,7 @@ public CliTokenSource(
     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,
@@ -84,15 +96,89 @@ public CliTokenSource(
       Environment env,
       List<String> fallbackCmd,
       List<String> forceCmd) {
-    super();
-    this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
+    this(
+        buildAttempts(forceCmd, cmd, fallbackCmd).stream()
+            .map(
+                a ->
+                    new CliCommand(
+                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
+                        a.fallbackTriggers,
+                        a.fallbackMessage))
+            .collect(Collectors.toList()),
+        tokenTypeField,
+        accessTokenField,
+        expiryField,
+        env,
+        true);
+  }
+
+  /** Creates a CliTokenSource from a pre-built attempt chain. */
+  static CliTokenSource fromAttempts(
+      List<CliCommand> attempts,
+      String tokenTypeField,
+      String accessTokenField,
+      String expiryField,
+      Environment env) {
+    return new CliTokenSource(
+        attempts.stream()
+            .map(
+                a ->
+                    new CliCommand(
+                        OSUtils.get(env).getCliExecutableCommand(a.cmd),
+                        a.fallbackTriggers,
+                        a.fallbackMessage))
+            .collect(Collectors.toList()),
+        tokenTypeField,
+        accessTokenField,
+        expiryField,
+        env,
+        true);
+  }
+
+  private CliTokenSource(
+      List<CliCommand> attempts,
+      String tokenTypeField,
+      String accessTokenField,
+      String expiryField,
+      Environment env,
+      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;
+  }
+
+  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) {
+    List<CliCommand> attempts = new ArrayList<>();
+
+    if (forceCmd != null) {
+      attempts.add(
+          new CliCommand(
+              forceCmd,
+              Arrays.asList(UNKNOWN_FORCE_REFRESH_FLAG, UNKNOWN_PROFILE_FLAG),
+              "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(UNKNOWN_PROFILE_FLAG),
+              "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;
   }
 
   /**
@@ -170,54 +256,39 @@ 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(flag);
-  }
-
-  private Token execProfileCmdWithFallback() {
-    try {
-      return execCliCommand(this.profileCmd);
-    } catch (IOException e) {
-      String textToCheck = getErrorText(e);
-      if (fallbackCmd != null && isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
-        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);
-        }
-      }
-      throw new DatabricksException(e.getMessage(), e);
+  private static boolean shouldFallback(CliCommand attempt, String errorText) {
+    if (errorText == null) {
+      return false;
     }
+    return attempt.fallbackTriggers.stream().anyMatch(errorText::contains);
   }
 
   @Override
   public Token getToken() {
-    if (forceCmd == null) {
-      return execProfileCmdWithFallback();
-    }
+    IOException lastException = null;
 
-    try {
-      return execCliCommand(this.forceCmd);
-    } catch (IOException e) {
-      String textToCheck = getErrorText(e);
-      if (isUnknownFlagError(textToCheck, UNKNOWN_FORCE_REFRESH_FLAG)
-          || isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
-        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();
+    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 (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);
   }
 }
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
@@ -69,15 +69,56 @@ List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
     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");
-    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<>();
+
+    List<String> profileCmd;
+    boolean hasHostFallback = false;
+
+    if (config.getProfile() != null) {
+      profileCmd = buildProfileArgs(cliPath, config);
+      hasHostFallback = config.getHost() != null;
+    } else {
+      profileCmd = buildHostArgs(cliPath, config);
+    }
+
+    attempts.add(
+        new CliTokenSource.CliCommand(
+            withForceRefresh(profileCmd),
+            Arrays.asList(UNKNOWN_FORCE_REFRESH_FLAG, UNKNOWN_PROFILE_FLAG),
+            "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(UNKNOWN_PROFILE_FLAG),
+              "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));
+    }
+
+    return attempts;
   }
 
   private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
@@ -90,23 +131,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
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
