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) + } + }) } }