Refactor CliTokenSource to use an ordered attempt chain

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -14,6 +14,7 @@
 ### Documentation
 
 ### Internal Changes
+* 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.

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

@@ -16,7 +16,10 @@
 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.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,13 +28,25 @@
 public class CliTokenSource implements TokenSource {
   private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
 
-  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.
+   */
+  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;
+
+    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.
@@ -49,34 +64,72 @@ String getFullOutput() {
     }
   }
 
+  private final List<CliCommand> commands;
+
+  // Index of the CLI command known to work, or -1 if not yet resolved. Once
+  // resolved it never changes — older CLIs don't gain new flags. We use
+  // AtomicInteger instead of synchronization because probing must be retryable
+  // on transient errors: concurrent callers may redundantly probe, but all
+  // converge to the same index.
+  private final AtomicInteger activeCommandIndex = new AtomicInteger(-1);
+
+  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(cmd, null, tokenTypeField, accessTokenField, expiryField, env);
   }
 
-  public CliTokenSource(
+  /** Creates a CliTokenSource from a pre-built command chain. */
+  static CliTokenSource fromCommands(
+      List<CliCommand> commands,
+      String tokenTypeField,
+      String accessTokenField,
+      String expiryField,
+      Environment env) {
+    return new CliTokenSource(null, commands, tokenTypeField, accessTokenField, expiryField, env);
+  }
+
+  private CliTokenSource(
       List<String> cmd,
+      List<CliCommand> commands,
       String tokenTypeField,
       String accessTokenField,
       String expiryField,
-      Environment env,
-      List<String> fallbackCmd,
-      List<String> secondFallbackCmd) {
-    this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
+      Environment env) {
+    if (commands != null && !commands.isEmpty()) {
+      this.commands =
+          commands.stream()
+              .map(
+                  a ->
+                      new CliCommand(
+                          OSUtils.get(env).getCliExecutableCommand(a.cmd),
+                          a.usedFlags,
+                          a.fallbackMessage))
+              .collect(Collectors.toList());
+    } else if (cmd != null) {
+      if (commands != null && commands.isEmpty()) {
+        LOG.warn("No CLI commands configured. Falling back to the default command.");
+      }
+      this.commands =
+          Collections.singletonList(
+              new CliCommand(
+                  OSUtils.get(env).getCliExecutableCommand(cmd), Collections.emptyList(), null));
+    } else {
+      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.secondFallbackCmd =
-        secondFallbackCmd != null
-            ? OSUtils.get(env).getCliExecutableCommand(secondFallbackCmd)
-            : null;
   }
 
   /**
@@ -137,8 +190,9 @@ 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.
+        // getMessage() carries the clean stderr message for user-facing errors;
+        // getFullOutput() includes both streams so isUnknownFlagError can detect
+        // "unknown flag:" regardless of which stream the CLI wrote it to.
         throw new CliCommandException("cannot get access token: " + stderr, stdout + "\n" + stderr);
       }
       JsonNode jsonNode = new ObjectMapper().readTree(stdout);
@@ -154,48 +208,61 @@ 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) {
-    return errorText != null && errorText.contains("unknown flag:");
+  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;
+      }
+    }
+    return false;
   }
 
   @Override
   public Token getToken() {
-    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 {
+    int idx = activeCommandIndex.get();
+    if (idx >= 0) {
+      try {
+        return execCliCommand(commands.get(idx).cmd);
+      } catch (IOException e) {
         throw new DatabricksException(e.getMessage(), e);
       }
     }
+    return probeAndExec();
+  }
 
-    try {
-      return execCliCommand(this.fallbackCmd);
-    } catch (IOException e) {
-      if (secondFallbackCmd != 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 {
+  /**
+   * Walks the command list from most-featured to simplest, looking for a CLI command that succeeds.
+   * When a command fails with "unknown flag" for one of its {@link CliCommand#usedFlags}, it logs a
+   * warning and tries the next. On success, {@link #activeCommandIndex} is stored so future calls
+   * skip probing.
+   */
+  private Token probeAndExec() {
+    for (int i = 0; i < commands.size(); i++) {
+      CliCommand command = commands.get(i);
+      try {
+        Token token = execCliCommand(command.cmd);
+        activeCommandIndex.set(i);
+        return token;
+      } catch (IOException e) {
+        if (i + 1 < commands.size() && isUnknownFlagError(getErrorText(e), command.usedFlags)) {
+          if (command.fallbackMessage != null) {
+            LOG.warn(command.fallbackMessage);
+          }
+          continue;
+        }
         throw new DatabricksException(e.getMessage(), e);
       }
     }
 
-    try {
-      return execCliCommand(this.secondFallbackCmd);
-    } catch (IOException e) {
-      throw new DatabricksException(e.getMessage(), e);
-    }
+    throw new DatabricksException("cannot get access token: all CLI commands failed");
   }
 }

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

@@ -80,6 +80,40 @@ private static List<String> withForceRefresh(List<String> cmd) {
     return forceCmd;
   }
 
+  List<CliTokenSource.CliCommand> buildCommands(String cliPath, DatabricksConfig config) {
+    List<CliTokenSource.CliCommand> commands = new ArrayList<>();
+
+    boolean hasProfile = config.getProfile() != null;
+    boolean hasHost = config.getHost() != null;
+
+    if (hasProfile) {
+      List<String> profileCmd = buildProfileArgs(cliPath, config);
+
+      commands.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."));
+
+      commands.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) {
+      commands.add(
+          new CliTokenSource.CliCommand(
+              buildHostArgs(cliPath, config), Collections.emptyList(), null));
+    }
+
+    return commands;
+  }
+
   private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
     String cliPath = config.getDatabricksCliPath();
     if (cliPath == null) {
@@ -90,29 +124,8 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
       return null;
     }
 
-    List<String> cmd;
-    List<String> fallbackCmd = null;
-    List<String> secondFallbackCmd = null;
-
-    if (config.getProfile() != null) {
-      List<String> profileArgs = buildProfileArgs(cliPath, config);
-      cmd = withForceRefresh(profileArgs);
-      fallbackCmd = profileArgs;
-      if (config.getHost() != null) {
-        secondFallbackCmd = buildHostArgs(cliPath, config);
-      }
-    } else {
-      cmd = buildHostArgs(cliPath, config);
-    }
-
-    return new CliTokenSource(
-        cmd,
-        "token_type",
-        "access_token",
-        "expiry",
-        config.getEnv(),
-        fallbackCmd,
-        secondFallbackCmd);
+    return CliTokenSource.fromCommands(
+        buildCommands(cliPath, config), "token_type", "access_token", "expiry", config.getEnv());
   }
 
   @Override