Generalize CLI token source into progressive command list

NEXT_CHANGELOG.md

@@ -23,6 +23,8 @@
 
 ### Internal Changes
 
+ * Pass `--force-refresh` to Databricks CLI `auth token` command to bypass the CLI's internal token cache.
+ * Generalize CLI token source into a progressive command list for forward-compatible flag support.
  * Use resolved host type from host metadata in `HostType()` method, falling back to URL pattern matching when metadata is unavailable.
  * Normalize internal token sources on `auth.TokenSource` for proper context propagation ([#1577](https://github.com/databricks/databricks-sdk-go/pull/1577)).
  * Fix `TestAzureGithubOIDCCredentials` hang caused by missing `HTTPTransport` stub: `EnsureResolved` now calls `resolveHostMetadata`, which makes a real network request when no transport is set ([#1550](https://github.com/databricks/databricks-sdk-go/pull/1550)).

config/cli_token_source.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/logger"
@@ -31,62 +32,114 @@ type cliTokenResponse struct {
 	Expiry      string `json:"expiry"`
 }
 
-type CliTokenSource struct {
-	// cmd is the primary command to execute (--profile when available, --host otherwise).
-	cmd []string
+// cliCommand is a single CLI invocation with an associated warning message
+// that is logged when this command fails and we fall back to the next one.
+type cliCommand struct {
+	args []string
+	// usedFlags lists all flags passed in this command. When the CLI reports
+	// "unknown flag: <flag>" for one of these, we fall back to the next command.
+	usedFlags      []string
+	warningMessage string
+}
 
-	// hostCmd is a fallback command using --host, used when the primary --profile
-	// command fails because the CLI is too old to support --profile.
-	hostCmd []string
+// CliTokenSource fetches OAuth tokens by shelling out to the Databricks CLI.
+// It holds a list of commands ordered from most feature-rich to simplest,
+// falling back progressively for older CLI versions that lack newer flags.
+type CliTokenSource struct {
+	commands []cliCommand
+	// activeCommandIndex is the 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 atomic.Int32 instead of sync.Once because
+	// probing must be retryable on transient errors: concurrent callers may
+	// redundantly probe, but all converge to the same index.
+	activeCommandIndex atomic.Int32
 }
 
 func NewCliTokenSource(cfg *Config) (*CliTokenSource, error) {
 	cliPath, err := findDatabricksCli(cfg.DatabricksCliPath)
 	if err != nil {
 		return nil, err
 	}
-	profileCmd, hostCmd := buildCliCommands(cliPath, cfg)
-	return &CliTokenSource{cmd: profileCmd, hostCmd: hostCmd}, nil
+	commands := buildCliCommands(cliPath, cfg)
+	if len(commands) == 0 {
+		return nil, fmt.Errorf("cannot configure CLI token source: neither profile nor host is set")
+	}
+	ts := &CliTokenSource{commands: commands}
+	ts.activeCommandIndex.Store(-1)
+	return ts, nil
 }
 
-// buildCliCommands constructs the CLI commands for fetching an auth token.
-// When cfg.Profile is set, the primary command uses --profile and a fallback
-// --host command is also returned for compatibility with older CLIs.
-// When cfg.Profile is empty, the primary command uses --host and no fallback
-// is needed.
-func buildCliCommands(cliPath string, cfg *Config) (primaryCmd []string, hostCmd []string) {
+// buildCliCommands constructs the list of CLI commands for fetching an auth
+// token, ordered from most feature-rich to simplest. The order defines the
+// fallback chain: when a command fails with an unknown flag error, the next
+// one is tried.
+//
+// When cfg.Profile is set, --force-refresh is based on the --profile command.
+// When cfg.Profile is empty, --force-refresh is based on the --host command
+// instead, so host-only configurations still benefit from cache bypass.
+func buildCliCommands(cliPath string, cfg *Config) []cliCommand {
+	var commands []cliCommand
 	if cfg.Profile != "" {
-		primary := []string{cliPath, "auth", "token", "--profile", cfg.Profile}
-		if cfg.Host != "" {
-			// Build a --host fallback for old CLIs that don't support --profile.
-			return primary, buildHostCommand(cliPath, cfg)
+		commands = append(commands, cliCommand{
+			args:           []string{cliPath, "auth", "token", "--profile", cfg.Profile, "--force-refresh"},
+			usedFlags:      []string{"--force-refresh", "--profile"},
+			warningMessage: "Databricks CLI does not support --force-refresh flag. The CLI's token cache may provide stale tokens. Please upgrade your CLI to the latest version.",
+		})
+		commands = append(commands, cliCommand{
+			args:           []string{cliPath, "auth", "token", "--profile", cfg.Profile},
+			usedFlags:      []string{"--profile"},
+			warningMessage: "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.",
+		})
+	}
+	if cfg.Host != "" {
+		hostArgs := []string{cliPath, "auth", "token", "--host", cfg.Host}
+		switch cfg.HostType() {
+		case AccountHost:
+			hostArgs = append(hostArgs, "--account-id", cfg.AccountID)
 		}
-		return primary, nil
+		if cfg.Profile == "" {
+			forceArgs := append(hostArgs[:len(hostArgs):len(hostArgs)], "--force-refresh")
+			commands = append(commands, cliCommand{
+				args:           forceArgs,
+				usedFlags:      []string{"--force-refresh"},
+				warningMessage: "Databricks CLI does not support --force-refresh flag. The CLI's token cache may provide stale tokens. Please upgrade your CLI to the latest version.",
+			})
+		}
+		commands = append(commands, cliCommand{args: hostArgs})
 	}
-	return buildHostCommand(cliPath, cfg), nil
+	return commands
 }
 
-// buildHostCommand constructs the legacy --host based CLI command.
-func buildHostCommand(cliPath string, cfg *Config) []string {
-	cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
-	switch cfg.HostType() {
-	case AccountHost:
-		cmd = append(cmd, "--account-id", cfg.AccountID)
+// Token fetches an OAuth token by shelling out to the Databricks CLI.
+func (c *CliTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
+	// If the working command has already been identified, call it directly.
+	// Otherwise, probe each command in order to find one that works.
+	if idx := int(c.activeCommandIndex.Load()); idx >= 0 {
+		return c.execCliCommand(ctx, c.commands[idx].args)
 	}
-	return cmd
+	return c.probeAndExec(ctx)
 }
 
-// Token fetches an OAuth token by shelling out to the Databricks CLI.
-// When a --profile command is configured, it is tried first. If the CLI
-// returns "unknown flag: --profile" (indicating an older CLI version),
-// the fallback --host command is used instead.
-func (c *CliTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
-	tok, err := c.execCliCommand(ctx, c.cmd)
-	if err != nil && c.hostCmd != nil && isUnknownFlagError(err) {
-		logger.Warnf(ctx, "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.")
-		return c.execCliCommand(ctx, c.hostCmd)
+// 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 [cliCommand.usedFlags], it logs a warning and tries the next. Any
+// other error stops probing immediately and is returned to the caller.
+// On success, [activeCommandIndex] is stored so future calls skip probing.
+func (c *CliTokenSource) probeAndExec(ctx context.Context) (*oauth2.Token, error) {
+	for i := range c.commands {
+		cmd := c.commands[i]
+		tok, err := c.execCliCommand(ctx, cmd.args)
+		if err == nil {
+			c.activeCommandIndex.Store(int32(i))
+			return tok, nil
+		}
+		lastCommand := i == len(c.commands)-1
+		if lastCommand || !isUnknownFlagError(err, cmd.usedFlags) {
+			return nil, err
+		}
+		logger.Warnf(ctx, cmd.warningMessage)
 	}
-	return tok, err
+	return nil, fmt.Errorf("cannot get access token: no CLI commands configured")
 }
 
 func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) {
@@ -95,7 +148,13 @@ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oa
 	if err != nil {
 		var exitErr *exec.ExitError
 		if errors.As(err, &exitErr) {
-			return nil, fmt.Errorf("cannot get access token: %s", strings.TrimSpace(string(exitErr.Stderr)))
+			// Format with %q: the CLI's stderr can be multi-line (includes usage
+			// text on unknown-flag errors) and quoting keeps it on one line so
+			// log aggregators don't split a single error across multiple entries.
+			// We intentionally discard exec.ExitError — the stderr text is the
+			// CLI's error contract; exit codes and process state are not useful.
+			return nil, fmt.Errorf("cannot get access token: %q",
+				strings.TrimSpace(string(exitErr.Stderr)))
 		}
 		return nil, fmt.Errorf("cannot get access token: %w", err)
 	}
@@ -114,11 +173,16 @@ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oa
 	}, nil
 }
 
-// isUnknownFlagError returns true if the error indicates the CLI does not
-// recognize the --profile flag. This happens with older CLI versions that
-// predate profile-based token lookup.
-func isUnknownFlagError(err error) bool {
-	return strings.Contains(err.Error(), "unknown flag: --profile")
+// isUnknownFlagError returns true if the error message indicates the CLI does
+// not recognize one of the given flags.
+func isUnknownFlagError(err error, flags []string) bool {
+	msg := err.Error()
+	for _, flag := range flags {
+		if strings.Contains(msg, "unknown flag: "+flag) {
+			return true
+		}
+	}
+	return false
 }
 
 // parseExpiry parses an expiry time string in multiple formats for cross-SDK compatibility.