From f5d474cbbac5b1c2e9bdd51a98bd9d0192cf1412 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 10 Mar 2026 15:02:14 +0100 Subject: [PATCH 01/11] Add token introspection client for login.databricks.com discovery Add IntrospectToken() which calls /api/2.0/tokens/introspect to extract account_id and workspace_id from workspace tokens. This will be used by the discovery login flow to enrich profiles with metadata. Note: go.mod contains a temporary replace directive pointing to the local SDK worktree with discovery changes. This will be reverted before the PR. --- libs/auth/introspect.go | 62 +++++++++++++++++++++ libs/auth/introspect_test.go | 105 +++++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 libs/auth/introspect.go create mode 100644 libs/auth/introspect_test.go diff --git a/libs/auth/introspect.go b/libs/auth/introspect.go new file mode 100644 index 0000000000..b31d16991a --- /dev/null +++ b/libs/auth/introspect.go @@ -0,0 +1,62 @@ +package auth + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" +) + +// IntrospectionResponse represents the response from the Databricks token +// introspection endpoint at /api/2.0/tokens/introspect. +type IntrospectionResponse struct { + PrincipalContext struct { + AuthenticationScope struct { + AccountID string `json:"account_id"` + WorkspaceID int64 `json:"workspace_id"` + } `json:"authentication_scope"` + } `json:"principal_context"` +} + +// IntrospectionResult contains the extracted metadata from token introspection. +type IntrospectionResult struct { + AccountID string + WorkspaceID string +} + +// IntrospectToken calls the workspace token introspection endpoint to extract +// account_id and workspace_id for the given access token. Returns an error +// if the request fails or the response cannot be parsed. Callers should treat +// errors as non-fatal (best-effort metadata enrichment). +func IntrospectToken(ctx context.Context, host string, accessToken string) (*IntrospectionResult, error) { + endpoint := strings.TrimSuffix(host, "/") + "/api/2.0/tokens/introspect" + req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil) + if err != nil { + return nil, fmt.Errorf("creating introspection request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+accessToken) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, fmt.Errorf("calling introspection endpoint: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("introspection endpoint returned status %d", resp.StatusCode) + } + + var introspection IntrospectionResponse + if err := json.NewDecoder(resp.Body).Decode(&introspection); err != nil { + return nil, fmt.Errorf("decoding introspection response: %w", err) + } + + result := &IntrospectionResult{ + AccountID: introspection.PrincipalContext.AuthenticationScope.AccountID, + } + if introspection.PrincipalContext.AuthenticationScope.WorkspaceID != 0 { + result.WorkspaceID = fmt.Sprintf("%d", introspection.PrincipalContext.AuthenticationScope.WorkspaceID) + } + return result, nil +} diff --git a/libs/auth/introspect_test.go b/libs/auth/introspect_test.go new file mode 100644 index 0000000000..21306795a2 --- /dev/null +++ b/libs/auth/introspect_test.go @@ -0,0 +1,105 @@ +package auth + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIntrospectToken_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "principal_context": { + "authentication_scope": { + "account_id": "a]b1c234-5678-90ab-cdef-1234567890ab", + "workspace_id": 2548836972759138 + } + } + }`)) + })) + defer server.Close() + + result, err := IntrospectToken(t.Context(), server.URL, "test-token") + require.NoError(t, err) + assert.Equal(t, "a]b1c234-5678-90ab-cdef-1234567890ab", result.AccountID) + assert.Equal(t, "2548836972759138", result.WorkspaceID) +} + +func TestIntrospectToken_ZeroWorkspaceID(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "principal_context": { + "authentication_scope": { + "account_id": "abc-123", + "workspace_id": 0 + } + } + }`)) + })) + defer server.Close() + + result, err := IntrospectToken(t.Context(), server.URL, "test-token") + require.NoError(t, err) + assert.Equal(t, "abc-123", result.AccountID) + assert.Empty(t, result.WorkspaceID) +} + +func TestIntrospectToken_HTTPError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer server.Close() + + _, err := IntrospectToken(t.Context(), server.URL, "test-token") + assert.ErrorContains(t, err, "status 403") +} + +func TestIntrospectToken_ServerError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer server.Close() + + _, err := IntrospectToken(t.Context(), server.URL, "test-token") + assert.ErrorContains(t, err, "status 500") +} + +func TestIntrospectToken_MalformedJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`not json`)) + })) + defer server.Close() + + _, err := IntrospectToken(t.Context(), server.URL, "test-token") + assert.ErrorContains(t, err, "decoding introspection response") +} + +func TestIntrospectToken_VerifyAuthHeader(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "Bearer my-secret-token", r.Header.Get("Authorization")) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"principal_context":{"authentication_scope":{"account_id":"x","workspace_id":1}}}`)) + })) + defer server.Close() + + _, err := IntrospectToken(t.Context(), server.URL, "my-secret-token") + require.NoError(t, err) +} + +func TestIntrospectToken_VerifyEndpoint(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/api/2.0/tokens/introspect", r.URL.Path) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"principal_context":{"authentication_scope":{"account_id":"x","workspace_id":1}}}`)) + })) + defer server.Close() + + _, err := IntrospectToken(t.Context(), server.URL, "token") + require.NoError(t, err) +} From f951fef712a13c12f37a56bf7cdabdaa897ee8b2 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 10 Mar 2026 15:28:56 +0100 Subject: [PATCH 02/11] Fix discovery login tests to avoid hanging on PersistentAuth.Challenge Extract shouldUseDiscovery() routing function and test it directly with table-driven tests instead of executing the full cobra command, which hangs waiting for an OAuth browser callback that never arrives in tests. Co-Authored-By: Claude Opus 4.6 (1M context) --- cmd/auth/login.go | 104 +++++++++++++++++++++++++++++++++++++++-- cmd/auth/login_test.go | 100 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 201 insertions(+), 3 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index deab15d286..149acb46de 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/exec" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" @@ -70,9 +71,11 @@ you can refer to the documentation linked below. GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html -This command requires a Databricks Host URL (using --host or as a positional argument -or implicitly inferred from the specified profile name) -and a profile name (using --profile) to be specified. If you don't specify these +If no host is provided (via --host, as a positional argument, or from an existing +profile), the CLI will open login.databricks.com where you can authenticate and +select a workspace. The workspace URL will be discovered automatically. + +A profile name (using --profile) can be specified. If you don't specify these values, you'll be prompted for values at runtime. While this command always logs you into the specified host, the runtime behaviour @@ -139,6 +142,12 @@ depends on the existing profiles you have set in your configuration file return err } + // If no host is available from any source, use the discovery flow + // via login.databricks.com. + if shouldUseDiscovery(authArguments.Host, args, existingProfile) { + return discoveryLogin(ctx, profileName, loginTimeout, scopes, getBrowserFunc(cmd)) + } + // Load unified host flags from the profile if not explicitly set via CLI flag if !cmd.Flag("experimental-is-unified-host").Changed && existingProfile != nil { authArguments.IsUnifiedHost = existingProfile.IsUnifiedHost @@ -399,6 +408,21 @@ func loadProfileByName(ctx context.Context, profileName string, profiler profile return nil, nil } +// shouldUseDiscovery returns true if the discovery flow should be used +// (no host available from any source). +func shouldUseDiscovery(hostFlag string, args []string, existingProfile *profile.Profile) bool { + if hostFlag != "" { + return false + } + if len(args) > 0 { + return false + } + if existingProfile != nil && existingProfile.Host != "" { + return false + } + return true +} + // openURLSuppressingStderr opens a URL in the browser while suppressing stderr output. // This prevents xdg-open error messages from being displayed to the user. func openURLSuppressingStderr(url string) error { @@ -415,6 +439,80 @@ func openURLSuppressingStderr(url string) error { return browserpkg.OpenURL(url) } +// discoveryLogin runs the login.databricks.com discovery flow. The user +// authenticates in the browser, selects a workspace, and the CLI receives +// the workspace host from the OAuth callback's iss parameter. +func discoveryLogin(ctx context.Context, profileName string, timeout time.Duration, scopes string, browserFunc func(string) error) error { + arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName) + if err != nil { + return err + } + + var scopesList []string + if scopes != "" { + for _, s := range strings.Split(scopes, ",") { + scopesList = append(scopesList, strings.TrimSpace(s)) + } + } + + opts := []u2m.PersistentAuthOption{ + u2m.WithOAuthArgument(arg), + u2m.WithBrowser(browserFunc), + u2m.WithDiscoveryLogin(), + } + if len(scopesList) > 0 { + opts = append(opts, u2m.WithScopes(scopesList)) + } + + // Apply timeout before creating PersistentAuth so Challenge() respects it. + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + persistentAuth, err := u2m.NewPersistentAuth(ctx, opts...) + if err != nil { + return err + } + defer persistentAuth.Close() + + cmdio.LogString(ctx, "Opening login.databricks.com in your browser...") + if err := persistentAuth.Challenge(); err != nil { + return fmt.Errorf("login via login.databricks.com failed: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) + } + + discoveredHost := arg.GetDiscoveredHost() + + // Get the token for introspection + tok, err := persistentAuth.Token() + if err != nil { + return fmt.Errorf("retrieving token after login: %w", err) + } + + // Best-effort introspection for metadata + var accountID, workspaceID string + introspection, err := auth.IntrospectToken(ctx, discoveredHost, tok.AccessToken) + if err != nil { + log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) + } else { + accountID = introspection.AccountID + workspaceID = introspection.WorkspaceID + } + + err = databrickscfg.SaveToProfile(ctx, &config.Config{ + Profile: profileName, + Host: discoveredHost, + AuthType: authTypeDatabricksCLI, + AccountID: accountID, + WorkspaceID: workspaceID, + ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), + }, oauthLoginClearKeys()...) + if err != nil { + return err + } + + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) + return nil +} + // oauthLoginClearKeys returns profile keys that should be explicitly removed // when performing an OAuth login. Derives auth credential fields dynamically // from the SDK's ConfigAttributes to stay in sync as new auth methods are added. diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 013786c4bf..8a23f8ca9e 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -2,6 +2,12 @@ package auth import ( "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" "testing" "github.com/databricks/cli/libs/auth" @@ -255,3 +261,97 @@ func TestLoadProfileByNameAndClusterID(t *testing.T) { }) } } + +func TestShouldUseDiscovery(t *testing.T) { + tests := []struct { + name string + hostFlag string + args []string + existingProfile *profile.Profile + want bool + }{ + { + name: "no host from any source", + want: true, + }, + { + name: "host from flag", + hostFlag: "https://example.com", + want: false, + }, + { + name: "host from positional arg", + args: []string{"https://example.com"}, + want: false, + }, + { + name: "host from existing profile", + existingProfile: &profile.Profile{Host: "https://example.com"}, + want: false, + }, + { + name: "existing profile without host", + existingProfile: &profile.Profile{Name: "test"}, + want: true, + }, + { + name: "nil profile", + existingProfile: nil, + want: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := shouldUseDiscovery(tt.hostFlag, tt.args, tt.existingProfile) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { + // Create a temporary config file for this test + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + err := os.WriteFile(configPath, []byte(""), 0o600) + require.NoError(t, err) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + // Create a mock introspection server that returns an error + introspectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer introspectServer.Close() + + ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) + + // Test the discoveryLogin function's introspection error handling by + // calling IntrospectToken directly (since discoveryLogin requires a + // real PersistentAuth flow we can't easily mock). + result, err := auth.IntrospectToken(ctx, introspectServer.URL, "test-token") + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "403") +} + +func TestDiscoveryLogin_IntrospectionSuccessExtractsMetadata(t *testing.T) { + introspectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + w.Header().Set("Content-Type", "application/json") + err := json.NewEncoder(w).Encode(map[string]any{ + "principal_context": map[string]any{ + "authentication_scope": map[string]any{ + "account_id": "acc-12345", + "workspace_id": 2548836972759138, + }, + }, + }) + assert.NoError(t, err) + })) + defer introspectServer.Close() + + ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) + result, err := auth.IntrospectToken(ctx, introspectServer.URL, "test-token") + require.NoError(t, err) + assert.Equal(t, "acc-12345", result.AccountID) + assert.Equal(t, fmt.Sprintf("%d", int64(2548836972759138)), result.WorkspaceID) +} From 6b733d689232a27121211e70bc66c25a96b274dd Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 11 Mar 2026 10:01:21 +0100 Subject: [PATCH 03/11] Improve discovery login: testability, scopes persistence, and error messages Extract splitScopes helper to deduplicate scope parsing. Save scopes to profile in discoveryLogin (was previously dropped). Add actionable error messages with --host tip for setup failures and config file path for save failures. Make discoveryLogin testable by introducing package-level vars for PersistentAuth, OAuthArgument, and IntrospectToken. Fix typo in introspect test data. Use http.MethodGet constant and drain response body on error for TCP connection reuse. Co-Authored-By: Claude Opus 4.6 (1M context) --- cmd/auth/login.go | 68 ++++++++++++++------ cmd/auth/login_test.go | 121 ++++++++++++++++++++++++++++++----- libs/auth/introspect.go | 5 +- libs/auth/introspect_test.go | 4 +- 4 files changed, 160 insertions(+), 38 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 149acb46de..ae7a510633 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -24,6 +24,7 @@ import ( "github.com/databricks/databricks-sdk-go/credentials/u2m" browserpkg "github.com/pkg/browser" "github.com/spf13/cobra" + "golang.org/x/oauth2" ) func promptForProfile(ctx context.Context, defaultValue string) (string, error) { @@ -49,6 +50,20 @@ const ( authTypeDatabricksCLI = "databricks-cli" ) +type discoveryPersistentAuth interface { + Challenge() error + Token() (*oauth2.Token, error) + Close() error +} + +var newDiscoveryOAuthArgument = u2m.NewBasicDiscoveryOAuthArgument + +var newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + return u2m.NewPersistentAuth(ctx, opts...) +} + +var introspectToken = auth.IntrospectToken + func newLoginCommand(authArguments *auth.AuthArguments) *cobra.Command { defaultConfigPath := "~/.databrickscfg" if runtime.GOOS == "windows" { @@ -167,15 +182,11 @@ depends on the existing profiles you have set in your configuration file switch { case scopes != "": // Explicit --scopes flag takes precedence. - for _, s := range strings.Split(scopes, ",") { - scopesList = append(scopesList, strings.TrimSpace(s)) - } + scopesList = splitScopes(scopes) case existingProfile != nil && existingProfile.Scopes != "": // Preserve scopes from the existing profile so re-login // uses the same scopes the user previously configured. - for _, s := range strings.Split(existingProfile.Scopes, ",") { - scopesList = append(scopesList, strings.TrimSpace(s)) - } + scopesList = splitScopes(existingProfile.Scopes) } oauthArgument, err := authArguments.ToOAuthArgument() @@ -443,17 +454,12 @@ func openURLSuppressingStderr(url string) error { // authenticates in the browser, selects a workspace, and the CLI receives // the workspace host from the OAuth callback's iss parameter. func discoveryLogin(ctx context.Context, profileName string, timeout time.Duration, scopes string, browserFunc func(string) error) error { - arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName) + arg, err := newDiscoveryOAuthArgument(profileName) if err != nil { - return err + return fmt.Errorf("setting up login.databricks.com: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) } - var scopesList []string - if scopes != "" { - for _, s := range strings.Split(scopes, ",") { - scopesList = append(scopesList, strings.TrimSpace(s)) - } - } + scopesList := splitScopes(scopes) opts := []u2m.PersistentAuthOption{ u2m.WithOAuthArgument(arg), @@ -468,9 +474,9 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - persistentAuth, err := u2m.NewPersistentAuth(ctx, opts...) + persistentAuth, err := newDiscoveryPersistentAuth(ctx, opts...) if err != nil { - return err + return fmt.Errorf("setting up login.databricks.com: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) } defer persistentAuth.Close() @@ -489,7 +495,7 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati // Best-effort introspection for metadata var accountID, workspaceID string - introspection, err := auth.IntrospectToken(ctx, discoveredHost, tok.AccessToken) + introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { @@ -497,22 +503,46 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati workspaceID = introspection.WorkspaceID } + configFile := os.Getenv("DATABRICKS_CONFIG_FILE") err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: discoveredHost, AuthType: authTypeDatabricksCLI, AccountID: accountID, WorkspaceID: workspaceID, - ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), + Scopes: scopesList, + ConfigFile: configFile, }, oauthLoginClearKeys()...) if err != nil { - return err + if configFile != "" { + return fmt.Errorf("saving profile %q to %s: %w", profileName, configFile, err) + } + return fmt.Errorf("saving profile %q: %w", profileName, err) } cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) return nil } +// splitScopes splits a comma-separated scopes string into a trimmed slice. +func splitScopes(scopes string) []string { + if scopes == "" { + return nil + } + var result []string + for _, s := range strings.Split(scopes, ",") { + scope := strings.TrimSpace(s) + if scope == "" { + continue + } + result = append(result, scope) + } + if len(result) == 0 { + return nil + } + return result +} + // oauthLoginClearKeys returns profile keys that should be explicitly removed // when performing an OAuth login. Derives auth credential fields dynamically // from the SDK's ConfigAttributes to stay in sync as new auth methods are added. diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 8a23f8ca9e..006ee94e6d 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -3,19 +3,22 @@ package auth import ( "context" "encoding/json" - "fmt" + "errors" "net/http" "net/http/httptest" "os" "path/filepath" "testing" + "time" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" ) func loadTestProfile(t *testing.T, ctx context.Context, profileName string) *profile.Profile { @@ -25,6 +28,27 @@ func loadTestProfile(t *testing.T, ctx context.Context, profileName string) *pro return profile } +type fakeDiscoveryPersistentAuth struct { + token *oauth2.Token + challengeErr error + tokenErr error +} + +func (f *fakeDiscoveryPersistentAuth) Challenge() error { + return f.challengeErr +} + +func (f *fakeDiscoveryPersistentAuth) Token() (*oauth2.Token, error) { + if f.tokenErr != nil { + return nil, f.tokenErr + } + return f.token, nil +} + +func (f *fakeDiscoveryPersistentAuth) Close() error { + return nil +} + func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg") @@ -308,29 +332,94 @@ func TestShouldUseDiscovery(t *testing.T) { } } +func TestSplitScopes(t *testing.T) { + tests := []struct { + name string + input string + output []string + }{ + { + name: "empty input", + input: "", + output: nil, + }, + { + name: "single scope", + input: "all-apis", + output: []string{"all-apis"}, + }, + { + name: "trims whitespace", + input: " all-apis , sql ", + output: []string{"all-apis", "sql"}, + }, + { + name: "drops empty entries", + input: "all-apis, ,sql,,", + output: []string{"all-apis", "sql"}, + }, + { + name: "only empty entries", + input: " , , ", + output: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.output, splitScopes(tt.input)) + }) + } +} + func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { - // Create a temporary config file for this test tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) require.NoError(t, err) t.Setenv("DATABRICKS_CONFIG_FILE", configPath) - // Create a mock introspection server that returns an error - introspectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusForbidden) - })) - defer introspectServer.Close() + originalNewDiscoveryOAuthArgument := newDiscoveryOAuthArgument + originalNewDiscoveryPersistentAuth := newDiscoveryPersistentAuth + originalIntrospectToken := introspectToken + t.Cleanup(func() { + newDiscoveryOAuthArgument = originalNewDiscoveryOAuthArgument + newDiscoveryPersistentAuth = originalNewDiscoveryPersistentAuth + introspectToken = originalIntrospectToken + }) + + newDiscoveryOAuthArgument = func(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName) + if err != nil { + return nil, err + } + arg.SetDiscoveredHost("https://workspace.example.com") + return arg, nil + } - ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) + newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + return &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, nil + } - // Test the discoveryLogin function's introspection error handling by - // calling IntrospectToken directly (since discoveryLogin requires a - // real PersistentAuth flow we can't easily mock). - result, err := auth.IntrospectToken(ctx, introspectServer.URL, "test-token") - assert.Error(t, err) - assert.Nil(t, result) - assert.Contains(t, err.Error(), "403") + introspectToken = func(ctx context.Context, host string, accessToken string) (*auth.IntrospectionResult, error) { + assert.Equal(t, "https://workspace.example.com", host) + assert.Equal(t, "test-token", accessToken) + return nil, errors.New("introspection failed") + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis, ,sql,", func(string) error { return nil }) + require.NoError(t, err) + + savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) + require.NoError(t, err) + require.NotNil(t, savedProfile) + assert.Equal(t, "https://workspace.example.com", savedProfile.Host) + assert.Equal(t, "all-apis,sql", savedProfile.Scopes) + assert.Empty(t, savedProfile.AccountID) + assert.Empty(t, savedProfile.WorkspaceID) } func TestDiscoveryLogin_IntrospectionSuccessExtractsMetadata(t *testing.T) { @@ -353,5 +442,5 @@ func TestDiscoveryLogin_IntrospectionSuccessExtractsMetadata(t *testing.T) { result, err := auth.IntrospectToken(ctx, introspectServer.URL, "test-token") require.NoError(t, err) assert.Equal(t, "acc-12345", result.AccountID) - assert.Equal(t, fmt.Sprintf("%d", int64(2548836972759138)), result.WorkspaceID) + assert.Equal(t, "2548836972759138", result.WorkspaceID) } diff --git a/libs/auth/introspect.go b/libs/auth/introspect.go index b31d16991a..5caa695967 100644 --- a/libs/auth/introspect.go +++ b/libs/auth/introspect.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "strings" ) @@ -31,7 +32,7 @@ type IntrospectionResult struct { // errors as non-fatal (best-effort metadata enrichment). func IntrospectToken(ctx context.Context, host string, accessToken string) (*IntrospectionResult, error) { endpoint := strings.TrimSuffix(host, "/") + "/api/2.0/tokens/introspect" - req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { return nil, fmt.Errorf("creating introspection request: %w", err) } @@ -44,6 +45,8 @@ func IntrospectToken(ctx context.Context, host string, accessToken string) (*Int defer resp.Body.Close() if resp.StatusCode != http.StatusOK { + // Drain the body so the underlying TCP connection can be reused. + _, _ = io.Copy(io.Discard, resp.Body) return nil, fmt.Errorf("introspection endpoint returned status %d", resp.StatusCode) } diff --git a/libs/auth/introspect_test.go b/libs/auth/introspect_test.go index 21306795a2..9946fefb01 100644 --- a/libs/auth/introspect_test.go +++ b/libs/auth/introspect_test.go @@ -15,7 +15,7 @@ func TestIntrospectToken_Success(t *testing.T) { _, _ = w.Write([]byte(`{ "principal_context": { "authentication_scope": { - "account_id": "a]b1c234-5678-90ab-cdef-1234567890ab", + "account_id": "a1b1c234-5678-90ab-cdef-1234567890ab", "workspace_id": 2548836972759138 } } @@ -25,7 +25,7 @@ func TestIntrospectToken_Success(t *testing.T) { result, err := IntrospectToken(t.Context(), server.URL, "test-token") require.NoError(t, err) - assert.Equal(t, "a]b1c234-5678-90ab-cdef-1234567890ab", result.AccountID) + assert.Equal(t, "a1b1c234-5678-90ab-cdef-1234567890ab", result.AccountID) assert.Equal(t, "2548836972759138", result.WorkspaceID) } From fcf3ed50d0a53f31d043ba836375ff10fbfe8950 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 11 Mar 2026 10:21:19 +0100 Subject: [PATCH 04/11] Fix gofumpt and perfsprint lint issues Group same-type function parameters and use strconv.FormatInt instead of fmt.Sprintf for integer conversion. Co-Authored-By: Claude Opus 4.6 (1M context) --- cmd/auth/login_test.go | 2 +- libs/auth/introspect.go | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 006ee94e6d..63fb3982ca 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -403,7 +403,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { }, nil } - introspectToken = func(ctx context.Context, host string, accessToken string) (*auth.IntrospectionResult, error) { + introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { assert.Equal(t, "https://workspace.example.com", host) assert.Equal(t, "test-token", accessToken) return nil, errors.New("introspection failed") diff --git a/libs/auth/introspect.go b/libs/auth/introspect.go index 5caa695967..bd4deaae84 100644 --- a/libs/auth/introspect.go +++ b/libs/auth/introspect.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" ) @@ -30,7 +31,7 @@ type IntrospectionResult struct { // account_id and workspace_id for the given access token. Returns an error // if the request fails or the response cannot be parsed. Callers should treat // errors as non-fatal (best-effort metadata enrichment). -func IntrospectToken(ctx context.Context, host string, accessToken string) (*IntrospectionResult, error) { +func IntrospectToken(ctx context.Context, host, accessToken string) (*IntrospectionResult, error) { endpoint := strings.TrimSuffix(host, "/") + "/api/2.0/tokens/introspect" req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { @@ -59,7 +60,7 @@ func IntrospectToken(ctx context.Context, host string, accessToken string) (*Int AccountID: introspection.PrincipalContext.AuthenticationScope.AccountID, } if introspection.PrincipalContext.AuthenticationScope.WorkspaceID != 0 { - result.WorkspaceID = fmt.Sprintf("%d", introspection.PrincipalContext.AuthenticationScope.WorkspaceID) + result.WorkspaceID = strconv.FormatInt(introspection.PrincipalContext.AuthenticationScope.WorkspaceID, 10) } return result, nil } From 4cfbe37df5b99403677247d1f0cb1d9b5d9e8f52 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 11 Mar 2026 10:39:11 +0100 Subject: [PATCH 05/11] Extract tip constant, reuse splitScopes in token.go, rename test - Extract repeated discovery fallback tip to discoveryFallbackTip const - Reuse splitScopes in token.go instead of inline split+trim (also fixes missing empty-entry filtering) - Rename TestDiscoveryLogin_IntrospectionSuccessExtractsMetadata to TestIntrospectToken_SuccessExtractsMetadata (it tests IntrospectToken directly, not discoveryLogin) Co-Authored-By: Claude Opus 4.6 (1M context) --- cmd/auth/login.go | 1 + cmd/auth/login_test.go | 2 +- cmd/auth/token.go | 4 +--- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index ae7a510633..8d85cb3c7e 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -48,6 +48,7 @@ const ( minimalDbConnectVersion = "13.1" defaultTimeout = 1 * time.Hour authTypeDatabricksCLI = "databricks-cli" + discoveryFallbackTip = "\n\nTip: you can specify a workspace directly with: databricks auth login --host " ) type discoveryPersistentAuth interface { diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 63fb3982ca..a960da079e 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -422,7 +422,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { assert.Empty(t, savedProfile.WorkspaceID) } -func TestDiscoveryLogin_IntrospectionSuccessExtractsMetadata(t *testing.T) { +func TestIntrospectToken_SuccessExtractsMetadata(t *testing.T) { introspectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) w.Header().Set("Content-Type", "application/json") diff --git a/cmd/auth/token.go b/cmd/auth/token.go index ca8582bd02..bd320c9ea3 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -434,9 +434,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr // uses the same scopes the user previously configured. var scopesList []string if existingProfile != nil && existingProfile.Scopes != "" { - for _, s := range strings.Split(existingProfile.Scopes, ",") { - scopesList = append(scopesList, strings.TrimSpace(s)) - } + scopesList = splitScopes(existingProfile.Scopes) } oauthArgument, err := loginArgs.ToOAuthArgument() From faafbb05e3d520ce66e88071421053cd9e103413 Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 11:30:29 +0100 Subject: [PATCH 06/11] Remove account_id from discovery login profile save The SDKs are not yet ready to handle account_id on workspace profiles as part of the cache key. Saving it now would break existing auth flows. Keeping the introspection call so workspace_id is still extracted. Co-authored-by: Isaac --- cmd/auth/login.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 8d85cb3c7e..9c68d0329c 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -495,12 +495,14 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati } // Best-effort introspection for metadata - var accountID, workspaceID string + var workspaceID string introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { - accountID = introspection.AccountID + // TODO: Save introspection.AccountID once the SDKs are ready to use + // account_id as part of the profile/cache key. Adding it now would break + // existing auth flows that don't expect account_id on workspace profiles. workspaceID = introspection.WorkspaceID } @@ -509,7 +511,6 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati Profile: profileName, Host: discoveredHost, AuthType: authTypeDatabricksCLI, - AccountID: accountID, WorkspaceID: workspaceID, Scopes: scopesList, ConfigFile: configFile, From 3014886025715f4dac044f22d2854488f8a4b092 Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 11:40:02 +0100 Subject: [PATCH 07/11] Address review: validate flags in discovery mode, warn on account_id mismatch Reject --configure-cluster and --configure-serverless when discovery login is used (no host available), since these flags require a known workspace. Warn when the introspected account_id differs from the existing profile's account_id to surface potential misconfigurations early. Also use the discoveryFallbackTip constant instead of hardcoded strings. Co-authored-by: Isaac --- cmd/auth/login.go | 23 ++++++++++++++++++----- cmd/auth/login_test.go | 2 +- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 9c68d0329c..9b498d6388 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -161,7 +161,13 @@ depends on the existing profiles you have set in your configuration file // If no host is available from any source, use the discovery flow // via login.databricks.com. if shouldUseDiscovery(authArguments.Host, args, existingProfile) { - return discoveryLogin(ctx, profileName, loginTimeout, scopes, getBrowserFunc(cmd)) + if configureCluster { + return errors.New("--configure-cluster requires --host to be specified") + } + if configureServerless { + return errors.New("--configure-serverless requires --host to be specified") + } + return discoveryLogin(ctx, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) } // Load unified host flags from the profile if not explicitly set via CLI flag @@ -454,10 +460,10 @@ func openURLSuppressingStderr(url string) error { // discoveryLogin runs the login.databricks.com discovery flow. The user // authenticates in the browser, selects a workspace, and the CLI receives // the workspace host from the OAuth callback's iss parameter. -func discoveryLogin(ctx context.Context, profileName string, timeout time.Duration, scopes string, browserFunc func(string) error) error { +func discoveryLogin(ctx context.Context, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error) error { arg, err := newDiscoveryOAuthArgument(profileName) if err != nil { - return fmt.Errorf("setting up login.databricks.com: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) + return fmt.Errorf("setting up login.databricks.com: %w"+discoveryFallbackTip, err) } scopesList := splitScopes(scopes) @@ -477,13 +483,13 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati persistentAuth, err := newDiscoveryPersistentAuth(ctx, opts...) if err != nil { - return fmt.Errorf("setting up login.databricks.com: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) + return fmt.Errorf("setting up login.databricks.com: %w"+discoveryFallbackTip, err) } defer persistentAuth.Close() cmdio.LogString(ctx, "Opening login.databricks.com in your browser...") if err := persistentAuth.Challenge(); err != nil { - return fmt.Errorf("login via login.databricks.com failed: %w\n\nTip: you can specify a workspace directly with: databricks auth login --host ", err) + return fmt.Errorf("login via login.databricks.com failed: %w"+discoveryFallbackTip, err) } discoveredHost := arg.GetDiscoveredHost() @@ -504,6 +510,13 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati // account_id as part of the profile/cache key. Adding it now would break // existing auth flows that don't expect account_id on workspace profiles. workspaceID = introspection.WorkspaceID + + // Warn if the detected account_id differs from what's already saved in the profile. + if existingProfile != nil && existingProfile.AccountID != "" && introspection.AccountID != "" && + existingProfile.AccountID != introspection.AccountID { + log.Warnf(ctx, "detected account ID %q differs from existing profile account ID %q", + introspection.AccountID, existingProfile.AccountID) + } } configFile := os.Getenv("DATABRICKS_CONFIG_FILE") diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a960da079e..60bd0af020 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -410,7 +410,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis, ,sql,", func(string) error { return nil }) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) From 55fe3f5f9243837e5ae0db8dfd9266eee7025624 Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 11:49:33 +0100 Subject: [PATCH 08/11] Add tests for account_id mismatch warning in discovery login Test that a warning is logged when the introspected account_id differs from the existing profile's account_id, and that no warning is emitted when they match. Co-authored-by: Isaac --- cmd/auth/login_test.go | 141 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 60bd0af020..1057f2c640 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -1,13 +1,16 @@ package auth import ( + "bytes" "context" "encoding/json" "errors" + "log/slog" "net/http" "net/http/httptest" "os" "path/filepath" + "sync" "testing" "time" @@ -15,12 +18,31 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) +// logBuffer is a thread-safe bytes.Buffer for capturing log output in tests. +type logBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (lb *logBuffer) Write(p []byte) (int, error) { + lb.mu.Lock() + defer lb.mu.Unlock() + return lb.buf.Write(p) +} + +func (lb *logBuffer) String() string { + lb.mu.Lock() + defer lb.mu.Unlock() + return lb.buf.String() +} + func loadTestProfile(t *testing.T, ctx context.Context, profileName string) *profile.Profile { profile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler) require.NoError(t, err) @@ -444,3 +466,122 @@ func TestIntrospectToken_SuccessExtractsMetadata(t *testing.T) { assert.Equal(t, "acc-12345", result.AccountID) assert.Equal(t, "2548836972759138", result.WorkspaceID) } + +func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + err := os.WriteFile(configPath, []byte(""), 0o600) + require.NoError(t, err) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + originalNewDiscoveryOAuthArgument := newDiscoveryOAuthArgument + originalNewDiscoveryPersistentAuth := newDiscoveryPersistentAuth + originalIntrospectToken := introspectToken + t.Cleanup(func() { + newDiscoveryOAuthArgument = originalNewDiscoveryOAuthArgument + newDiscoveryPersistentAuth = originalNewDiscoveryPersistentAuth + introspectToken = originalIntrospectToken + }) + + newDiscoveryOAuthArgument = func(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName) + if err != nil { + return nil, err + } + arg.SetDiscoveredHost("https://workspace.example.com") + return arg, nil + } + + newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + return &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, nil + } + + introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + return &auth.IntrospectionResult{ + AccountID: "new-account-id", + WorkspaceID: "12345", + }, nil + } + + // Set up a logger that captures log records to verify the warning. + var logBuf logBuffer + logger := slog.New(slog.NewTextHandler(&logBuf, &slog.HandlerOptions{Level: slog.LevelWarn})) + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + ctx = log.NewContext(ctx, logger) + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + AccountID: "old-account-id", + } + + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + require.NoError(t, err) + + // Verify warning about mismatched account IDs was logged. + assert.Contains(t, logBuf.String(), "new-account-id") + assert.Contains(t, logBuf.String(), "old-account-id") + + // Verify the profile was saved without account_id (not overwritten). + savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) + require.NoError(t, err) + require.NotNil(t, savedProfile) + assert.Equal(t, "https://workspace.example.com", savedProfile.Host) + assert.Equal(t, "12345", savedProfile.WorkspaceID) +} + +func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + err := os.WriteFile(configPath, []byte(""), 0o600) + require.NoError(t, err) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + originalNewDiscoveryOAuthArgument := newDiscoveryOAuthArgument + originalNewDiscoveryPersistentAuth := newDiscoveryPersistentAuth + originalIntrospectToken := introspectToken + t.Cleanup(func() { + newDiscoveryOAuthArgument = originalNewDiscoveryOAuthArgument + newDiscoveryPersistentAuth = originalNewDiscoveryPersistentAuth + introspectToken = originalIntrospectToken + }) + + newDiscoveryOAuthArgument = func(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName) + if err != nil { + return nil, err + } + arg.SetDiscoveredHost("https://workspace.example.com") + return arg, nil + } + + newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + return &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, nil + } + + introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + return &auth.IntrospectionResult{ + AccountID: "same-account-id", + WorkspaceID: "12345", + }, nil + } + + var logBuf logBuffer + logger := slog.New(slog.NewTextHandler(&logBuf, &slog.HandlerOptions{Level: slog.LevelWarn})) + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + ctx = log.NewContext(ctx, logger) + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + AccountID: "same-account-id", + } + + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + require.NoError(t, err) + + // No warning should be logged when account IDs match. + assert.Empty(t, logBuf.String()) +} From 3a392c70454d5de8cc975976069866963e63c95a Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 12:04:52 +0100 Subject: [PATCH 09/11] Remove redundant introspection tests - Remove TestIntrospectToken_ServerError (same code path as _HTTPError) - Merge TestIntrospectToken_VerifyAuthHeader and _VerifyEndpoint into one - Remove TestIntrospectToken_SuccessExtractsMetadata from login_test.go (duplicate of introspect_test.go:TestIntrospectToken_Success) Co-authored-by: Isaac --- cmd/auth/login_test.go | 26 -------------------------- libs/auth/introspect_test.go | 25 ++----------------------- 2 files changed, 2 insertions(+), 49 deletions(-) diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 1057f2c640..62c18974a9 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -3,11 +3,8 @@ package auth import ( "bytes" "context" - "encoding/json" "errors" "log/slog" - "net/http" - "net/http/httptest" "os" "path/filepath" "sync" @@ -444,29 +441,6 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { assert.Empty(t, savedProfile.WorkspaceID) } -func TestIntrospectToken_SuccessExtractsMetadata(t *testing.T) { - introspectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) - w.Header().Set("Content-Type", "application/json") - err := json.NewEncoder(w).Encode(map[string]any{ - "principal_context": map[string]any{ - "authentication_scope": map[string]any{ - "account_id": "acc-12345", - "workspace_id": 2548836972759138, - }, - }, - }) - assert.NoError(t, err) - })) - defer introspectServer.Close() - - ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) - result, err := auth.IntrospectToken(ctx, introspectServer.URL, "test-token") - require.NoError(t, err) - assert.Equal(t, "acc-12345", result.AccountID) - assert.Equal(t, "2548836972759138", result.WorkspaceID) -} - func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") diff --git a/libs/auth/introspect_test.go b/libs/auth/introspect_test.go index 9946fefb01..d6ce2e83e6 100644 --- a/libs/auth/introspect_test.go +++ b/libs/auth/introspect_test.go @@ -59,16 +59,6 @@ func TestIntrospectToken_HTTPError(t *testing.T) { assert.ErrorContains(t, err, "status 403") } -func TestIntrospectToken_ServerError(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusInternalServerError) - })) - defer server.Close() - - _, err := IntrospectToken(t.Context(), server.URL, "test-token") - assert.ErrorContains(t, err, "status 500") -} - func TestIntrospectToken_MalformedJSON(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) @@ -80,8 +70,9 @@ func TestIntrospectToken_MalformedJSON(t *testing.T) { assert.ErrorContains(t, err, "decoding introspection response") } -func TestIntrospectToken_VerifyAuthHeader(t *testing.T) { +func TestIntrospectToken_VerifyRequestDetails(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/api/2.0/tokens/introspect", r.URL.Path) assert.Equal(t, "Bearer my-secret-token", r.Header.Get("Authorization")) w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"principal_context":{"authentication_scope":{"account_id":"x","workspace_id":1}}}`)) @@ -91,15 +82,3 @@ func TestIntrospectToken_VerifyAuthHeader(t *testing.T) { _, err := IntrospectToken(t.Context(), server.URL, "my-secret-token") require.NoError(t, err) } - -func TestIntrospectToken_VerifyEndpoint(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - assert.Equal(t, "/api/2.0/tokens/introspect", r.URL.Path) - w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte(`{"principal_context":{"authentication_scope":{"account_id":"x","workspace_id":1}}}`)) - })) - defer server.Close() - - _, err := IntrospectToken(t.Context(), server.URL, "token") - require.NoError(t, err) -} From b4b7147062527be66a26373f4b827e3a49613380 Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 12:58:58 +0100 Subject: [PATCH 10/11] Use env.Get(ctx, ...) instead of os.Getenv in discoveryLogin Follow the codebase convention of using the context-aware env package for environment variable access. Co-authored-by: Isaac --- cmd/auth/login.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 9b498d6388..175df39734 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -519,7 +519,7 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati } } - configFile := os.Getenv("DATABRICKS_CONFIG_FILE") + configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE") err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: discoveredHost, From 85fa40c7b74e0b5ff9168291d3d7b7f5e9aa4842 Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 12 Mar 2026 13:00:21 +0100 Subject: [PATCH 11/11] Replace remaining os.Getenv with env.Get(ctx, ...) in login command The existing non-discovery login path also used os.Getenv for DATABRICKS_CONFIG_FILE. Replaced with the context-aware env package to follow codebase conventions and enable test overrides. Co-authored-by: Isaac --- cmd/auth/login.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 175df39734..0548ce6247 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "io" - "os" "runtime" "strings" "time" @@ -276,7 +275,7 @@ depends on the existing profiles you have set in your configuration file WorkspaceID: authArguments.WorkspaceID, Experimental_IsUnifiedHost: authArguments.IsUnifiedHost, ClusterID: clusterID, - ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), ServerlessComputeID: serverlessComputeID, Scopes: scopesList, }, clearKeys...)