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

@@ -22,6 +22,7 @@
 
 ### Internal Changes
 
+ * Generalize CLI token source into a progressive command list for forward-compatible flag support.
  * 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)).
  * Bump golang.org/x/crypto from 0.21.0 to 0.45.0 in /examples/slog ([#1566](https://github.com/databricks/databricks-sdk-go/pull/1566)).

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,89 +32,110 @@ type cliTokenResponse struct {
 	Expiry      string `json:"expiry"`
 }
 
-// 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.
-type CliTokenSource struct {
-	// forceCmd uses --profile with --force-refresh to bypass the CLI's token cache.
-	// Nil when cfg.Profile is empty (--force-refresh requires --profile support).
-	forceCmd []string
+// cliError wraps stderr output from a failed CLI invocation.
+type cliError struct {
+	stderr string
+}
 
-	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
-	profileCmd []string
+func (e *cliError) Error() string {
+	return e.stderr
+}
+
+// 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 uses --host as a fallback for CLIs that predate --profile support.
-	// Nil when cfg.Host is empty.
-	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
 	}
-	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,
-// a plain --profile variant, and (when host is available) a --host fallback.
-// When cfg.Profile is empty, only --host is returned — the CLI must support
-// --profile before --force-refresh can be used (monotonic feature assumption).
-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.
+func buildCliCommands(cliPath string, cfg *Config) []cliCommand {
+	var commands []cliCommand
 	if cfg.Profile != "" {
-		profileCmd = []string{cliPath, "auth", "token", "--profile", cfg.Profile}
-		forceCmd = append(profileCmd, "--force-refresh")
+		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)
-	}
-	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)
+		args := []string{cliPath, "auth", "token", "--host", cfg.Host}
+		switch cfg.HostType() {
+		case AccountHost:
+			args = append(args, "--account-id", cfg.AccountID)
+		}
+		commands = append(commands, cliCommand{args: args})
 	}
-	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.
+// If the working command has already been identified, it is called directly.
+// Otherwise, [probeAndExec] tries each command in order to find one that works.
 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 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.
+// 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) {
@@ -122,7 +144,8 @@ 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)))
+			return nil, fmt.Errorf("cannot get access token: %w",
+				&cliError{stderr: strings.TrimSpace(string(exitErr.Stderr))})
 		}
 		return nil, fmt.Errorf("cannot get access token: %w", err)
 	}
@@ -141,14 +164,19 @@ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oa
 	}, nil
 }
 
-// isUnknownFlagError returns true if the error indicates the CLI does not
-// recognize a flag. Pass a specific flag (e.g. "--profile") to check for that
-// flag, or pass "" to match any "unknown flag:" error.
-func isUnknownFlagError(err error, flag string) bool {
-	if flag == "" {
-		return strings.Contains(err.Error(), "unknown flag:")
+// isUnknownFlagError returns true if the error wraps a [cliError] whose stderr
+// indicates the CLI does not recognize one of the given flags.
+func isUnknownFlagError(err error, flags []string) bool {
+	var cliErr *cliError
+	if !errors.As(err, &cliErr) {
+		return false
+	}
+	for _, flag := range flags {
+		if strings.Contains(cliErr.stderr, "unknown flag: "+flag) {
+			return true
+		}
 	}
-	return strings.Contains(err.Error(), "unknown flag: "+flag)
+	return false
 }
 
 // parseExpiry parses an expiry time string in multiple formats for cross-SDK compatibility.