From 4ed1b4ee473c1f3dde2127d9ab99d028b83a851e Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sat, 21 Feb 2026 19:20:50 +0000 Subject: [PATCH 1/6] Implement CLICredentials to read tokens from the local cache The SDK's `u2m` credentials strategy (auth type "databricks-cli") shells out to `databricks auth token` as a subprocess to obtain OAuth tokens. When the CLI itself is the process, this creates a circular dependency: the CLI spawns a child copy of itself just to read the token cache. Implement `CLICredentials.Configure()` to read OAuth tokens directly from the local token cache via `u2m.PersistentAuth`, bypassing the subprocess entirely. The `init()` function registers this strategy in the CLI's custom credentials chain, which runs on every CLI invocation regardless of the command being executed. Co-Authored-By: Claude Opus 4.6 --- go.mod | 2 +- go.sum | 4 +- libs/auth/credentials.go | 94 +++++++++++++++ libs/auth/credentials_test.go | 211 ++++++++++++++++++++++++++++++++++ 4 files changed, 308 insertions(+), 3 deletions(-) create mode 100644 libs/auth/credentials.go create mode 100644 libs/auth/credentials_test.go diff --git a/go.mod b/go.mod index c71055fc4f..a923b793f6 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( github.com/charmbracelet/bubbletea v1.3.10 // MIT github.com/charmbracelet/huh v0.8.0 github.com/charmbracelet/lipgloss v1.1.0 // MIT - github.com/databricks/databricks-sdk-go v0.110.0 // Apache 2.0 + github.com/databricks/databricks-sdk-go v0.110.1-0.20260221140112-be1d4d821dd1 // Apache 2.0 github.com/fatih/color v1.18.0 // MIT github.com/google/uuid v1.6.0 // BSD-3-Clause github.com/gorilla/mux v1.8.1 // BSD 3-Clause diff --git a/go.sum b/go.sum index 1dae205541..556512706d 100644 --- a/go.sum +++ b/go.sum @@ -75,8 +75,8 @@ github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s= github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI= -github.com/databricks/databricks-sdk-go v0.110.0 h1:npiIFyXkRfRrgLBoUVwP9ZgePmjtPuwVQmMt3Bd72M0= -github.com/databricks/databricks-sdk-go v0.110.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= +github.com/databricks/databricks-sdk-go v0.110.1-0.20260221140112-be1d4d821dd1 h1:Rmfg+YEO1ARiEPzQqC2eftDwvXuu/yvWAPhdRuWPk9w= +github.com/databricks/databricks-sdk-go v0.110.1-0.20260221140112-be1d4d821dd1/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go new file mode 100644 index 0000000000..47a5aac30c --- /dev/null +++ b/libs/auth/credentials.go @@ -0,0 +1,94 @@ +package auth + +import ( + "context" + "errors" + + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/config/credentials" + sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" + "github.com/databricks/databricks-sdk-go/credentials/u2m" +) + +func init() { + // Sets the credentials chain for the CLI. + // + // The CLI relies on its own credentials chain to authenticate the user. + // This guarantees that the CLI remain stable despite the evolution of + // the SDK while allowing the customization of some strategies such as + // "databricks-cli" which has a different behavior than the SDK. + config.DefaultCredentialStrategyProvider = func() config.CredentialsStrategy { + // Order in which strategies are tested. Iteration proceeds from most + // specific to most generic, and the first strategy to return a non-nil + // credentials provider is selected. + // + // Modifying this order could break authentication for users whose + // environments are compatible with multiple strategies and who rely + // on the current priority for tie-breaking. + return config.NewCredentialsChain( + config.PatCredentials{}, + config.BasicCredentials{}, + config.M2mCredentials{}, + CLICredentials{}, // custom + config.MetadataServiceCredentials{}, + // OIDC Strategies. + config.GitHubOIDCCredentials{}, + config.AzureDevOpsOIDCCredentials{}, + config.EnvOIDCCredentials{}, + config.FileOIDCCredentials{}, + // Azure strategies. + config.AzureGithubOIDCCredentials{}, + config.AzureMsiCredentials{}, + config.AzureClientSecretCredentials{}, + config.AzureCliCredentials{}, + // Google strategies. + config.GoogleCredentials{}, + config.GoogleDefaultCredentials{}, + ) + } +} + +// CLICredentials is a credentials strategy that reads OAuth tokens directly +// from the local token store. It replaces the SDK's default "databricks-cli" +// strategy, which shells out to `databricks auth token` as a subprocess. +type CLICredentials struct { + PersistentAuthOptions []u2m.PersistentAuthOption +} + +// Name implements [config.CredentialsStrategy]. +func (c CLICredentials) Name() string { + return "databricks-cli" +} + +// Configure implements [config.CredentialsStrategy]. +func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (credentials.CredentialsProvider, error) { + if cfg.Host == "" { + return nil, errors.New("no host provided") + } + oauthArg, err := authArgumentsFromConfig(cfg).ToOAuthArgument() + if err != nil { + return nil, err + } + opts := append(c.PersistentAuthOptions, u2m.WithOAuthArgument(oauthArg)) + persistentAuth, err := u2m.NewPersistentAuth(ctx, opts...) + if err != nil { + return nil, err + } + ts := sdkauth.NewCachedTokenSource( + authconv.AuthTokenSource(persistentAuth), + sdkauth.WithAsyncRefresh(!cfg.DisableOAuthRefreshToken), + ) + cp := credentials.NewOAuthCredentialsProviderFromTokenSource(ts) + return cp, nil +} + +// authArgumentsFromConfig converts an SDK config to AuthArguments. +func authArgumentsFromConfig(cfg *config.Config) AuthArguments { + return AuthArguments{ + Host: cfg.Host, + AccountID: cfg.AccountID, + WorkspaceID: cfg.WorkspaceID, + IsUnifiedHost: cfg.Experimental_IsUnifiedHost, + } +} diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go new file mode 100644 index 0000000000..32c4acbda4 --- /dev/null +++ b/libs/auth/credentials_test.go @@ -0,0 +1,211 @@ +package auth_test + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +type inMemoryTokenCache struct { + Tokens map[string]*oauth2.Token +} + +func (c *inMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) { + t, ok := c.Tokens[key] + if !ok { + return nil, cache.ErrNotFound + } + return t, nil +} + +func (c *inMemoryTokenCache) Store(key string, t *oauth2.Token) error { + c.Tokens[key] = t + return nil +} + +var _ cache.TokenCache = (*inMemoryTokenCache)(nil) + +type mockEndpointSupplier struct{} + +func (m *mockEndpointSupplier) GetAccountOAuthEndpoints(_ context.Context, accountHost, _ string) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + TokenEndpoint: accountHost + "/token", + AuthorizationEndpoint: accountHost + "/authorize", + }, nil +} + +func (m *mockEndpointSupplier) GetWorkspaceOAuthEndpoints(_ context.Context, workspaceHost string) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + TokenEndpoint: workspaceHost + "/token", + AuthorizationEndpoint: workspaceHost + "/authorize", + }, nil +} + +func (m *mockEndpointSupplier) GetUnifiedOAuthEndpoints(_ context.Context, host, _ string) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + TokenEndpoint: host + "/token", + AuthorizationEndpoint: host + "/authorize", + }, nil +} + +var _ u2m.OAuthEndpointSupplier = (*mockEndpointSupplier)(nil) + +func TestCLICredentialsName(t *testing.T) { + c := auth.CLICredentials{} + assert.Equal(t, "databricks-cli", c.Name()) +} + +func TestCLICredentialsConfigure(t *testing.T) { + refreshSuccess := fixtures.HTTPFixture{ + MatchAny: true, + Status: 200, + Response: map[string]string{ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": "3600", + }, + } + refreshFailure := fixtures.HTTPFixture{ + MatchAny: true, + Status: 401, + Response: map[string]string{ + "error": "invalid_request", + "error_description": "Refresh token is invalid", + }, + } + + tests := []struct { + name string + cfg *config.Config + cache map[string]*oauth2.Token + http []fixtures.HTTPFixture + wantNil bool + wantErr string + }{ + { + name: "empty host returns error", + cfg: &config.Config{}, + wantErr: "no host provided", + }, + { + name: "workspace host with valid token", + cfg: &config.Config{ + Host: "https://myworkspace.cloud.databricks.com", + }, + cache: map[string]*oauth2.Token{ + "https://myworkspace.cloud.databricks.com": { + AccessToken: "valid-token", + Expiry: time.Now().Add(time.Hour), + }, + }, + }, + { + name: "account host with valid token", + cfg: &config.Config{ + Host: "https://accounts.cloud.databricks.com", + AccountID: "test-account-id", + }, + cache: map[string]*oauth2.Token{ + "https://accounts.cloud.databricks.com/oidc/accounts/test-account-id": { + AccessToken: "valid-account-token", + Expiry: time.Now().Add(time.Hour), + }, + }, + }, + { + name: "no cached token", + cfg: &config.Config{ + Host: "https://myworkspace.cloud.databricks.com", + }, + cache: map[string]*oauth2.Token{}, + }, + { + name: "expired token with successful refresh", + cfg: &config.Config{ + Host: "https://myworkspace.cloud.databricks.com", + }, + cache: map[string]*oauth2.Token{ + "https://myworkspace.cloud.databricks.com": { + RefreshToken: "valid-refresh", + Expiry: time.Now().Add(-time.Hour), + }, + }, + http: []fixtures.HTTPFixture{refreshSuccess}, + }, + { + name: "expired token with failed refresh", + cfg: &config.Config{ + Host: "https://myworkspace.cloud.databricks.com", + }, + cache: map[string]*oauth2.Token{ + "https://myworkspace.cloud.databricks.com": { + RefreshToken: "bad-refresh", + Expiry: time.Now().Add(-time.Hour), + }, + }, + http: []fixtures.HTTPFixture{refreshFailure}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + + // Pre-populate the cache using the real cache keys when a cache + // map is provided. Test cases use hard-coded keys for readability; + // we re-key them here if the config produces a valid OAuthArgument. + tokenCache := &inMemoryTokenCache{Tokens: make(map[string]*oauth2.Token)} + if tt.cache != nil && tt.cfg.Host != "" { + key, err := func() (string, error) { + arg, err := auth.AuthArguments{ + Host: tt.cfg.Host, + AccountID: tt.cfg.AccountID, + }.ToOAuthArgument() + if err != nil { + return "", err + } + return arg.GetCacheKey(), nil + }() + if err == nil { + for _, tok := range tt.cache { + tokenCache.Tokens[key] = tok + break // only one token per test case + } + } + } + + opts := []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tokenCache), + u2m.WithOAuthEndpointSupplier(&mockEndpointSupplier{}), + } + if len(tt.http) > 0 { + transport := fixtures.SliceTransport(tt.http) + opts = append(opts, u2m.WithHttpClient(&http.Client{Transport: transport})) + } + + c := auth.CLICredentials{PersistentAuthOptions: opts} + cp, err := c.Configure(ctx, tt.cfg) + + if tt.wantErr != "" { + assert.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + if tt.wantNil { + assert.Nil(t, cp) + return + } + require.NotNil(t, cp) + }) + } +} From 94e78c4323d0b8c1e106ba3784d7b9a19d7d0a16 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 15:16:26 +0000 Subject: [PATCH 2/6] update --- libs/auth/credentials.go | 95 +++++++----- libs/auth/credentials_test.go | 271 ++++++++++++++++------------------ 2 files changed, 180 insertions(+), 186 deletions(-) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 47a5aac30c..f61e5c4600 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -6,46 +6,49 @@ import ( "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" "github.com/databricks/databricks-sdk-go/credentials/u2m" ) +// The CLI relies on its own credentials chain to authenticate the user. +// This guarantees that the CLI remain stable despite the evolution of +// the SDK while allowing the customization of some strategies such as +// "databricks-cli" which has a different behavior than the SDK. +// +// Order in which strategies are tested. Iteration proceeds from most +// specific to most generic, and the first strategy to return a non-nil +// credentials provider is selected. +// +// Modifying this order could break authentication for users whose +// environments are compatible with multiple strategies and who rely +// on the current priority for tie-breaking. +var credentialChain = []config.CredentialsStrategy{ + config.PatCredentials{}, + config.BasicCredentials{}, + config.M2mCredentials{}, + CLICredentials{}, // custom + config.MetadataServiceCredentials{}, + // OIDC Strategies. + config.GitHubOIDCCredentials{}, + config.AzureDevOpsOIDCCredentials{}, + config.EnvOIDCCredentials{}, + config.FileOIDCCredentials{}, + // Azure strategies. + config.AzureGithubOIDCCredentials{}, + config.AzureMsiCredentials{}, + config.AzureClientSecretCredentials{}, + config.AzureCliCredentials{}, + // Google strategies. + config.GoogleCredentials{}, + config.GoogleDefaultCredentials{}, +} + func init() { // Sets the credentials chain for the CLI. - // - // The CLI relies on its own credentials chain to authenticate the user. - // This guarantees that the CLI remain stable despite the evolution of - // the SDK while allowing the customization of some strategies such as - // "databricks-cli" which has a different behavior than the SDK. config.DefaultCredentialStrategyProvider = func() config.CredentialsStrategy { - // Order in which strategies are tested. Iteration proceeds from most - // specific to most generic, and the first strategy to return a non-nil - // credentials provider is selected. - // - // Modifying this order could break authentication for users whose - // environments are compatible with multiple strategies and who rely - // on the current priority for tie-breaking. - return config.NewCredentialsChain( - config.PatCredentials{}, - config.BasicCredentials{}, - config.M2mCredentials{}, - CLICredentials{}, // custom - config.MetadataServiceCredentials{}, - // OIDC Strategies. - config.GitHubOIDCCredentials{}, - config.AzureDevOpsOIDCCredentials{}, - config.EnvOIDCCredentials{}, - config.FileOIDCCredentials{}, - // Azure strategies. - config.AzureGithubOIDCCredentials{}, - config.AzureMsiCredentials{}, - config.AzureClientSecretCredentials{}, - config.AzureCliCredentials{}, - // Google strategies. - config.GoogleCredentials{}, - config.GoogleDefaultCredentials{}, - ) + return config.NewCredentialsChain(credentialChain...) } } @@ -53,7 +56,9 @@ func init() { // from the local token store. It replaces the SDK's default "databricks-cli" // strategy, which shells out to `databricks auth token` as a subprocess. type CLICredentials struct { - PersistentAuthOptions []u2m.PersistentAuthOption + // persistentAuth is a function to override the default implementation + // of the persistent auth client. It exists for testing purposes only. + persistentAuthFn func(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) } // Name implements [config.CredentialsStrategy]. @@ -61,28 +66,38 @@ func (c CLICredentials) Name() string { return "databricks-cli" } +var errNoHost = errors.New("no host provided") + // Configure implements [config.CredentialsStrategy]. func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (credentials.CredentialsProvider, error) { if cfg.Host == "" { - return nil, errors.New("no host provided") + return nil, errNoHost } oauthArg, err := authArgumentsFromConfig(cfg).ToOAuthArgument() if err != nil { return nil, err } - opts := append(c.PersistentAuthOptions, u2m.WithOAuthArgument(oauthArg)) - persistentAuth, err := u2m.NewPersistentAuth(ctx, opts...) + ts, err := c.persistentAuth(ctx, u2m.WithOAuthArgument(oauthArg)) if err != nil { return nil, err } - ts := sdkauth.NewCachedTokenSource( - authconv.AuthTokenSource(persistentAuth), - sdkauth.WithAsyncRefresh(!cfg.DisableOAuthRefreshToken), + cp := credentials.NewOAuthCredentialsProviderFromTokenSource( + sdkauth.NewCachedTokenSource(ts, sdkauth.WithAsyncRefresh(!cfg.DisableOAuthRefreshToken)), ) - cp := credentials.NewOAuthCredentialsProviderFromTokenSource(ts) return cp, nil } +func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + if c.persistentAuthFn != nil { + return c.persistentAuthFn(ctx, opts...) + } + ts, err := u2m.NewPersistentAuth(ctx, opts...) + if err != nil { + return nil, err + } + return authconv.AuthTokenSource(ts), nil +} + // authArgumentsFromConfig converts an SDK config to AuthArguments. func authArgumentsFromConfig(cfg *config.Config) AuthArguments { return AuthArguments{ diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 32c4acbda4..8114f3611a 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -1,211 +1,190 @@ -package auth_test +package auth import ( "context" + "errors" "net/http" + "slices" "testing" - "time" - "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/credentials/u2m" - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" - "github.com/databricks/databricks-sdk-go/httpclient/fixtures" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) -type inMemoryTokenCache struct { - Tokens map[string]*oauth2.Token -} - -func (c *inMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) { - t, ok := c.Tokens[key] - if !ok { - return nil, cache.ErrNotFound +// TestCredentialChainOrder purely exists as an extra measure to catch +// accidental change in the ordering. +func TestCredentialChainOrder(t *testing.T) { + names := make([]string, len(credentialChain)) + for i, s := range credentialChain { + names[i] = s.Name() + } + want := []string{ + "pat", + "basic", + "oauth-m2m", + "databricks-cli", + "metadata-service", + "github-oidc", + "azure-devops-oidc", + "env-oidc", + "file-oidc", + "github-oidc-azure", + "azure-msi", + "azure-client-secret", + "azure-cli", + "google-credentials", + "google-id", + } + if !slices.Equal(names, want) { + t.Errorf("credential chain order: want %v, got %v", want, names) } - return t, nil -} - -func (c *inMemoryTokenCache) Store(key string, t *oauth2.Token) error { - c.Tokens[key] = t - return nil -} - -var _ cache.TokenCache = (*inMemoryTokenCache)(nil) - -type mockEndpointSupplier struct{} - -func (m *mockEndpointSupplier) GetAccountOAuthEndpoints(_ context.Context, accountHost, _ string) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: accountHost + "/token", - AuthorizationEndpoint: accountHost + "/authorize", - }, nil -} - -func (m *mockEndpointSupplier) GetWorkspaceOAuthEndpoints(_ context.Context, workspaceHost string) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: workspaceHost + "/token", - AuthorizationEndpoint: workspaceHost + "/authorize", - }, nil -} - -func (m *mockEndpointSupplier) GetUnifiedOAuthEndpoints(_ context.Context, host, _ string) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: host + "/token", - AuthorizationEndpoint: host + "/authorize", - }, nil } -var _ u2m.OAuthEndpointSupplier = (*mockEndpointSupplier)(nil) - func TestCLICredentialsName(t *testing.T) { - c := auth.CLICredentials{} - assert.Equal(t, "databricks-cli", c.Name()) -} - -func TestCLICredentialsConfigure(t *testing.T) { - refreshSuccess := fixtures.HTTPFixture{ - MatchAny: true, - Status: 200, - Response: map[string]string{ - "access_token": "refreshed-access-token", - "token_type": "Bearer", - "expires_in": "3600", - }, - } - refreshFailure := fixtures.HTTPFixture{ - MatchAny: true, - Status: 401, - Response: map[string]string{ - "error": "invalid_request", - "error_description": "Refresh token is invalid", - }, + c := CLICredentials{} + if got := c.Name(); got != "databricks-cli" { + t.Errorf("Name(): want %q, got %q", "databricks-cli", got) } +} +func TestAuthArgumentsFromConfig(t *testing.T) { tests := []struct { - name string - cfg *config.Config - cache map[string]*oauth2.Token - http []fixtures.HTTPFixture - wantNil bool - wantErr string + name string + cfg *config.Config + want AuthArguments }{ { - name: "empty host returns error", - cfg: &config.Config{}, - wantErr: "no host provided", + name: "empty config", + cfg: &config.Config{}, + want: AuthArguments{}, }, { - name: "workspace host with valid token", + name: "workspace host only", cfg: &config.Config{ Host: "https://myworkspace.cloud.databricks.com", }, - cache: map[string]*oauth2.Token{ - "https://myworkspace.cloud.databricks.com": { - AccessToken: "valid-token", - Expiry: time.Now().Add(time.Hour), - }, + want: AuthArguments{ + Host: "https://myworkspace.cloud.databricks.com", }, }, { - name: "account host with valid token", + name: "account host with account ID", cfg: &config.Config{ Host: "https://accounts.cloud.databricks.com", AccountID: "test-account-id", }, - cache: map[string]*oauth2.Token{ - "https://accounts.cloud.databricks.com/oidc/accounts/test-account-id": { - AccessToken: "valid-account-token", - Expiry: time.Now().Add(time.Hour), - }, + want: AuthArguments{ + Host: "https://accounts.cloud.databricks.com", + AccountID: "test-account-id", }, }, { - name: "no cached token", + name: "all fields", cfg: &config.Config{ - Host: "https://myworkspace.cloud.databricks.com", + Host: "https://myhost.com", + AccountID: "acc-123", + WorkspaceID: "ws-456", + Experimental_IsUnifiedHost: true, + }, + want: AuthArguments{ + Host: "https://myhost.com", + AccountID: "acc-123", + WorkspaceID: "ws-456", + IsUnifiedHost: true, }, - cache: map[string]*oauth2.Token{}, }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := authArgumentsFromConfig(tt.cfg) + if got != tt.want { + t.Errorf("want %v, got %v", tt.want, got) + } + }) + } +} + +func TestCLICredentialsConfigure(t *testing.T) { + testErr := errors.New("test error") + + tests := []struct { + name string + cfg *config.Config + persistentAuthFn func(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) + wantErr error + wantToken string + }{ { - name: "expired token with successful refresh", + name: "empty host returns error", + cfg: &config.Config{}, + wantErr: errNoHost, + }, + { + name: "persistentAuthFn error is propagated", cfg: &config.Config{ Host: "https://myworkspace.cloud.databricks.com", }, - cache: map[string]*oauth2.Token{ - "https://myworkspace.cloud.databricks.com": { - RefreshToken: "valid-refresh", - Expiry: time.Now().Add(-time.Hour), - }, + persistentAuthFn: func(_ context.Context, _ ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + return nil, testErr }, - http: []fixtures.HTTPFixture{refreshSuccess}, + wantErr: testErr, }, { - name: "expired token with failed refresh", + name: "workspace host", cfg: &config.Config{ Host: "https://myworkspace.cloud.databricks.com", }, - cache: map[string]*oauth2.Token{ - "https://myworkspace.cloud.databricks.com": { - RefreshToken: "bad-refresh", - Expiry: time.Now().Add(-time.Hour), - }, + persistentAuthFn: func(_ context.Context, _ ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "workspace-token"}, nil + }), nil + }, + wantToken: "workspace-token", + }, + { + name: "account host", + cfg: &config.Config{ + Host: "https://accounts.cloud.databricks.com", + AccountID: "test-account-id", + }, + persistentAuthFn: func(_ context.Context, _ ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "account-token"}, nil + }), nil }, - http: []fixtures.HTTPFixture{refreshFailure}, + wantToken: "account-token", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { ctx := context.Background() + c := CLICredentials{persistentAuthFn: tt.persistentAuthFn} - // Pre-populate the cache using the real cache keys when a cache - // map is provided. Test cases use hard-coded keys for readability; - // we re-key them here if the config produces a valid OAuthArgument. - tokenCache := &inMemoryTokenCache{Tokens: make(map[string]*oauth2.Token)} - if tt.cache != nil && tt.cfg.Host != "" { - key, err := func() (string, error) { - arg, err := auth.AuthArguments{ - Host: tt.cfg.Host, - AccountID: tt.cfg.AccountID, - }.ToOAuthArgument() - if err != nil { - return "", err - } - return arg.GetCacheKey(), nil - }() - if err == nil { - for _, tok := range tt.cache { - tokenCache.Tokens[key] = tok - break // only one token per test case - } - } - } + got, err := c.Configure(ctx, tt.cfg) - opts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(tokenCache), - u2m.WithOAuthEndpointSupplier(&mockEndpointSupplier{}), + if !errors.Is(err, tt.wantErr) { + t.Fatalf("want error %v, got %v", tt.wantErr, err) } - if len(tt.http) > 0 { - transport := fixtures.SliceTransport(tt.http) - opts = append(opts, u2m.WithHttpClient(&http.Client{Transport: transport})) + if tt.wantErr != nil { + return } - c := auth.CLICredentials{PersistentAuthOptions: opts} - cp, err := c.Configure(ctx, tt.cfg) - - if tt.wantErr != "" { - assert.ErrorContains(t, err, tt.wantErr) - return + // Verify the credentials provider sets the correct Bearer token. + req, err := http.NewRequest("GET", tt.cfg.Host, nil) + if err != nil { + t.Fatalf("creating request: %v", err) } - require.NoError(t, err) - if tt.wantNil { - assert.Nil(t, cp) - return + if err := got.SetHeaders(req); err != nil { + t.Fatalf("SetHeaders: want no error, got %v", err) + } + want := "Bearer " + tt.wantToken + if gotHeader := req.Header.Get("Authorization"); gotHeader != want { + t.Errorf("Authorization header: want %q, got %q", want, gotHeader) } - require.NotNil(t, cp) }) } } From 5ab55dc91c4406f4f79d8808c6c6e9125108c0d4 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 15:37:28 +0000 Subject: [PATCH 3/6] Remove redundant sdkauth import alias Co-Authored-By: Claude Opus 4.6 --- libs/auth/credentials.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index f61e5c4600..2ec839bc2a 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -7,7 +7,6 @@ import ( "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/credentials" "github.com/databricks/databricks-sdk-go/config/experimental/auth" - sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" "github.com/databricks/databricks-sdk-go/credentials/u2m" ) @@ -82,7 +81,7 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return nil, err } cp := credentials.NewOAuthCredentialsProviderFromTokenSource( - sdkauth.NewCachedTokenSource(ts, sdkauth.WithAsyncRefresh(!cfg.DisableOAuthRefreshToken)), + auth.NewCachedTokenSource(ts, auth.WithAsyncRefresh(!cfg.DisableOAuthRefreshToken)), ) return cp, nil } From 5ba8a3c6b1c3593037918dfa0e4d35b063c19b00 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 17:48:05 +0000 Subject: [PATCH 4/6] Wrap credential chain to preserve "default" name in error messages The SDK's credentialsChain has an empty name when no strategy succeeds, producing " auth: cannot configure..." instead of "default auth: ...". Wrap the chain like the SDK's DefaultCredentials to provide the fallback. Co-Authored-By: Claude Opus 4.6 --- libs/auth/credentials.go | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 2ec839bc2a..9aa825bda9 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -47,10 +47,27 @@ var credentialChain = []config.CredentialsStrategy{ func init() { // Sets the credentials chain for the CLI. config.DefaultCredentialStrategyProvider = func() config.CredentialsStrategy { - return config.NewCredentialsChain(credentialChain...) + return &defaultCredentials{chain: config.NewCredentialsChain(credentialChain...)} } } +// defaultCredentials wraps the CLI credential chain and provides "default" +// as the fallback name, matching the SDK's DefaultCredentials behavior. +type defaultCredentials struct { + chain config.CredentialsStrategy +} + +func (d *defaultCredentials) Name() string { + if name := d.chain.Name(); name != "" { + return name + } + return "default" +} + +func (d *defaultCredentials) Configure(ctx context.Context, cfg *config.Config) (credentials.CredentialsProvider, error) { + return d.chain.Configure(ctx, cfg) +} + // CLICredentials is a credentials strategy that reads OAuth tokens directly // from the local token store. It replaces the SDK's default "databricks-cli" // strategy, which shells out to `databricks auth token` as a subprocess. From 91123a483f3e725230619d7249f1edf5f9c0ad8a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 18:45:04 +0000 Subject: [PATCH 5/6] Clean up comments in credentials.go Co-Authored-By: Claude Opus 4.6 --- libs/auth/credentials.go | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 9aa825bda9..5b1e2d9d7a 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -11,14 +11,11 @@ import ( "github.com/databricks/databricks-sdk-go/credentials/u2m" ) -// The CLI relies on its own credentials chain to authenticate the user. -// This guarantees that the CLI remain stable despite the evolution of -// the SDK while allowing the customization of some strategies such as -// "databricks-cli" which has a different behavior than the SDK. -// -// Order in which strategies are tested. Iteration proceeds from most -// specific to most generic, and the first strategy to return a non-nil -// credentials provider is selected. +// The credentials chain used by the CLI. It is a custom implementation +// that differs from the SDK's default credentials chain. This guarantees +// that the CLI remain stable despite the evolution of the SDK while +// allowing the customization of some strategies such as "databricks-cli" +// which has a different behavior than the SDK. // // Modifying this order could break authentication for users whose // environments are compatible with multiple strategies and who rely @@ -103,6 +100,9 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return cp, nil } +// persistentAuth returns a token source. It is a convenience function that +// overrides the default implementation of the persistent auth client if +// an alternative implementation is provided for testing. func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { if c.persistentAuthFn != nil { return c.persistentAuthFn(ctx, opts...) From e515fdc4285fba75dcbc043ea0408262ab6e8380 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 24 Feb 2026 22:54:30 +0000 Subject: [PATCH 6/6] Add Profile to AuthArguments and document scope limitation Co-Authored-By: Claude Opus 4.6 --- libs/auth/credentials.go | 8 ++++++++ libs/auth/credentials_test.go | 2 ++ 2 files changed, 10 insertions(+) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 5b1e2d9d7a..6adf0ab9c0 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -82,6 +82,13 @@ func (c CLICredentials) Name() string { var errNoHost = errors.New("no host provided") // Configure implements [config.CredentialsStrategy]. +// +// IMPORTANT: This credentials strategy ignores the scopes specified in the +// config and purely relies on the scopes from the loaded CLI token. This can +// lead to mismatches if the token was obtained with different scopes than the +// ones configured in the current profile. This is a temporary limitation that +// will be addressed in a future release by adding support for dynamic token +// downscoping. func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (credentials.CredentialsProvider, error) { if cfg.Host == "" { return nil, errNoHost @@ -121,5 +128,6 @@ func authArgumentsFromConfig(cfg *config.Config) AuthArguments { AccountID: cfg.AccountID, WorkspaceID: cfg.WorkspaceID, IsUnifiedHost: cfg.Experimental_IsUnifiedHost, + Profile: cfg.Profile, } } diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 8114f3611a..0d8973f853 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -86,12 +86,14 @@ func TestAuthArgumentsFromConfig(t *testing.T) { Host: "https://myhost.com", AccountID: "acc-123", WorkspaceID: "ws-456", + Profile: "my-profile", Experimental_IsUnifiedHost: true, }, want: AuthArguments{ Host: "https://myhost.com", AccountID: "acc-123", WorkspaceID: "ws-456", + Profile: "my-profile", IsUnifiedHost: true, }, },