Refactor CliTokenSource to use an ordered attempt chain

NEXT_CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### New Features and Improvements
 * Added automatic detection of AI coding agents (Antigravity, Claude Code, Cline, Codex, Copilot CLI, Cursor, Gemini CLI, OpenCode) in the user-agent string. The SDK now appends `agent/<name>` to HTTP request headers when running inside a known AI agent environment.
+* Pass `--force-refresh` to Databricks CLI `auth token` command so the SDK always receives a fresh token instead of a potentially stale one from the CLI's internal cache. Falls back gracefully on older CLIs that do not support this flag.
 
 ### Bug Fixes
 * Fixed Databricks CLI authentication to detect when the cached token's scopes don't match the SDK's configured scopes. Previously, a scope mismatch was silently ignored, causing requests to use wrong permissions. The SDK now raises an error with instructions to re-authenticate.
@@ -13,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,16 +28,25 @@
 public class CliTokenSource implements TokenSource {
   private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
 
-  private List<String> cmd;
-  private String tokenTypeField;
-  private String accessTokenField;
-  private String expiryField;
-  private Environment env;
-  // fallbackCmd is tried when the primary command fails with "unknown flag: --profile",
-  // indicating the CLI is too old to support --profile. Can be removed once support
-  // for CLI versions predating --profile is dropped.
-  // See: https://github.com/databricks/databricks-sdk-go/pull/1497
-  private List<String> fallbackCmd;
+  /**
+   * 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.
@@ -52,30 +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);
+    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) {
-    super();
-    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;
   }
 
   /**
@@ -136,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);
@@ -153,28 +208,61 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
     }
   }
 
+  private static String getErrorText(IOException e) {
+    return e instanceof CliCommandException
+        ? ((CliCommandException) e).getFullOutput()
+        : e.getMessage();
+  }
+
+  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) {
-      String textToCheck =
-          e instanceof CliCommandException
-              ? ((CliCommandException) e).getFullOutput()
-              : e.getMessage();
-      if (fallbackCmd != null
-          && textToCheck != null
-          && textToCheck.contains("unknown flag: --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);
+    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();
+  }
+
+  /**
+   * 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);
       }
-      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

@@ -69,6 +69,51 @@ List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
     return cmd;
   }
 
+  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<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) {
@@ -79,25 +124,8 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
       return null;
     }
 
-    List<String> cmd;
-    List<String> fallbackCmd = null;
-
-    if (config.getProfile() != null) {
-      // When profile is set, use --profile as the primary command.
-      // The profile contains the full config (host, account_id, etc.).
-      cmd =
-          new ArrayList<>(
-              Arrays.asList(cliPath, "auth", "token", "--profile", config.getProfile()));
-      // Build a --host fallback for older CLIs that don't support --profile.
-      if (config.getHost() != null) {
-        fallbackCmd = buildHostArgs(cliPath, config);
-      }
-    } else {
-      cmd = buildHostArgs(cliPath, config);
-    }
-
-    return new CliTokenSource(
-        cmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd);
+    return CliTokenSource.fromCommands(
+        buildCommands(cliPath, config), "token_type", "access_token", "expiry", config.getEnv());
   }
 
   @Override