Add --force-refresh flag to CLI token source

Pass --force-refresh to the Databricks CLI auth token command to bypass
the CLI's internal token cache. The SDK manages its own token caching,
so the CLI serving stale tokens from its cache causes unnecessary
refresh failures.

Commands are now tried in order: forceCmd (--profile + --force-refresh),
profileCmd (--profile), hostCmd (--host), falling back progressively
for older CLI versions that don't support newer flags.

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

Author: mihaimitrea-db

NEXT_CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### New Features and Improvements
 
+* Pass `--force-refresh` to Databricks CLI `auth token` command to bypass the CLI's internal token cache.
 * Added `HostMetadataResolver` hook to allow callers to customize host metadata resolution, e.g. with caching ([#1572](https://github.com/databricks/databricks-sdk-go/pull/1572)).
 * Added `NewLimitIterator` to `listing` package for lazy iteration with a cap on output items ([#1555](https://github.com/databricks/databricks-sdk-go/pull/1555)).
 

config/cli_token_source.go

@@ -31,12 +31,19 @@ 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 {
-	// cmd is the primary command to execute (--profile when available, --host otherwise).
-	cmd []string
+	// 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
 
-	// hostCmd is a fallback command using --host, used when the primary --profile
-	// command fails because the CLI is too old to support --profile.
+	// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
+	profileCmd []string
+
+	// hostCmd uses --host as a fallback for CLIs that predate --profile support.
+	// Nil when cfg.Host is empty.
 	hostCmd []string
 }
 
@@ -45,25 +52,25 @@ func NewCliTokenSource(cfg *Config) (*CliTokenSource, error) {
 	if err != nil {
 		return nil, err
 	}
-	profileCmd, hostCmd := buildCliCommands(cliPath, cfg)
-	return &CliTokenSource{cmd: profileCmd, hostCmd: hostCmd}, nil
+	forceCmd, profileCmd, hostCmd := buildCliCommands(cliPath, cfg)
+	return &CliTokenSource{forceCmd: forceCmd, profileCmd: profileCmd, hostCmd: hostCmd}, 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) {
+// 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
 	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)
-		}
-		return primary, nil
+		profileCmd = []string{cliPath, "auth", "token", "--profile", cfg.Profile}
+		forceCmd = append(profileCmd, "--force-refresh")
 	}
-	return buildHostCommand(cliPath, cfg), nil
+	if cfg.Host != "" {
+		hostCmd = buildHostCommand(cliPath, cfg)
+	}
+	return forceCmd, profileCmd, hostCmd
 }
 
 // buildHostCommand constructs the legacy --host based CLI command.
@@ -77,16 +84,37 @@ func buildHostCommand(cliPath string, cfg *Config) []string {
 }
 
 // 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.
+// 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) {
-	tok, err := c.execCliCommand(ctx, c.cmd)
-	if err != nil && c.hostCmd != nil && isUnknownFlagError(err) {
+	if c.forceCmd != nil {
+		tok, err := c.execCliCommand(ctx, c.forceCmd)
+		if err == nil {
+			return tok, nil
+		}
+		if !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.")
+	}
+
+	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 c.execCliCommand(ctx, c.hostCmd)
 	}
-	return tok, err
+
+	return nil, fmt.Errorf("no CLI command configured")
 }
 
 func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) {
@@ -115,10 +143,13 @@ func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oa
 }
 
 // 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")
+// 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:")
+	}
+	return strings.Contains(err.Error(), "unknown flag: "+flag)
 }
 
 // parseExpiry parses an expiry time string in multiple formats for cross-SDK compatibility.

config/cli_token_source_test.go

@@ -136,20 +136,21 @@ func TestBuildCliCommands(t *testing.T) {
 	)
 
 	testCases := []struct {
-		name        string
-		cfg         *Config
-		wantCmd     []string
-		wantHostCmd []string
+		name           string
+		cfg            *Config
+		wantForceCmd   []string
+		wantProfileCmd []string
+		wantHostCmd    []string
 	}{
 		{
-			name:    "workspace host",
-			cfg:     &Config{Host: host},
-			wantCmd: []string{cliPath, "auth", "token", "--host", host},
+			name:        "workspace host only",
+			cfg:         &Config{Host: host},
+			wantHostCmd: []string{cliPath, "auth", "token", "--host", host},
 		},
 		{
-			name:    "account host",
-			cfg:     &Config{Host: accountHost, AccountID: accountID},
-			wantCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
+			name:        "account host only",
+			cfg:         &Config{Host: accountHost, AccountID: accountID},
+			wantHostCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
 		},
 		{
 			name: "former unified host treated as workspace",
@@ -158,26 +159,31 @@ func TestBuildCliCommands(t *testing.T) {
 				AccountID:   accountID,
 				WorkspaceID: workspaceID,
 			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost},
+			wantHostCmd: []string{cliPath, "auth", "token", "--host", unifiedHost},
 		},
 		{
-			name:        "profile uses --profile with --host fallback",
-			cfg:         &Config{Profile: "my-profile", Host: host},
-			wantCmd:     []string{cliPath, "auth", "token", "--profile", "my-profile"},
-			wantHostCmd: []string{cliPath, "auth", "token", "--host", host},
+			name:           "profile with host — all three commands",
+			cfg:            &Config{Profile: "my-profile", Host: host},
+			wantForceCmd:   []string{cliPath, "auth", "token", "--profile", "my-profile", "--force-refresh"},
+			wantProfileCmd: []string{cliPath, "auth", "token", "--profile", "my-profile"},
+			wantHostCmd:    []string{cliPath, "auth", "token", "--host", host},
 		},
 		{
-			name:    "profile without host — no fallback",
-			cfg:     &Config{Profile: "my-profile"},
-			wantCmd: []string{cliPath, "auth", "token", "--profile", "my-profile"},
+			name:           "profile without host — no host fallback",
+			cfg:            &Config{Profile: "my-profile"},
+			wantForceCmd:   []string{cliPath, "auth", "token", "--profile", "my-profile", "--force-refresh"},
+			wantProfileCmd: []string{cliPath, "auth", "token", "--profile", "my-profile"},
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			gotCmd, gotHostCmd := buildCliCommands(cliPath, tc.cfg)
-			if !slices.Equal(gotCmd, tc.wantCmd) {
-				t.Errorf("primary cmd = %v, want %v", gotCmd, tc.wantCmd)
+			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)
 			}
 			if !slices.Equal(gotHostCmd, tc.wantHostCmd) {
 				t.Errorf("host cmd = %v, want %v", gotHostCmd, tc.wantHostCmd)
@@ -204,9 +210,8 @@ func TestNewCliTokenSource(t *testing.T) {
 		if err != nil {
 			t.Fatalf("NewCliTokenSource() unexpected error: %v", err)
 		}
-		// Verify CLI path was resolved and used
-		if ts.cmd[0] != validCliPath {
-			t.Errorf("cmd[0] = %q, want %q", ts.cmd[0], validCliPath)
+		if ts.hostCmd[0] != validCliPath {
+			t.Errorf("hostCmd[0] = %q, want %q", ts.hostCmd[0], validCliPath)
 		}
 	})
 
@@ -262,7 +267,7 @@ func TestCliTokenSource_Token(t *testing.T) {
 				t.Fatalf("failed to create mock script: %v", err)
 			}
 
-			ts := &CliTokenSource{cmd: []string{mockScript}}
+			ts := &CliTokenSource{hostCmd: []string{mockScript}}
 			token, err := ts.Token(context.Background())
 
 			if tc.wantErrMsg != "" {
@@ -281,42 +286,122 @@ func TestCliTokenSource_Token(t *testing.T) {
 	}
 }
 
-func TestCliTokenSource_Token_FallbackOnUnknownFlag(t *testing.T) {
+func TestCliTokenSource_Token_ForceRefreshFallbackToProfile(t *testing.T) {
+	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: "profile-token",
+		TokenType:   "Bearer",
+		Expiry:      expiry,
+	})
+
+	tempDir := t.TempDir()
+
+	forceScript := filepath.Join(tempDir, "force_cli.sh")
+	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)
+	}
+
+	profileScript := filepath.Join(tempDir, "profile_cli.sh")
+	if err := os.WriteFile(profileScript, []byte("#!/bin/sh\necho '"+string(validResponse)+"'"), 0755); err != nil {
+		t.Fatalf("failed to create profile script: %v", err)
+	}
+
+	ts := &CliTokenSource{
+		forceCmd:   []string{forceScript},
+		profileCmd: []string{profileScript},
+	}
+	token, err := ts.Token(context.Background())
+	if err != nil {
+		t.Fatalf("Token() error = %v, want fallback to profileCmd to succeed", err)
+	}
+	if token.AccessToken != "profile-token" {
+		t.Errorf("AccessToken = %q, want %q", token.AccessToken, "profile-token")
+	}
+}
+
+func TestCliTokenSource_Token_ProfileFallbackToHost(t *testing.T) {
 	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: "fallback-token",
+		AccessToken: "host-token",
 		TokenType:   "Bearer",
 		Expiry:      expiry,
 	})
 
 	tempDir := t.TempDir()
 
-	// Primary script simulates an old CLI that doesn't know --profile.
 	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)
 	}
 
-	// Fallback script succeeds with --host.
 	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{
-		cmd:     []string{profileScript},
-		hostCmd: []string{hostScript},
+		profileCmd: []string{profileScript},
+		hostCmd:    []string{hostScript},
 	}
 	token, err := ts.Token(context.Background())
 	if err != nil {
-		t.Fatalf("Token() error = %v, want fallback to succeed", err)
+		t.Fatalf("Token() error = %v, want fallback to hostCmd to succeed", err)
 	}
-	if token.AccessToken != "fallback-token" {
-		t.Errorf("AccessToken = %q, want %q", token.AccessToken, "fallback-token")
+	if token.AccessToken != "host-token" {
+		t.Errorf("AccessToken = %q, want %q", token.AccessToken, "host-token")
+	}
+}
+
+func TestCliTokenSource_Token_ForceRefreshFallbackToHostOnProfileError(t *testing.T) {
+	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",
+		TokenType:   "Bearer",
+		Expiry:      expiry,
+	})
+
+	tempDir := t.TempDir()
+
+	// forceCmd fails with --profile unknown (very old CLI).
+	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 {
+		t.Fatalf("failed to create force script: %v", err)
+	}
+
+	// profileCmd also fails with --profile unknown.
+	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},
+		hostCmd:    []string{hostScript},
+	}
+	token, err := ts.Token(context.Background())
+	if err != nil {
+		t.Fatalf("Token() error = %v, want fallback through to hostCmd to succeed", err)
+	}
+	if token.AccessToken != "host-token" {
+		t.Errorf("AccessToken = %q, want %q", token.AccessToken, "host-token")
 	}
 }
 
@@ -327,21 +412,27 @@ func TestCliTokenSource_Token_NoFallbackOnRealError(t *testing.T) {
 
 	tempDir := t.TempDir()
 
-	// Primary script fails with a real auth error (not unknown flag).
+	// forceCmd fails with a real auth error (not unknown flag).
+	forceScript := filepath.Join(tempDir, "force_cli.sh")
+	if err := os.WriteFile(forceScript, []byte("#!/bin/sh\necho 'cache: databricks OAuth is not configured for this host' >&2\nexit 1"), 0755); err != nil {
+		t.Fatalf("failed to create force script: %v", err)
+	}
+
+	// profileCmd and hostCmd should not be called.
 	profileScript := filepath.Join(tempDir, "profile_cli.sh")
-	if err := os.WriteFile(profileScript, []byte("#!/bin/sh\necho 'cache: databricks OAuth is not configured for this host' >&2\nexit 1"), 0755); err != nil {
+	if err := os.WriteFile(profileScript, []byte("#!/bin/sh\necho 'should not reach here' >&2\nexit 1"), 0755); err != nil {
 		t.Fatalf("failed to create profile script: %v", err)
 	}
 
-	// Fallback script would succeed, but should not be called.
 	hostScript := filepath.Join(tempDir, "host_cli.sh")
 	if err := os.WriteFile(hostScript, []byte("#!/bin/sh\necho 'should not reach here' >&2\nexit 1"), 0755); err != nil {
 		t.Fatalf("failed to create host script: %v", err)
 	}
 
 	ts := &CliTokenSource{
-		cmd:     []string{profileScript},
-		hostCmd: []string{hostScript},
+		forceCmd:   []string{forceScript},
+		profileCmd: []string{profileScript},
+		hostCmd:    []string{hostScript},
 	}
 	_, err := ts.Token(context.Background())
 	if err == nil {