Fix CLI token source --profile fallback with version detection

[mihaimitrea-db]

Summary

Fix the broken --profile fallback in CliTokenSource by replacing error-based detection with version-based CLI detection at init time.

Why

The --profile flag on databricks auth token is a global Cobra flag (defined as a persistent flag on the root command). Old CLIs (< v0.207.1) silently accept it — they don't report "unknown flag: --profile" but instead fail later with "cannot fetch credentials". This means the existing isUnknownFlagError check (config/cli_token_source.go:120) never matches, and the --host fallback is dead code.

This was verified by testing against CLI v0.207.0 vs v0.207.1:

• v0.207.0: databricks auth token --profile workspace → Error: init: cannot fetch credentials (not "unknown flag")
• v0.207.1: databricks auth token --profile workspace → returns a valid token

Approaches considered

Three approaches were evaluated for detecting whether the installed CLI supports --profile:

• Error-based detection (try-and-retry) — the current approach on main. Run databricks auth token --profile <name> and check whether the error contains "unknown flag: --profile". This is broken: because --profile is a global Cobra flag, old CLIs accept it silently and fail with a different error ("cannot fetch credentials"), so the fallback to --host never triggers.

• --help flag parsing (databricks auth token --help + substring matching) was rejected because the --help output format is not a stable API. More importantly, --profile would appear in --help output even on old CLIs that don't actually implement profile-based token lookup — it shows up because it's a global persistent flag, not because the auth token subcommand uses it. This approach has the same fundamental flaw as error-based detection.

• Version detection (databricks version + semver comparison) — the approach taken here. Run databricks version at init time, parse the semver (e.g., "Databricks CLI v0.207.1"), and compare against known minimum versions for each flag. This is reliable because the version string is a stable output format, and the mapping between flags and CLI versions is well-defined (databricks/cli#855 for --profile in v0.207.1). If version detection fails, the SDK falls back to the most conservative command (--host only).

References

• --profile support added in CLI v0.207.1: databricks/cli#855 (Oct 2023)

What changed

Interface changes

None. CliTokenSource is not part of the public API surface.

NewCliTokenSource now takes context.Context as its first parameter, needed for exec.CommandContext when running databricks version. This is consistent with every CredentialsStrategy.Configure method in the codebase, and the single caller (auth_u2m.go) already has ctx in scope.

Behavioral changes

• When cfg.Profile is set but the CLI is too old (< v0.207.1), the SDK now correctly falls back to --host. Previously this fallback was dead code.
• A warning is logged when the --profile flag is not supported and the SDK falls back to --host.

Internal changes

• cliVersion type: semver parsing with AtLeast() comparison and String() formatting.
• getCliVersion(ctx, cliPath): runs databricks version and parses the output.
• parseCliVersion: parses "Databricks CLI v0.207.1" → cliVersion{0, 207, 1}.
• resolveCliCommand: bridges version detection and command building. Falls back to zero version on detection failure.
• buildCliCommand: pure function — takes a version, returns a single resolved command. No exec calls, easy to test.
• CliTokenSource simplified to a single cmd []string field. No hostCmd, no runtime fallback.
• Token() is now one line: return c.execCliCommand(ctx, c.cmd).
• Removed: isUnknownFlagError, buildCliCommands (plural), buildHostCommand, hostCmd field.

How is this tested?

Manual tests on versions 0.207.0 and 0.207.1

Unit tests in config/cli_token_source_test.go:

• TestParseCliVersion — standard versions, patch versions, malformed output, empty string, missing prefix.
• TestCliVersion_AtLeast — equal, higher/lower patch/minor/major, zero vs zero, zero vs nonzero.
• TestBuildCliCommand — table-driven: version x config → expected command. Covers: host-only, account host, profile+new CLI (uses --profile), profile+old CLI (falls back to --host), profile-only+old CLI (nil), zero version (detection failed, falls back to --host), neither profile nor host (nil).
• TestNewCliTokenSource — success with host, success with profile, CLI not found, neither profile nor host.
• TestCliTokenSource_Token — success, CLI error, invalid JSON.

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (4cb4c26 -> 82d9599)

NEXT_CHANGELOG.md

@@ -6,5 +6,5 @@
  
 + * 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)).
-  * 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)).
-  * Bump golang.org/x/net from 0.23.0 to 0.33.0 in /examples/slog ([#1127](https://github.com/databricks/databricks-sdk-go/pull/1127)).
\ No newline at end of file
+  * 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)).
\ No newline at end of file

config/cli_token_source.go

@@ -18,13 +18,13 @@
 +	flag           string
 +	warningMessage string
 +}
-+
+ 
+-	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
+-	profileCmd []string
 +var cliFeatureFlags = []cliFeatureFlag{
 +	{"--force-refresh", "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."},
 +}
- 
--	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
--	profileCmd []string
++
 +const profileFlagWarning = "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version."
  
 -	// hostCmd uses --host as a fallback for CLIs that predate --profile support.
@@ -83,7 +83,7 @@
 +				warning = cliFeatureFlags[i].warningMessage
 +			}
 +			commands = append(commands, cliCommand{args: args, warningMessage: warning})
-+			}
++		}
  	}
 -	if cfg.Host != "" {
 -		hostCmd = buildHostCommand(cliPath, cfg)

_{Reproduce locally: git range-diff f626fed..4cb4c26 67f79d8..82d9599 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (82d9599 -> 68d45f4)

config/cli_token_source.go

@@ -18,13 +18,13 @@
 +	flag           string
 +	warningMessage string
 +}
- 
--	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
--	profileCmd []string
++
 +var cliFeatureFlags = []cliFeatureFlag{
 +	{"--force-refresh", "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."},
 +}
-+
+ 
+-	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
+-	profileCmd []string
 +const profileFlagWarning = "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version."
  
 -	// hostCmd uses --host as a fallback for CLIs that predate --profile support.
@@ -143,10 +143,8 @@
 -		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 c.execCliCommand(ctx, c.hostCmd)
--	}
--
- 	return nil, fmt.Errorf("no CLI command configured")
+-	return c.execCliCommand(ctx, c.hostCmd)
++	return nil, fmt.Errorf("cannot get access token: no CLI commands configured")
  }
- 
\ No newline at end of file
+ 
+ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) {
\ No newline at end of file

_{Reproduce locally: git range-diff 67f79d8..82d9599 db4df21..68d45f4 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (68d45f4 -> 4a5079f)

config/cli_token_source.go

@@ -116,6 +116,17 @@
  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 c.profileCmd != nil {
+-		tok, err := c.execCliCommand(ctx, c.profileCmd)
 +	for i := c.activeCommandIndex; i < len(c.commands); i++ {
 +		cmd := c.commands[i]
 +		tok, err := c.execCliCommand(ctx, cmd.args)
@@ -123,26 +134,18 @@
 +			c.activeCommandIndex = i
  			return tok, nil
  		}
--		if !isUnknownFlagError(err, "") {
+-		if !isUnknownFlagError(err, "--profile") {
 +		lastCommand := i == len(c.commands)-1
 +		if lastCommand || !isUnknownFlagError(err, "") {
  			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.")
-+		logger.Warnf(ctx, cmd.warningMessage)
- 	}
--
--	if c.profileCmd != nil {
--		tok, err := c.execCliCommand(ctx, c.profileCmd)
--		if err == nil {
--			return tok, nil
--		}
--		if !isUnknownFlagError(err, "--profile") {
--			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")
  }

config/cli_token_source_test.go

@@ -76,13 +76,13 @@
 -			gotForceCmd, gotProfileCmd, gotHostCmd := buildCliCommands(cliPath, tc.cfg)
 -			if !slices.Equal(gotForceCmd, tc.wantForceCmd) {
 -				t.Errorf("force cmd = %v, want %v", gotForceCmd, tc.wantForceCmd)
--			}
--			if !slices.Equal(gotProfileCmd, tc.wantProfileCmd) {
--				t.Errorf("profile cmd = %v, want %v", gotProfileCmd, tc.wantProfileCmd)
 +			got := buildCliCommands(cliPath, tc.cfg)
 +			if len(got) != len(tc.wantCommands) {
 +				t.Fatalf("got %d commands, want %d", len(got), len(tc.wantCommands))
  			}
+-			if !slices.Equal(gotProfileCmd, tc.wantProfileCmd) {
+-				t.Errorf("profile cmd = %v, want %v", gotProfileCmd, tc.wantProfileCmd)
+-			}
 -			if !slices.Equal(gotHostCmd, tc.wantHostCmd) {
 -				t.Errorf("host cmd = %v, want %v", gotHostCmd, tc.wantHostCmd)
 +			for i, want := range tc.wantCommands {
@@ -208,15 +208,15 @@
  	_, err := ts.Token(context.Background())
  	if err == nil {
  		t.Fatal("Token() error = nil, want error")
- 		t.Errorf("Token() error = %v, want error containing original auth failure", err)
  	}
  }
-+
+ 
+-func TestCliTokenSource_Token_NilHostCmdReturnsError(t *testing.T) {
 +func TestCliTokenSource_Token_ActiveCommandIndexPersists(t *testing.T) {
-+	if runtime.GOOS == "windows" {
-+		t.Skip("Skipping shell script test on Windows")
-+	}
-+
+ 	if runtime.GOOS == "windows" {
+ 		t.Skip("Skipping shell script test on Windows")
+ 	}
+ 
 +	expiry := time.Now().Add(1 * time.Hour).Format(time.RFC3339)
 +	validResponse, _ := json.Marshal(cliTokenResponse{
 +		AccessToken: "host-token",
@@ -224,18 +224,25 @@
 +		Expiry:      expiry,
 +	})
 +
-+	tempDir := t.TempDir()
-+
-+	forceScript := filepath.Join(tempDir, "force_cli.sh")
+ 	tempDir := t.TempDir()
+ 
+ 	forceScript := filepath.Join(tempDir, "force_cli.sh")
+-	if err := os.WriteFile(forceScript, []byte("#!/bin/sh\necho 'Error: unknown flag: --profile' >&2\nexit 1"), 0755); err != nil {
 +	if err := os.WriteFile(forceScript, []byte("#!/bin/sh\necho 'Error: unknown flag: --force-refresh' >&2\nexit 1"), 0755); err != nil {
-+		t.Fatalf("failed to create force script: %v", err)
-+	}
-+
+ 		t.Fatalf("failed to create force script: %v", err)
+ 	}
+ 
+-	profileScript := filepath.Join(tempDir, "profile_cli.sh")
+-	if err := os.WriteFile(profileScript, []byte("#!/bin/sh\necho 'Error: unknown flag: --profile' >&2\nexit 1"), 0755); err != nil {
+-		t.Fatalf("failed to create profile script: %v", err)
 +	hostScript := filepath.Join(tempDir, "host_cli.sh")
 +	if err := os.WriteFile(hostScript, []byte("#!/bin/sh\necho '"+string(validResponse)+"'"), 0755); err != nil {
 +		t.Fatalf("failed to create host script: %v", err)
-+	}
-+
+ 	}
+ 
+-	ts := &CliTokenSource{
+-		forceCmd:   []string{forceScript},
+-		profileCmd: []string{profileScript},
 +	ts := &CliTokenSource{commands: []cliCommand{
 +		{args: []string{forceScript}, warningMessage: "force-refresh not supported"},
 +		{args: []string{hostScript}},
@@ -245,13 +252,18 @@
 +	token, err := ts.Token(context.Background())
 +	if err != nil {
 +		t.Fatalf("first Token() error = %v", err)
-+	}
+ 	}
+-	_, err := ts.Token(context.Background())
+-	if err == nil {
+-		t.Fatal("Token() error = nil, want error")
 +	if token.AccessToken != "host-token" {
 +		t.Errorf("first AccessToken = %q, want %q", token.AccessToken, "host-token")
 +	}
 +	if ts.activeCommandIndex != 1 {
 +		t.Errorf("activeCommandIndex = %d, want 1", ts.activeCommandIndex)
-+	}
+ 	}
+-	if !strings.Contains(err.Error(), "no CLI commands available") {
+-		t.Errorf("Token() error = %v, want error containing %q", err, "no CLI commands available")
 +
 +	// Second call: starts at activeCommandIndex, skipping the force command.
 +	token, err = ts.Token(context.Background())
@@ -260,5 +272,5 @@
 +	}
 +	if token.AccessToken != "host-token" {
 +		t.Errorf("second AccessToken = %q, want %q", token.AccessToken, "host-token")
-+	}
-+}
\ No newline at end of file
+ 	}
+ }
\ No newline at end of file

_{Reproduce locally: git range-diff db4df21..68d45f4 97a1007..4a5079f | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (4a5079f -> 6f4fead)

config/cli_token_source.go

@@ -1,6 +1,13 @@
 diff --git a/config/cli_token_source.go b/config/cli_token_source.go
 --- a/config/cli_token_source.go
 +++ b/config/cli_token_source.go
+ 	"path/filepath"
+ 	"runtime"
+ 	"strings"
++	"sync/atomic"
+ 	"time"
+ 
+ 	"github.com/databricks/databricks-sdk-go/logger"
  	Expiry      string `json:"expiry"`
  }
  
@@ -42,7 +49,7 @@
 +// falling back progressively for older CLI versions that lack newer flags.
 +type CliTokenSource struct {
 +	commands           []cliCommand
-+	activeCommandIndex int
++	activeCommandIndex atomic.Int32
  }
  
  func NewCliTokenSource(cfg *Config) (*CliTokenSource, error) {
@@ -127,11 +134,11 @@
 -
 -	if c.profileCmd != nil {
 -		tok, err := c.execCliCommand(ctx, c.profileCmd)
-+	for i := c.activeCommandIndex; i < len(c.commands); i++ {
++	for i := int(c.activeCommandIndex.Load()); i < len(c.commands); i++ {
 +		cmd := c.commands[i]
 +		tok, err := c.execCliCommand(ctx, cmd.args)
  		if err == nil {
-+			c.activeCommandIndex = i
++			c.activeCommandIndex.Store(int32(i))
  			return tok, nil
  		}
 -		if !isUnknownFlagError(err, "--profile") {

config/cli_token_source_test.go

@@ -259,8 +259,8 @@
 +	if token.AccessToken != "host-token" {
 +		t.Errorf("first AccessToken = %q, want %q", token.AccessToken, "host-token")
 +	}
-+	if ts.activeCommandIndex != 1 {
-+		t.Errorf("activeCommandIndex = %d, want 1", ts.activeCommandIndex)
++	if ts.activeCommandIndex.Load() != 1 {
++		t.Errorf("activeCommandIndex = %d, want 1", ts.activeCommandIndex.Load())
  	}
 -	if !strings.Contains(err.Error(), "no CLI commands available") {
 -		t.Errorf("Token() error = %v, want error containing %q", err, "no CLI commands available")

_{Reproduce locally: git range-diff 97a1007..4a5079f 69a7c95..6f4fead | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (6f4fead -> 2218270)

config/cli_token_source.go

@@ -11,43 +11,29 @@
  	Expiry      string `json:"expiry"`
  }
  
--// CliTokenSource fetches OAuth tokens by shelling out to the Databricks CLI.
++// 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
++	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.
--type CliTokenSource struct {
++// 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 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
-+// cliFeatureFlag defines a CLI feature flag and the warning to log when
-+// falling back because the CLI does not support it. Ordered newest-first:
-+// commands are built by progressively stripping these flags for older CLIs.
-+type cliFeatureFlag struct {
-+	flag           string
-+	warningMessage string
-+}
-+
-+var cliFeatureFlags = []cliFeatureFlag{
-+	{"--force-refresh", "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."},
-+}
- 
+-
 -	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
 -	profileCmd []string
-+const profileFlagWarning = "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version."
- 
+-
 -	// hostCmd uses --host as a fallback for CLIs that predate --profile support.
 -	// Nil when cfg.Host is empty.
 -	hostCmd []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
-+	warningMessage 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 atomic.Int32
  }
@@ -69,48 +55,47 @@
 -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. When cfg.Profile is set,
-+// commands include feature flags from [cliFeatureFlags] (stripped progressively)
-+// plus a --host fallback when host is available. When cfg.Profile is empty,
-+// only --host is returned — the CLI must support --profile before any feature
-+// flags can be used (monotonic feature assumption).
++// 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")
-+		baseArgs := []string{cliPath, "auth", "token", "--profile", cfg.Profile}
-+		for i := 0; i <= len(cliFeatureFlags); i++ {
-+			args := append([]string{}, baseArgs...)
-+			for _, f := range cliFeatureFlags[i:] {
-+				args = append(args, f.flag)
-+			}
-+			warning := profileFlagWarning
-+			if i < len(cliFeatureFlags) {
-+				warning = cliFeatureFlags[i].warningMessage
-+			}
-+			commands = append(commands, cliCommand{args: args, warningMessage: warning})
-+		}
- 	}
+-	}
 -	if cfg.Host != "" {
 -		hostCmd = buildHostCommand(cliPath, cfg)
-+	hostArgs := buildHostCommand(cliPath, cfg)
-+	if hostArgs != nil {
-+		commands = append(commands, cliCommand{args: hostArgs})
- 	}
+-	}
 -	return forceCmd, profileCmd, hostCmd
++		commands = append(commands, cliCommand{
++			args:           []string{cliPath, "auth", "token", "--profile", cfg.Profile, "--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:           []string{cliPath, "auth", "token", "--profile", cfg.Profile},
++			warningMessage: "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.",
++		})
++	}
++	commands = appendHostCommand(commands, cliPath, cfg)
 +	return commands
  }
  
- // buildHostCommand constructs the legacy --host based CLI command.
-+// Returns nil when cfg.Host is empty.
- func buildHostCommand(cliPath string, cfg *Config) []string {
+-// buildHostCommand constructs the legacy --host based CLI command.
+-func buildHostCommand(cliPath string, cfg *Config) []string {
+-	cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
++func appendHostCommand(commands []cliCommand, cliPath string, cfg *Config) []cliCommand {
 +	if cfg.Host == "" {
-+		return nil
++		return commands
 +	}
- 	cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
++	args := []string{cliPath, "auth", "token", "--host", cfg.Host}
  	switch cfg.HostType() {
  	case AccountHost:
+-		cmd = append(cmd, "--account-id", cfg.AccountID)
++		args = append(args, "--account-id", cfg.AccountID)
+ 	}
+-	return cmd
++	return append(commands, cliCommand{args: args})
  }
  
  // Token fetches an OAuth token by shelling out to the Databricks CLI.

_{Reproduce locally: git range-diff 69a7c95..6f4fead 69a7c95..2218270 | Disable: git config gitstack.push-range-diff false}

[mihaimitrea-db]

Range-diff: stack/force-refresh-flag (2218270 -> 76e74ca)

config/cli_token_source.go

@@ -15,6 +15,7 @@
 +// that is logged when this command fails and we fall back to the next one.
 +type cliCommand struct {
 +	args           []string
++	flags          []string
 +	warningMessage string
 +}
 +
@@ -70,10 +71,12 @@
 -	return forceCmd, profileCmd, hostCmd
 +		commands = append(commands, cliCommand{
 +			args:           []string{cliPath, "auth", "token", "--profile", cfg.Profile, "--force-refresh"},
++			flags:          []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},
++			flags:          []string{"--profile"},
 +			warningMessage: "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.",
 +		})
 +	}
@@ -108,17 +111,6 @@
  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 c.profileCmd != nil {
--		tok, err := c.execCliCommand(ctx, c.profileCmd)
 +	for i := int(c.activeCommandIndex.Load()); i < len(c.commands); i++ {
 +		cmd := c.commands[i]
 +		tok, err := c.execCliCommand(ctx, cmd.args)
@@ -126,20 +118,52 @@
 +			c.activeCommandIndex.Store(int32(i))
  			return tok, nil
  		}
--		if !isUnknownFlagError(err, "--profile") {
+-		if !isUnknownFlagError(err, "--force-refresh") && !isUnknownFlagError(err, "--profile") {
 +		lastCommand := i == len(c.commands)-1
-+		if lastCommand || !isUnknownFlagError(err, "") {
++		if lastCommand || !isUnknownFlagError(err, cmd.flags) {
  			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.")
++		logger.Warnf(ctx, cmd.warningMessage)
+ 	}
+-
+-	if c.profileCmd != nil {
+-		tok, err := c.execCliCommand(ctx, c.profileCmd)
+-		if err == nil {
+-			return tok, nil
+-		}
+-		if !isUnknownFlagError(err, "--profile") {
+-			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) {
\ No newline at end of file
+ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) {
+ }
+ 
+ // 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:")
++// 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 strings.Contains(err.Error(), "unknown flag: "+flag)
++	return false
+ }
+ 
+ // parseExpiry parses an expiry time string in multiple formats for cross-SDK compatibility.
\ No newline at end of file

config/cli_token_source_test.go

@@ -121,7 +121,7 @@
 -		profileCmd: []string{profileScript},
 -	}
 +	ts := &CliTokenSource{commands: []cliCommand{
-+		{args: []string{forceScript}, warningMessage: "force-refresh not supported"},
++		{args: []string{forceScript}, flags: []string{"--force-refresh"}, warningMessage: "force-refresh not supported"},
 +		{args: []string{profileScript}},
 +	}}
  	token, err := ts.Token(context.Background())
@@ -135,7 +135,7 @@
 -		hostCmd:    []string{hostScript},
 -	}
 +	ts := &CliTokenSource{commands: []cliCommand{
-+		{args: []string{profileScript}, warningMessage: "profile not supported"},
++		{args: []string{profileScript}, flags: []string{"--profile"}, warningMessage: "profile not supported"},
 +		{args: []string{hostScript}},
 +	}}
  	token, err := ts.Token(context.Background())
@@ -171,8 +171,8 @@
 -		hostCmd:    []string{hostScript},
 -	}
 +	ts := &CliTokenSource{commands: []cliCommand{
-+		{args: []string{forceScript}, warningMessage: "force-refresh not supported"},
-+		{args: []string{profileScript}, warningMessage: "profile not supported"},
++		{args: []string{forceScript}, flags: []string{"--force-refresh", "--profile"}, warningMessage: "force-refresh not supported"},
++		{args: []string{profileScript}, flags: []string{"--profile"}, warningMessage: "profile not supported"},
 +		{args: []string{hostScript}},
 +	}}
  	token, err := ts.Token(context.Background())
@@ -201,8 +201,8 @@
 -		hostCmd:    []string{hostScript},
 -	}
 +	ts := &CliTokenSource{commands: []cliCommand{
-+		{args: []string{forceScript}, warningMessage: "force-refresh not supported"},
-+		{args: []string{profileScript}, warningMessage: "profile not supported"},
++		{args: []string{forceScript}, flags: []string{"--force-refresh", "--profile"}, warningMessage: "force-refresh not supported"},
++		{args: []string{profileScript}, flags: []string{"--profile"}, warningMessage: "profile not supported"},
 +		{args: []string{hostScript}},
 +	}}
  	_, err := ts.Token(context.Background())
@@ -244,7 +244,7 @@
 -		forceCmd:   []string{forceScript},
 -		profileCmd: []string{profileScript},
 +	ts := &CliTokenSource{commands: []cliCommand{
-+		{args: []string{forceScript}, warningMessage: "force-refresh not supported"},
++		{args: []string{forceScript}, flags: []string{"--force-refresh"}, warningMessage: "force-refresh not supported"},
 +		{args: []string{hostScript}},
 +	}}
 +

_{Reproduce locally: git range-diff 69a7c95..2218270 69a7c95..76e74ca | Disable: git config gitstack.push-range-diff false}

[renaudhartert-db]

Could we use golang.org/x/mod/semver instead? It's already part of the go.mod.

[renaudhartert-db]

Suggested change

	// --profile is a global Cobra flag — old CLIs accept it silently but
	// fail with "cannot fetch credentials" instead of "unknown flag".
	// We use version detection to decide --profile vs --host.
	// Flag --profile is a global CLI flag and is recognized for all commands even
	// the ones that do not support it. Only use flag --profile in CLI versions that
	// are known to support it in command `auth token`.

[renaudhartert-db]

Suggested change

	logger.Warnf(ctx, "Profile %q was specified but Databricks CLI %s does not support --profile (requires >= %s). Falling back to --host.", cfg.Profile, ver, cliVersionForProfile)
	logger.Warnf(ctx, "Databricks CLI %s does not support --profile (requires >= %s). Falling back to --host.", ver, cliVersionForProfile)

[renaudhartert-db]

This is equivalent but less bugprone in case cmd is initialized as an empty slice.

Suggested change

	if cmd == nil && cfg.Host != "" {
	if len(cmd) == 0 && cfg.Host != "" {

[renaudhartert-db]

[optional] I believe a more readable and idiomatic way to write this code would be to organize it so that (i) the "normal branch" (i.e. use --profile) is on the lowest level of indentation, and (ii) each branch is terminal. Something like this:

func buildCliCommand(ctx context.Context, cliPath string, cfg *Config, ver cliVersion) []string {
	if cfg.Profile == "" {
		return buildCliHostCommand(ctx context.Context, cliPath, cfg)
	}

	if !ver.AtLeast(cliVersionForProfile) {
		logger.Warnf(ctx, "CLI version XXX does not support --profile")
		return buildCliHostCommand(ctx context.Context, cliPath, cfg)
	}

	cmd := []string{cliPath, "auth", "token", "--profile", cfg.Profile}
	if cfg.HostType() == AccountHost {
		cmd = append(cmd, "--account-id", cfg.AccountID)
	}
	return cmd
}

func buildCliHostCommand(ctx context.Context, cliPath string, cfg *Config, ver cliVersion) []string {
	cmd = []string{cliPath, "auth", "token", "--host", cfg.Host}
	if cfg.HostType() == AccountHost {
		cmd = append(cmd, "--account-id", cfg.AccountID)
	}
	return cmd
}

It is a little longer but much easier to follow. Fallin back to using --host is now clearly the exception path and one does not need to understand how the host command is built to understand that.

[renaudhartert-db]

Why aren't they useful? Naively, I'd imagine that one might want to access this information for debugging what is happening with the CLI.

I could understand that we discard it because we do not want the error to be part of the SDK API contract but then I wonder why we are not discarding it at line 155 too.

[renaudhartert-db]

Could we have this as table tests and merge them with the one above so that we have one comprehensive test suite for NewCliTokenSource?

[github-actions]

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-go

Inputs:

• PR number: 1605
• Commit SHA: eb40138c964569f236225858d16b02bd517548b1

Checks will be approved automatically on success.