From a3140009bed15894c45ba2d45f6769c307173f3f Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Tue, 7 Apr 2026 07:25:35 +0000 Subject: [PATCH] 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 --- NEXT_CHANGELOG.md | 1 + config/cli_token_source.go | 95 +++++++++----- config/cli_token_source_test.go | 212 ++++++++++++++++++++------------ 3 files changed, 198 insertions(+), 110 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index c05710010..36482f112 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -23,6 +23,7 @@ ### Internal Changes + * Pass `--force-refresh` to Databricks CLI `auth token` command to bypass the CLI's internal token cache. * 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)). diff --git a/config/cli_token_source.go b/config/cli_token_source.go index cb502130b..2b0780af3 100644 --- a/config/cli_token_source.go +++ b/config/cli_token_source.go @@ -31,12 +31,22 @@ 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 - - // hostCmd is a fallback command using --host, used when the primary --profile - // command fails because the CLI is too old to support --profile. + // 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 } @@ -45,25 +55,29 @@ 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 +// (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 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} + } + 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 buildHostCommand(cliPath, cfg), nil + return forceCmd, profileCmd, hostCmd } // buildHostCommand constructs the legacy --host based CLI command. @@ -77,16 +91,36 @@ 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, "--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) + 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.") - return c.execCliCommand(ctx, c.hostCmd) } - return tok, err + + if c.hostCmd == nil { + return nil, fmt.Errorf("cannot get access token: no CLI commands available") + } + return c.execCliCommand(ctx, c.hostCmd) } func (c *CliTokenSource) execCliCommand(ctx context.Context, args []string) (*oauth2.Token, error) { @@ -95,7 +129,7 @@ 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))) + return nil, fmt.Errorf("cannot get access token: %q", strings.TrimSpace(string(exitErr.Stderr))) } return nil, fmt.Errorf("cannot get access token: %w", err) } @@ -115,10 +149,9 @@ 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 the given flag (e.g. "--profile", "--force-refresh"). +func isUnknownFlagError(err error, flag string) bool { + return strings.Contains(err.Error(), "unknown flag: "+flag) } // parseExpiry parses an expiry time string in multiple formats for cross-SDK compatibility. diff --git a/config/cli_token_source_test.go b/config/cli_token_source_test.go index 153a3dffb..99db794ea 100644 --- a/config/cli_token_source_test.go +++ b/config/cli_token_source_test.go @@ -136,20 +136,23 @@ 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 — force-refresh based on host", + cfg: &Config{Host: host}, + wantForceCmd: []string{cliPath, "auth", "token", "--host", host, "--force-refresh"}, + 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 — force-refresh based on host with account-id", + cfg: &Config{Host: accountHost, AccountID: accountID}, + wantForceCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID, "--force-refresh"}, + wantHostCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID}, }, { name: "former unified host treated as workspace", @@ -158,26 +161,32 @@ func TestBuildCliCommands(t *testing.T) { AccountID: accountID, WorkspaceID: workspaceID, }, - wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost}, + wantForceCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--force-refresh"}, + 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 — force-refresh based on profile", + 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 +213,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 +270,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,73 +289,119 @@ func TestCliTokenSource_Token(t *testing.T) { } } -func TestCliTokenSource_Token_FallbackOnUnknownFlag(t *testing.T) { +func TestCliTokenSource_Token_Fallback(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", + successResponse, _ := json.Marshal(cliTokenResponse{ + AccessToken: "test-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) - } + success := "#!/bin/sh\necho '" + string(successResponse) + "'" + unknownForceRefresh := "#!/bin/sh\necho 'Error: unknown flag: --force-refresh' >&2\nexit 1" + unknownProfile := "#!/bin/sh\necho 'Error: unknown flag: --profile' >&2\nexit 1" + realError := "#!/bin/sh\necho 'cache: databricks OAuth is not configured for this host' >&2\nexit 1" + unreachable := "#!/bin/sh\necho 'should not reach here' >&2\nexit 1" - ts := &CliTokenSource{ - cmd: []string{profileScript}, - hostCmd: []string{hostScript}, - } - token, err := ts.Token(context.Background()) - if err != nil { - t.Fatalf("Token() error = %v, want fallback to succeed", err) - } - if token.AccessToken != "fallback-token" { - t.Errorf("AccessToken = %q, want %q", token.AccessToken, "fallback-token") - } -} - -func TestCliTokenSource_Token_NoFallbackOnRealError(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("Skipping shell script test on Windows") + testCases := []struct { + name string + forceScript string // empty means nil forceCmd + profileScript string // empty means nil profileCmd + hostScript string // empty means nil hostCmd + wantToken string + wantErrMsg string + }{ + { + name: "force-refresh succeeds", + forceScript: success, + profileScript: unreachable, + hostScript: unreachable, + wantToken: "test-token", + }, + { + name: "force-refresh falls back to profile", + forceScript: unknownForceRefresh, + profileScript: success, + wantToken: "test-token", + }, + { + name: "force-refresh falls back to host (no profile)", + forceScript: unknownForceRefresh, + hostScript: success, + wantToken: "test-token", + }, + { + name: "profile falls back to host", + profileScript: unknownProfile, + hostScript: success, + wantToken: "test-token", + }, + { + name: "full fallback chain: force -> profile -> host", + forceScript: unknownProfile, + profileScript: unknownProfile, + hostScript: success, + wantToken: "test-token", + }, + { + name: "real error stops fallback", + forceScript: realError, + profileScript: unreachable, + hostScript: unreachable, + wantErrMsg: "databricks OAuth is not configured", + }, + { + name: "nil hostCmd after profile failure returns error", + forceScript: unknownProfile, + profileScript: unknownProfile, + wantErrMsg: "no CLI commands available", + }, } - tempDir := t.TempDir() - - // Primary script fails with a real auth error (not unknown flag). - 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 { - t.Fatalf("failed to create profile script: %v", err) - } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tempDir := t.TempDir() + var ts CliTokenSource - // 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) - } + if tc.forceScript != "" { + path := filepath.Join(tempDir, "force_cli.sh") + if err := os.WriteFile(path, []byte(tc.forceScript), 0755); err != nil { + t.Fatalf("failed to create force script: %v", err) + } + ts.forceCmd = []string{path} + } + if tc.profileScript != "" { + path := filepath.Join(tempDir, "profile_cli.sh") + if err := os.WriteFile(path, []byte(tc.profileScript), 0755); err != nil { + t.Fatalf("failed to create profile script: %v", err) + } + ts.profileCmd = []string{path} + } + if tc.hostScript != "" { + path := filepath.Join(tempDir, "host_cli.sh") + if err := os.WriteFile(path, []byte(tc.hostScript), 0755); err != nil { + t.Fatalf("failed to create host script: %v", err) + } + ts.hostCmd = []string{path} + } - ts := &CliTokenSource{ - cmd: []string{profileScript}, - hostCmd: []string{hostScript}, - } - _, err := ts.Token(context.Background()) - if err == nil { - t.Fatal("Token() error = nil, want error") - } - if !strings.Contains(err.Error(), "databricks OAuth is not configured") { - t.Errorf("Token() error = %v, want error containing original auth failure", err) + token, err := ts.Token(context.Background()) + if tc.wantErrMsg != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErrMsg) { + t.Errorf("Token() error = %v, want error containing %q", err, tc.wantErrMsg) + } + return + } + if err != nil { + t.Fatalf("Token() error = %v", err) + } + if token.AccessToken != tc.wantToken { + t.Errorf("AccessToken = %q, want %q", token.AccessToken, tc.wantToken) + } + }) } }