Generalize CLI token source into progressive command list

Replace the three explicit command fields (forceCmd, profileCmd, hostCmd)
with a []cliCommand list and an activeCommandIndex. Feature flags are
defined statically in cliFeatureFlags and stripped progressively to
build commands for older CLI versions.

Token() now iterates from activeCommandIndex, falling back on unknown
flag errors and caching the working command index for subsequent calls.
Adding future flags (e.g. --scopes) requires only a one-line addition
to the cliFeatureFlags slice.

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

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -24,6 +24,7 @@
 ### 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,96 +32,114 @@ type cliTokenResponse struct {
 	Expiry      string `json:"expiry"`
 }
 
+// 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
+}
+
 // CliTokenSource fetches OAuth tokens by shelling out to the Databricks CLI.
-// Commands are tried in order: forceCmd -> profileCmd -> hostCmd, progressively
-// falling back to simpler invocations for older CLI versions.
+// 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 {
-	// forceCmd appends --force-refresh to the base command (profileCmd when a
-	// profile is configured, hostCmd otherwise) to bypass the CLI's token cache.
-	// Nil when neither profile nor host is set.
-	// CLI support: >= v0.296.0 (databricks/cli#4767).
-	forceCmd []string
-
-	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
-	// CLI support: >= v0.207.1 (databricks/cli#855).
-	profileCmd []string
-
-	// hostCmd uses --host as a fallback for CLIs that predate --profile support.
-	// Nil when cfg.Host is empty.
-	hostCmd []string
+	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
 	}
-	forceCmd, profileCmd, hostCmd := buildCliCommands(cliPath, cfg)
-	return &CliTokenSource{forceCmd: forceCmd, profileCmd: 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, three commands are built: a --force-refresh variant
-// (based on profileCmd), a plain --profile variant, and (when host is available)
-// a --host fallback. When cfg.Profile is empty, --force-refresh is based on the
-// --host command instead.
-func buildCliCommands(cliPath string, cfg *Config) ([]string, []string, []string) {
-	var forceCmd, profileCmd, 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 != "" {
-		profileCmd = []string{cliPath, "auth", "token", "--profile", cfg.Profile}
+		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 != "" {
-		hostCmd = buildHostCommand(cliPath, cfg)
-	}
-	if profileCmd != nil {
-		forceCmd = append(profileCmd, "--force-refresh")
-	} else if hostCmd != nil {
-		forceCmd = append(hostCmd, "--force-refresh")
-	}
-	return forceCmd, profileCmd, hostCmd
-}
-
-// 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)
+		hostArgs := []string{cliPath, "auth", "token", "--host", cfg.Host}
+		switch cfg.HostType() {
+		case AccountHost:
+			hostArgs = append(hostArgs, "--account-id", cfg.AccountID)
+		}
+		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 cmd
+	return commands
 }
 
 // Token fetches an OAuth token by shelling out to the Databricks CLI.
-// Commands are tried in order: forceCmd -> profileCmd -> hostCmd, skipping nil
-// entries. Each command falls through to the next on "unknown flag" errors,
-// logging a warning about the unsupported feature.
 func (c *CliTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
-	if c.forceCmd != nil {
-		tok, err := c.execCliCommand(ctx, c.forceCmd)
-		if err == nil {
-			return tok, nil
-		}
-		if !isUnknownFlagError(err, "--force-refresh") && !isUnknownFlagError(err, "--profile") {
-			return nil, err
-		}
-		logger.Warnf(ctx, "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.")
+	// 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 c.probeAndExec(ctx)
+}
 
-	if c.profileCmd != nil {
-		tok, err := c.execCliCommand(ctx, c.profileCmd)
+// 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
 		}
-		if !isUnknownFlagError(err, "--profile") {
+		lastCommand := i == len(c.commands)-1
+		if lastCommand || !isUnknownFlagError(err, cmd.usedFlags) {
 			return nil, err
 		}
-		logger.Warnf(ctx, "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.")
-	}
-
-	if c.hostCmd == nil {
-		return nil, fmt.Errorf("cannot get access token: no CLI commands available")
+		logger.Warnf(ctx, cmd.warningMessage)
 	}
-	return c.execCliCommand(ctx, c.hostCmd)
+	return nil, fmt.Errorf("cannot get access token: no CLI commands configured")
 }
 
 func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) {
@@ -129,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: %q", 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)
 	}
@@ -148,10 +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 given flag (e.g. "--profile", "--force-refresh").
-func isUnknownFlagError(err error, flag string) bool {
-	return strings.Contains(err.Error(), "unknown flag: "+flag)
+// 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.