From f5d474cbbac5b1c2e9bdd51a98bd9d0192cf1412 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 10 Mar 2026 15:02:14 +0100 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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...) From 3f33cce49965752662855748b595a10afe4bc4be Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 15:40:59 +0100 Subject: [PATCH 12/20] Validate discovered host is non-empty after discovery login After GetDiscoveredHost(), check that the returned host is not empty before saving it to the profile. An empty host would produce a broken profile entry. Return a clear error with the fallback tip so the user knows they can retry with --host. --- cmd/auth/login.go | 3 +++ cmd/auth/login_test.go | 25 +++++++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 0548ce6247..245cf833f1 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -492,6 +492,9 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati } discoveredHost := arg.GetDiscoveredHost() + if discoveredHost == "" { + return fmt.Errorf("login succeeded but no workspace host was discovered" + discoveryFallbackTip) + } // Get the token for introspection tok, err := persistentAuth.Token() diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 62c18974a9..dde76ddb4a 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -559,3 +559,28 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { // No warning should be logged when account IDs match. assert.Empty(t, logBuf.String()) } + +func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { + originalNewDiscoveryOAuthArgument := newDiscoveryOAuthArgument + originalNewDiscoveryPersistentAuth := newDiscoveryPersistentAuth + t.Cleanup(func() { + newDiscoveryOAuthArgument = originalNewDiscoveryOAuthArgument + newDiscoveryPersistentAuth = originalNewDiscoveryPersistentAuth + }) + + // Return arg without calling SetDiscoveredHost, so GetDiscoveredHost returns "". + newDiscoveryOAuthArgument = func(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + return u2m.NewBasicDiscoveryOAuthArgument(profileName) + } + + newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + return &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, nil + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err := discoveryLogin(ctx, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + require.Error(t, err) + assert.Contains(t, err.Error(), "no workspace host was discovered") +} From 7d5624dd6eb489699bfbf9db914e71e2e7181158 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 16:02:55 +0100 Subject: [PATCH 13/20] Reject --account-id, --workspace-id, --experimental-is-unified-host without --host in discovery login These flags are meaningless in discovery mode (login.databricks.com) and silently ignored before this change. For example, 'databricks auth login --account-id ' would go through workspace discovery but getProfileName() would generate an ACCOUNT- profile name for a workspace token. Now the CLI returns a clear error telling the user to provide --host when any of these flags are explicitly set. --- cmd/auth/login.go | 23 +++++++++++++++++++ cmd/auth/login_test.go | 51 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 245cf833f1..d5845f9298 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -166,6 +166,9 @@ depends on the existing profiles you have set in your configuration file if configureServerless { return errors.New("--configure-serverless requires --host to be specified") } + if err := validateDiscoveryFlagCompatibility(cmd); err != nil { + return err + } return discoveryLogin(ctx, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) } @@ -440,6 +443,26 @@ func shouldUseDiscovery(hostFlag string, args []string, existingProfile *profile return true } +// discoveryIncompatibleFlags lists flags that require --host and are incompatible +// with the discovery login flow via login.databricks.com. +var discoveryIncompatibleFlags = []string{ + "account-id", + "workspace-id", + "experimental-is-unified-host", +} + +// validateDiscoveryFlagCompatibility returns an error if any flags that require +// --host were explicitly set. These flags are meaningless in discovery mode +// and could lead to incorrect profile configuration. +func validateDiscoveryFlagCompatibility(cmd *cobra.Command) error { + for _, name := range discoveryIncompatibleFlags { + if cmd.Flag(name).Changed { + return fmt.Errorf("--%s requires --host to be specified", name) + } + } + return nil +} + // 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 { diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index dde76ddb4a..1c0b65ae32 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" @@ -391,6 +392,56 @@ func TestSplitScopes(t *testing.T) { } } +func TestValidateDiscoveryFlagCompatibility(t *testing.T) { + tests := []struct { + name string + setFlag string + flagVal string + wantErr string + }{ + { + name: "account-id is incompatible", + setFlag: "account-id", + flagVal: "abc123", + wantErr: "--account-id requires --host to be specified", + }, + { + name: "workspace-id is incompatible", + setFlag: "workspace-id", + flagVal: "12345", + wantErr: "--workspace-id requires --host to be specified", + }, + { + name: "experimental-is-unified-host is incompatible", + setFlag: "experimental-is-unified-host", + flagVal: "true", + wantErr: "--experimental-is-unified-host requires --host to be specified", + }, + { + name: "no flags set is ok", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cmd := &cobra.Command{} + cmd.Flags().String("account-id", "", "") + cmd.Flags().String("workspace-id", "", "") + cmd.Flags().Bool("experimental-is-unified-host", false, "") + + if tt.setFlag != "" { + require.NoError(t, cmd.Flags().Set(tt.setFlag, tt.flagVal)) + } + + err := validateDiscoveryFlagCompatibility(cmd) + if tt.wantErr != "" { + assert.EqualError(t, err, tt.wantErr) + } else { + assert.NoError(t, err) + } + }) + } +} + func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") From c1e0fc4b4be80ce96c70bd1d34646b31ba01eb16 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 17:11:02 +0100 Subject: [PATCH 14/20] Fix lint: use errors.New per perfsprint linter rule --- 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 d5845f9298..91565050d7 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -516,7 +516,7 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { - return fmt.Errorf("login succeeded but no workspace host was discovered" + discoveryFallbackTip) + return errors.New("login succeeded but no workspace host was discovered" + discoveryFallbackTip) } // Get the token for introspection From c6391afe50cfbe05eb8b8fd3e19646fb0aea165a Mon Sep 17 00:00:00 2001 From: simon Date: Mon, 16 Mar 2026 20:21:36 +0100 Subject: [PATCH 15/20] Accept HTTP client parameter in IntrospectToken Allow callers to pass a configured *http.Client so corporate proxy and custom CA certificate settings are respected. Falls back to http.DefaultClient when nil for backward compatibility. Co-authored-by: Isaac --- libs/auth/introspect.go | 7 +++++-- libs/auth/introspect_test.go | 10 +++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/libs/auth/introspect.go b/libs/auth/introspect.go index bd4deaae84..6558801014 100644 --- a/libs/auth/introspect.go +++ b/libs/auth/introspect.go @@ -31,7 +31,10 @@ 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, accessToken string) (*IntrospectionResult, error) { +func IntrospectToken(ctx context.Context, host, accessToken string, httpClient *http.Client) (*IntrospectionResult, error) { + if httpClient == nil { + httpClient = http.DefaultClient + } endpoint := strings.TrimSuffix(host, "/") + "/api/2.0/tokens/introspect" req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) if err != nil { @@ -39,7 +42,7 @@ func IntrospectToken(ctx context.Context, host, accessToken string) (*Introspect } req.Header.Set("Authorization", "Bearer "+accessToken) - resp, err := http.DefaultClient.Do(req) + resp, err := httpClient.Do(req) if err != nil { return nil, fmt.Errorf("calling introspection endpoint: %w", err) } diff --git a/libs/auth/introspect_test.go b/libs/auth/introspect_test.go index d6ce2e83e6..655669390b 100644 --- a/libs/auth/introspect_test.go +++ b/libs/auth/introspect_test.go @@ -23,7 +23,7 @@ func TestIntrospectToken_Success(t *testing.T) { })) defer server.Close() - result, err := IntrospectToken(t.Context(), server.URL, "test-token") + result, err := IntrospectToken(t.Context(), server.URL, "test-token", nil) require.NoError(t, err) assert.Equal(t, "a1b1c234-5678-90ab-cdef-1234567890ab", result.AccountID) assert.Equal(t, "2548836972759138", result.WorkspaceID) @@ -43,7 +43,7 @@ func TestIntrospectToken_ZeroWorkspaceID(t *testing.T) { })) defer server.Close() - result, err := IntrospectToken(t.Context(), server.URL, "test-token") + result, err := IntrospectToken(t.Context(), server.URL, "test-token", nil) require.NoError(t, err) assert.Equal(t, "abc-123", result.AccountID) assert.Empty(t, result.WorkspaceID) @@ -55,7 +55,7 @@ func TestIntrospectToken_HTTPError(t *testing.T) { })) defer server.Close() - _, err := IntrospectToken(t.Context(), server.URL, "test-token") + _, err := IntrospectToken(t.Context(), server.URL, "test-token", nil) assert.ErrorContains(t, err, "status 403") } @@ -66,7 +66,7 @@ func TestIntrospectToken_MalformedJSON(t *testing.T) { })) defer server.Close() - _, err := IntrospectToken(t.Context(), server.URL, "test-token") + _, err := IntrospectToken(t.Context(), server.URL, "test-token", nil) assert.ErrorContains(t, err, "decoding introspection response") } @@ -79,6 +79,6 @@ func TestIntrospectToken_VerifyRequestDetails(t *testing.T) { })) defer server.Close() - _, err := IntrospectToken(t.Context(), server.URL, "my-secret-token") + _, err := IntrospectToken(t.Context(), server.URL, "my-secret-token", nil) require.NoError(t, err) } From 28bbb689bd88b27ddf5a73f6fe299dbd0396d340 Mon Sep 17 00:00:00 2001 From: simon Date: Mon, 16 Mar 2026 20:21:51 +0100 Subject: [PATCH 16/20] Fix discovery login: HTTP client, scopes, stale fields, error messages Addresses four issues found in code review: 1. Pass a dedicated HTTP client (cloned transport) to IntrospectToken instead of nil, so proxy/TLS settings are inherited without sharing global state. 2. Reuse existing profile scopes during discovery relogin when --scopes is not explicitly set. Add regression tests for both fallback and override cases. 3. Add workspace_id to the clearKeys list so stale values from older hostless/unified profiles are removed. Fresh introspection values are re-written after clearing. Add regression tests for stale field cleanup and fresh workspace_id writes. 4. Replace string concatenation with discoveryFallbackTip with a discoveryErr helper for consistent, readable error construction. Co-authored-by: Isaac --- cmd/auth/login.go | 45 ++++++-- cmd/auth/login_test.go | 242 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 277 insertions(+), 10 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index faea99e775..ff357f7606 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "net/http" "runtime" "strings" "time" @@ -50,6 +51,15 @@ const ( discoveryFallbackTip = "\n\nTip: you can specify a workspace directly with: databricks auth login --host " ) +// discoveryErr wraps an error (or creates a new one) and appends the +// discovery fallback tip so users know they can bypass login.databricks.com. +func discoveryErr(msg string, err error) error { + if err != nil { + return fmt.Errorf("%s: %w%s", msg, err, discoveryFallbackTip) + } + return fmt.Errorf("%s%s", msg, discoveryFallbackTip) +} + type discoveryPersistentAuth interface { Challenge() error Token() (*oauth2.Token, error) @@ -487,10 +497,13 @@ func openURLSuppressingStderr(url string) 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"+discoveryFallbackTip, err) + return discoveryErr("setting up login.databricks.com", err) } scopesList := splitScopes(scopes) + if len(scopesList) == 0 && existingProfile != nil && existingProfile.Scopes != "" { + scopesList = splitScopes(existingProfile.Scopes) + } opts := []u2m.PersistentAuthOption{ u2m.WithOAuthArgument(arg), @@ -507,18 +520,18 @@ 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"+discoveryFallbackTip, err) + return discoveryErr("setting up login.databricks.com", 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"+discoveryFallbackTip, err) + return discoveryErr("login via login.databricks.com failed", err) } discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { - return errors.New("login succeeded but no workspace host was discovered" + discoveryFallbackTip) + return discoveryErr("login succeeded but no workspace host was discovered", nil) } // Get the token for introspection @@ -527,9 +540,15 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati return fmt.Errorf("retrieving token after login: %w", err) } - // Best-effort introspection for metadata + // Best-effort introspection for metadata. Use a dedicated HTTP client + // (not http.DefaultClient) so corporate proxy and custom CA settings + // from the default transport are inherited without sharing global state. + httpClient := &http.Client{} + if t, ok := http.DefaultTransport.(*http.Transport); ok { + httpClient.Transport = t.Clone() + } var workspaceID string - introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken) + introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken, httpClient) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { @@ -547,6 +566,18 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati } configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE") + clearKeys := oauthLoginClearKeys() + // Clear stale routing fields that may exist from a previous login to a + // different host type. Discovery always produces a workspace-level profile. + // workspace_id is cleared and re-written only if introspection returned a + // fresh value, preventing stale IDs from surviving failed introspection. + clearKeys = append(clearKeys, + "account_id", + "workspace_id", + "experimental_is_unified_host", + "cluster_id", + "serverless_compute_id", + ) err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: discoveredHost, @@ -554,7 +585,7 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati WorkspaceID: workspaceID, Scopes: scopesList, ConfigFile: configFile, - }, oauthLoginClearKeys()...) + }, clearKeys...) if err != nil { if configFile != "" { return fmt.Errorf("saving profile %q to %s: %w", profileName, configFile, err) diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a2762508b9..9ce5f7392a 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -5,6 +5,7 @@ import ( "context" "errors" "log/slog" + "net/http" "os" "path/filepath" "sync" @@ -485,7 +486,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { }, nil } - introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { assert.Equal(t, "https://workspace.example.com", host) assert.Equal(t, "test-token", accessToken) return nil, errors.New("introspection failed") @@ -535,7 +536,7 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { }, nil } - introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { return &auth.IntrospectionResult{ AccountID: "new-account-id", WorkspaceID: "12345", @@ -599,7 +600,7 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { }, nil } - introspectToken = func(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { return &auth.IntrospectionResult{ AccountID: "same-account-id", WorkspaceID: "12345", @@ -647,3 +648,238 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } + +func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(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, _ *http.Client) (*auth.IntrospectionResult, error) { + return nil, errors.New("introspection failed") + } + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + Host: "https://old-workspace.example.com", + Scopes: "sql,clusters", + } + + // No --scopes flag (empty string), should fall back to existing profile scopes. + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, 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, "sql,clusters", savedProfile.Scopes) +} + +func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(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, _ *http.Client) (*auth.IntrospectionResult, error) { + return nil, errors.New("introspection failed") + } + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + Host: "https://old-workspace.example.com", + Scopes: "sql,clusters", + } + + // Explicit --scopes flag should override existing profile scopes. + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis", existingProfile, 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, "all-apis", savedProfile.Scopes) +} + +func TestDiscoveryLogin_ClearsStaleRoutingFieldsFromUnifiedProfile(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + + // Pre-populate a profile that looks like an older hostless/unified login. + initialConfig := `[DISCOVERY] +host = https://old-unified.databricks.com +account_id = old-account +workspace_id = 999999 +experimental_is_unified_host = true +auth_type = databricks-cli +` + err := os.WriteFile(configPath, []byte(initialConfig), 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://new-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 + } + + // Introspection fails, so workspace_id should be cleared (not left stale). + introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { + return nil, errors.New("introspection unavailable") + } + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + Host: "https://old-unified.databricks.com", + AccountID: "old-account", + WorkspaceID: "999999", + IsUnifiedHost: true, + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, 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://new-workspace.example.com", savedProfile.Host) + // Stale routing fields must be cleared. + assert.Empty(t, savedProfile.AccountID, "stale account_id should be cleared") + assert.Empty(t, savedProfile.WorkspaceID, "stale workspace_id should be cleared on introspection failure") + assert.False(t, savedProfile.IsUnifiedHost, "stale experimental_is_unified_host should be cleared") +} + +func TestDiscoveryLogin_IntrospectionWritesFreshWorkspaceID(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + + // Pre-populate with stale workspace_id. + initialConfig := `[DISCOVERY] +host = https://old.example.com +workspace_id = 111111 +auth_type = databricks-cli +` + err := os.WriteFile(configPath, []byte(initialConfig), 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://new-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 + } + + // Introspection succeeds with a fresh workspace_id. + introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { + return &auth.IntrospectionResult{ + AccountID: "fresh-account", + WorkspaceID: "222222", + }, nil + } + + existingProfile := &profile.Profile{ + Name: "DISCOVERY", + Host: "https://old.example.com", + WorkspaceID: "111111", + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, 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://new-workspace.example.com", savedProfile.Host) + assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") +} From 9b683f9fafafef1e5cf7fbaa8e2e78c6ed1bf74b Mon Sep 17 00:00:00 2001 From: simon Date: Mon, 16 Mar 2026 20:50:45 +0100 Subject: [PATCH 17/20] Use SDK config for introspection HTTP client TLS settings Build the introspection HTTP client from a resolved SDK config.Config instead of cloning http.DefaultTransport. This picks up TLS settings (InsecureSkipVerify, custom transport) from environment variables, matching how the rest of the SDK handles HTTP connections. Add httpClientFromConfig helper with tests, and add a test verifying IntrospectToken uses the supplied HTTP client. Co-authored-by: Isaac --- cmd/auth/login.go | 30 ++++++++++++++++++++++++------ cmd/auth/login_test.go | 25 +++++++++++++++++++++++++ libs/auth/introspect_test.go | 28 ++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 6 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index ff357f7606..373c41d2e9 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -2,6 +2,7 @@ package auth import ( "context" + "crypto/tls" "errors" "fmt" "io" @@ -540,13 +541,15 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati return fmt.Errorf("retrieving token after login: %w", err) } - // Best-effort introspection for metadata. Use a dedicated HTTP client - // (not http.DefaultClient) so corporate proxy and custom CA settings - // from the default transport are inherited without sharing global state. - httpClient := &http.Client{} - if t, ok := http.DefaultTransport.(*http.Transport); ok { - httpClient.Transport = t.Clone() + // Best-effort introspection for metadata. Build the HTTP client from a + // resolved SDK config so that TLS settings (InsecureSkipVerify, custom + // transport) from environment variables are respected. + introspectCfg := &config.Config{ + Host: discoveredHost, + Token: tok.AccessToken, } + _ = introspectCfg.EnsureResolved() + httpClient := httpClientFromConfig(introspectCfg) var workspaceID string introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken, httpClient) if err != nil { @@ -657,3 +660,18 @@ func getBrowserFunc(cmd *cobra.Command) func(url string) error { } } } + +// httpClientFromConfig builds an *http.Client that respects the TLS settings +// from a resolved SDK config. This picks up InsecureSkipVerify from environment +// variables and any custom transport set on the config. +func httpClientFromConfig(cfg *config.Config) *http.Client { + c := &http.Client{} + if t, ok := cfg.HTTPTransport.(*http.Transport); ok && t != nil { + c.Transport = t.Clone() + } else if cfg.InsecureSkipVerify { + c.Transport = &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec + } + } + return c +} diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 9ce5f7392a..a2e23984ac 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -17,6 +17,7 @@ import ( "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/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" @@ -883,3 +884,27 @@ auth_type = databricks-cli assert.Equal(t, "https://new-workspace.example.com", savedProfile.Host) assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") } + +func TestHttpClientFromConfig_DefaultTransport(t *testing.T) { + cfg := &config.Config{Host: "https://example.com"} + c := httpClientFromConfig(cfg) + assert.Nil(t, c.Transport, "default config should produce nil transport (uses http.DefaultTransport)") +} + +func TestHttpClientFromConfig_InsecureSkipVerify(t *testing.T) { + cfg := &config.Config{Host: "https://example.com", InsecureSkipVerify: true} + c := httpClientFromConfig(cfg) + require.NotNil(t, c.Transport) + transport, ok := c.Transport.(*http.Transport) + require.True(t, ok) + assert.True(t, transport.TLSClientConfig.InsecureSkipVerify) +} + +func TestHttpClientFromConfig_CustomTransport(t *testing.T) { + custom := &http.Transport{} + cfg := &config.Config{Host: "https://example.com", HTTPTransport: custom} + c := httpClientFromConfig(cfg) + require.NotNil(t, c.Transport) + // Should be a clone, not the same pointer. + assert.NotSame(t, custom, c.Transport) +} diff --git a/libs/auth/introspect_test.go b/libs/auth/introspect_test.go index 655669390b..3ea5f10717 100644 --- a/libs/auth/introspect_test.go +++ b/libs/auth/introspect_test.go @@ -82,3 +82,31 @@ func TestIntrospectToken_VerifyRequestDetails(t *testing.T) { _, err := IntrospectToken(t.Context(), server.URL, "my-secret-token", nil) require.NoError(t, err) } + +// roundTripFunc adapts a function into an http.RoundTripper. +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func TestIntrospectToken_UsesSuppliedHTTPClient(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":"x","workspace_id":1}}}`)) + })) + defer server.Close() + + var transportUsed bool + customClient := &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + transportUsed = true + return http.DefaultTransport.RoundTrip(req) + }), + } + + result, err := IntrospectToken(t.Context(), server.URL, "test-token", customClient) + require.NoError(t, err) + assert.True(t, transportUsed, "expected IntrospectToken to use the supplied HTTP client") + assert.Equal(t, "x", result.AccountID) +} From 67512b88fb3097fbf697edacfd7f97687d65df59 Mon Sep 17 00:00:00 2001 From: simon Date: Mon, 16 Mar 2026 21:02:28 +0100 Subject: [PATCH 18/20] Simplify introspection HTTP client to clone DefaultTransport Replace httpClientFromConfig with httpClientForIntrospection that clones http.DefaultTransport. The previous approach tried to load TLS settings via EnsureResolved on a config with Host+Token already set, which short-circuits the config-file loader and never picks up profile TLS settings. The InsecureSkipVerify branch also created a zero-value transport missing proxy, dialer, and HTTP/2 settings. Since discovery login may not have a profile yet, the correct approach is to clone DefaultTransport which inherits system CA certs, proxy settings, timeouts, and connection pooling. Co-authored-by: Isaac --- cmd/auth/login.go | 31 ++++++++----------------------- cmd/auth/login_test.go | 24 ++++-------------------- 2 files changed, 12 insertions(+), 43 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 373c41d2e9..b6af8538ae 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -2,7 +2,6 @@ package auth import ( "context" - "crypto/tls" "errors" "fmt" "io" @@ -541,15 +540,8 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati return fmt.Errorf("retrieving token after login: %w", err) } - // Best-effort introspection for metadata. Build the HTTP client from a - // resolved SDK config so that TLS settings (InsecureSkipVerify, custom - // transport) from environment variables are respected. - introspectCfg := &config.Config{ - Host: discoveredHost, - Token: tok.AccessToken, - } - _ = introspectCfg.EnsureResolved() - httpClient := httpClientFromConfig(introspectCfg) + // Best-effort introspection for metadata. + httpClient := httpClientForIntrospection() var workspaceID string introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken, httpClient) if err != nil { @@ -661,17 +653,10 @@ func getBrowserFunc(cmd *cobra.Command) func(url string) error { } } -// httpClientFromConfig builds an *http.Client that respects the TLS settings -// from a resolved SDK config. This picks up InsecureSkipVerify from environment -// variables and any custom transport set on the config. -func httpClientFromConfig(cfg *config.Config) *http.Client { - c := &http.Client{} - if t, ok := cfg.HTTPTransport.(*http.Transport); ok && t != nil { - c.Transport = t.Clone() - } else if cfg.InsecureSkipVerify { - c.Transport = &http.Transport{ - TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec - } - } - return c +// httpClientForIntrospection returns an *http.Client suitable for the +// best-effort token introspection call. It clones http.DefaultTransport +// to inherit system CA certs, proxy settings, and timeouts. +func httpClientForIntrospection() *http.Client { + transport := http.DefaultTransport.(*http.Transport).Clone() + return &http.Client{Transport: transport} } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a2e23984ac..937f6dca52 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -17,7 +17,6 @@ import ( "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/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" @@ -885,26 +884,11 @@ auth_type = databricks-cli assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") } -func TestHttpClientFromConfig_DefaultTransport(t *testing.T) { - cfg := &config.Config{Host: "https://example.com"} - c := httpClientFromConfig(cfg) - assert.Nil(t, c.Transport, "default config should produce nil transport (uses http.DefaultTransport)") -} - -func TestHttpClientFromConfig_InsecureSkipVerify(t *testing.T) { - cfg := &config.Config{Host: "https://example.com", InsecureSkipVerify: true} - c := httpClientFromConfig(cfg) +func TestHttpClientForIntrospection(t *testing.T) { + c := httpClientForIntrospection() require.NotNil(t, c.Transport) transport, ok := c.Transport.(*http.Transport) require.True(t, ok) - assert.True(t, transport.TLSClientConfig.InsecureSkipVerify) -} - -func TestHttpClientFromConfig_CustomTransport(t *testing.T) { - custom := &http.Transport{} - cfg := &config.Config{Host: "https://example.com", HTTPTransport: custom} - c := httpClientFromConfig(cfg) - require.NotNil(t, c.Transport) - // Should be a clone, not the same pointer. - assert.NotSame(t, custom, c.Transport) + // Should be a clone of DefaultTransport, not the same pointer. + assert.NotSame(t, http.DefaultTransport, transport) } From 186a91ded95459fa63d7100de41b26c6e40c4c4c Mon Sep 17 00:00:00 2001 From: simon Date: Sun, 22 Mar 2026 08:29:05 +0100 Subject: [PATCH 19/20] Add acceptance test for discovery login flow Adds an e2e acceptance test that verifies `databricks auth login` without `--host` uses the discovery flow via login.databricks.com. The test uses a custom browser script (discovery_browser.py) that simulates the login.databricks.com redirect, injecting the testserver as the workspace issuer via the `iss` callback parameter. Verifies: - Login succeeds and profile is saved with discovered host - Introspection endpoint is called (workspace_id in profile) - Profile is listed with correct host Co-authored-by: Isaac --- acceptance/bin/discovery_browser.py | 66 +++++++++++++++++++ .../auth/login/discovery/out.databrickscfg | 10 +++ .../cmd/auth/login/discovery/out.test.toml | 5 ++ .../cmd/auth/login/discovery/output.txt | 14 ++++ acceptance/cmd/auth/login/discovery/script | 14 ++++ acceptance/cmd/auth/login/discovery/test.toml | 18 +++++ 6 files changed, 127 insertions(+) create mode 100755 acceptance/bin/discovery_browser.py create mode 100644 acceptance/cmd/auth/login/discovery/out.databrickscfg create mode 100644 acceptance/cmd/auth/login/discovery/out.test.toml create mode 100644 acceptance/cmd/auth/login/discovery/output.txt create mode 100644 acceptance/cmd/auth/login/discovery/script create mode 100644 acceptance/cmd/auth/login/discovery/test.toml diff --git a/acceptance/bin/discovery_browser.py b/acceptance/bin/discovery_browser.py new file mode 100755 index 0000000000..9e80c25a9e --- /dev/null +++ b/acceptance/bin/discovery_browser.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +""" +Simulates the login.databricks.com discovery flow for acceptance tests. + +When the CLI opens this "browser" with the login.databricks.com URL, +the script extracts the OAuth parameters from the destination_url, +constructs a callback to localhost with an iss parameter pointing +at the testserver, and fetches it. + +Usage: discovery_browser.py +""" + +import os +import sys +import urllib.parse +import urllib.request + +if len(sys.argv) < 2: + sys.stderr.write("Usage: discovery_browser.py \n") + sys.exit(1) + +url = sys.argv[1] +parsed = urllib.parse.urlparse(url) +top_params = urllib.parse.parse_qs(parsed.query) + +destination_url = top_params.get("destination_url", [None])[0] +if not destination_url: + sys.stderr.write(f"No destination_url found in: {url}\n") + sys.exit(1) + +dest_parsed = urllib.parse.urlparse(destination_url) +dest_params = urllib.parse.parse_qs(dest_parsed.query) + +redirect_uri = dest_params.get("redirect_uri", [None])[0] +state = dest_params.get("state", [None])[0] + +if not redirect_uri or not state: + sys.stderr.write(f"Missing redirect_uri or state in destination_url: {destination_url}\n") + sys.exit(1) + +# The testserver's host acts as the workspace issuer. +testserver_host = os.environ.get("DATABRICKS_HOST", "") +if not testserver_host: + sys.stderr.write("DATABRICKS_HOST not set\n") + sys.exit(1) + +issuer = testserver_host.rstrip("/") + "/oidc" + +# Build the callback URL with code, state, and iss (the workspace issuer). +callback_params = urllib.parse.urlencode({ + "code": "oauth-code", + "state": state, + "iss": issuer, +}) +callback_url = f"{redirect_uri}?{callback_params}" + +try: + response = urllib.request.urlopen(callback_url) + if response.status != 200: + sys.stderr.write(f"Callback failed: {callback_url} (status {response.status})\n") + sys.exit(1) +except Exception as e: + sys.stderr.write(f"Callback failed: {callback_url} ({e})\n") + sys.exit(1) + +sys.exit(0) diff --git a/acceptance/cmd/auth/login/discovery/out.databrickscfg b/acceptance/cmd/auth/login/discovery/out.databrickscfg new file mode 100644 index 0000000000..d6e17b7595 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery/out.databrickscfg @@ -0,0 +1,10 @@ +; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified. +[DEFAULT] + +[discovery-test] +host = [DATABRICKS_URL] +workspace_id = 12345 +auth_type = databricks-cli + +[__settings__] +default_profile = discovery-test diff --git a/acceptance/cmd/auth/login/discovery/out.test.toml b/acceptance/cmd/auth/login/discovery/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/discovery/output.txt b/acceptance/cmd/auth/login/discovery/output.txt new file mode 100644 index 0000000000..c687b07fd5 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery/output.txt @@ -0,0 +1,14 @@ + +>>> [CLI] auth login --profile discovery-test +Opening login.databricks.com in your browser... +Profile discovery-test was successfully saved + +>>> [CLI] auth profiles +Name Host Valid +discovery-test (Default) [DATABRICKS_URL] YES + +>>> print_requests.py --get //tokens/introspect +{ + "method": "GET", + "path": "/api/2.0/tokens/introspect" +} diff --git a/acceptance/cmd/auth/login/discovery/script b/acceptance/cmd/auth/login/discovery/script new file mode 100644 index 0000000000..4ae4c682d9 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery/script @@ -0,0 +1,14 @@ +sethome "./home" + +# Use the discovery browser script that simulates login.databricks.com +export BROWSER="discovery_browser.py" + +trace $CLI auth login --profile discovery-test + +trace $CLI auth profiles + +# Verify the introspection endpoint was called (workspace_id in profile confirms this too). +trace print_requests.py --get //tokens/introspect + +# Track the .databrickscfg file that was created to surface changes. +mv "./home/.databrickscfg" "./out.databrickscfg" diff --git a/acceptance/cmd/auth/login/discovery/test.toml b/acceptance/cmd/auth/login/discovery/test.toml new file mode 100644 index 0000000000..d430a86e03 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery/test.toml @@ -0,0 +1,18 @@ +Ignore = [ + "home" +] +RecordRequests = true + +# Override the introspection endpoint so we can verify it gets called. +[[Server]] +Pattern = "GET /api/2.0/tokens/introspect" +Response.Body = ''' +{ + "principal_context": { + "authentication_scope": { + "account_id": "test-account-123", + "workspace_id": 12345 + } + } +} +''' From fa4d62b7122289e869ae6445aefe0adc95dd1877 Mon Sep 17 00:00:00 2001 From: simon Date: Sun, 22 Mar 2026 08:31:03 +0100 Subject: [PATCH 20/20] Address review: discoveryClient interface, consolidate flags, cleanup Replace package-level var overrides (newDiscoveryOAuthArgument, newDiscoveryPersistentAuth, introspectToken) with a discoveryClient interface injected into discoveryLogin(). Tests use fakeDiscoveryClient instead of save/restore globals, reducing boilerplate. Other review feedback: - Consolidate --configure-cluster and --configure-serverless into discoveryIncompatibleFlags (removes separate if blocks) - Remove httpClientForIntrospection(); default http.Client respects system proxy and CA config already - Remove superfluous empty-string guard in splitScopes - Improve comment on key clearing rationale --- acceptance/bin/discovery_browser.py | 12 +- cmd/auth/login.go | 60 +++-- cmd/auth/login_test.go | 333 +++++++++++----------------- 3 files changed, 163 insertions(+), 242 deletions(-) diff --git a/acceptance/bin/discovery_browser.py b/acceptance/bin/discovery_browser.py index 9e80c25a9e..42099fa06d 100755 --- a/acceptance/bin/discovery_browser.py +++ b/acceptance/bin/discovery_browser.py @@ -47,11 +47,13 @@ issuer = testserver_host.rstrip("/") + "/oidc" # Build the callback URL with code, state, and iss (the workspace issuer). -callback_params = urllib.parse.urlencode({ - "code": "oauth-code", - "state": state, - "iss": issuer, -}) +callback_params = urllib.parse.urlencode( + { + "code": "oauth-code", + "state": state, + "iss": issuer, + } +) callback_url = f"{redirect_uri}?{callback_params}" try: diff --git a/cmd/auth/login.go b/cmd/auth/login.go index b6af8538ae..379ce1a4c2 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "io" - "net/http" "runtime" "strings" "time" @@ -66,13 +65,27 @@ type discoveryPersistentAuth interface { Close() error } -var newDiscoveryOAuthArgument = u2m.NewBasicDiscoveryOAuthArgument +// discoveryClient abstracts the external dependencies of discoveryLogin so +// they can be replaced in tests without package-level variable mutation. +type discoveryClient interface { + NewOAuthArgument(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) + NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) + IntrospectToken(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) +} + +type defaultDiscoveryClient struct{} + +func (d *defaultDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + return u2m.NewBasicDiscoveryOAuthArgument(profileName) +} -var newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { +func (d *defaultDiscoveryClient) NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { return u2m.NewPersistentAuth(ctx, opts...) } -var introspectToken = auth.IntrospectToken +func (d *defaultDiscoveryClient) IntrospectToken(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + return auth.IntrospectToken(ctx, host, accessToken, nil) +} func newLoginCommand(authArguments *auth.AuthArguments) *cobra.Command { defaultConfigPath := "~/.databrickscfg" @@ -170,16 +183,10 @@ 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) { - if configureCluster { - return errors.New("--configure-cluster requires --host to be specified") - } - if configureServerless { - return errors.New("--configure-serverless requires --host to be specified") - } if err := validateDiscoveryFlagCompatibility(cmd); err != nil { return err } - return discoveryLogin(ctx, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) + return discoveryLogin(ctx, &defaultDiscoveryClient{}, profileName, loginTimeout, scopes, existingProfile, getBrowserFunc(cmd)) } // Load unified host flags from the profile if not explicitly set via CLI flag @@ -461,6 +468,8 @@ var discoveryIncompatibleFlags = []string{ "account-id", "workspace-id", "experimental-is-unified-host", + "configure-cluster", + "configure-serverless", } // validateDiscoveryFlagCompatibility returns an error if any flags that require @@ -494,8 +503,8 @@ 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, existingProfile *profile.Profile, browserFunc func(string) error) error { - arg, err := newDiscoveryOAuthArgument(profileName) +func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error) error { + arg, err := dc.NewOAuthArgument(profileName) if err != nil { return discoveryErr("setting up login.databricks.com", err) } @@ -518,7 +527,7 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - persistentAuth, err := newDiscoveryPersistentAuth(ctx, opts...) + persistentAuth, err := dc.NewPersistentAuth(ctx, opts...) if err != nil { return discoveryErr("setting up login.databricks.com", err) } @@ -541,9 +550,8 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati } // Best-effort introspection for metadata. - httpClient := httpClientForIntrospection() var workspaceID string - introspection, err := introspectToken(ctx, discoveredHost, tok.AccessToken, httpClient) + introspection, err := dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { @@ -562,10 +570,11 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE") clearKeys := oauthLoginClearKeys() - // Clear stale routing fields that may exist from a previous login to a - // different host type. Discovery always produces a workspace-level profile. - // workspace_id is cleared and re-written only if introspection returned a - // fresh value, preventing stale IDs from surviving failed introspection. + // Discovery login always produces a workspace-level profile pointing at the + // discovered host. Any previous routing metadata (account_id, workspace_id, + // is_unified_host, cluster_id, serverless_compute_id) from a prior login to + // a different host type must be cleared so they don't leak into the new + // profile. workspace_id is re-added only when introspection succeeds. clearKeys = append(clearKeys, "account_id", "workspace_id", @@ -594,9 +603,6 @@ func discoveryLogin(ctx context.Context, profileName string, timeout time.Durati // 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) @@ -652,11 +658,3 @@ func getBrowserFunc(cmd *cobra.Command) func(url string) error { } } } - -// httpClientForIntrospection returns an *http.Client suitable for the -// best-effort token introspection call. It clones http.DefaultTransport -// to inherit system CA certs, proxy settings, and timeouts. -func httpClientForIntrospection() *http.Client { - transport := http.DefaultTransport.(*http.Transport).Clone() - return &http.Client{Transport: transport} -} diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 937f6dca52..670ff7c211 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -5,7 +5,6 @@ import ( "context" "errors" "log/slog" - "net/http" "os" "path/filepath" "sync" @@ -70,6 +69,41 @@ func (f *fakeDiscoveryPersistentAuth) Close() error { return nil } +type fakeDiscoveryClient struct { + oauthArg *u2m.BasicDiscoveryOAuthArgument + oauthArgErr error + persistentAuth discoveryPersistentAuth + persistentAuthErr error + introspection *auth.IntrospectionResult + introspectionErr error + // For assertions + introspectHost string + introspectToken string +} + +func (f *fakeDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { + if f.oauthArgErr != nil { + return nil, f.oauthArgErr + } + return f.oauthArg, nil +} + +func (f *fakeDiscoveryClient) NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { + if f.persistentAuthErr != nil { + return nil, f.persistentAuthErr + } + return f.persistentAuth, nil +} + +func (f *fakeDiscoveryClient) IntrospectToken(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) { + f.introspectHost = host + f.introspectToken = accessToken + if f.introspectionErr != nil { + return nil, f.introspectionErr + } + return f.introspection, nil +} + func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg") @@ -430,6 +464,18 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) { flagVal: "true", wantErr: "--experimental-is-unified-host requires --host to be specified", }, + { + name: "configure-cluster is incompatible", + setFlag: "configure-cluster", + flagVal: "true", + wantErr: "--configure-cluster requires --host to be specified", + }, + { + name: "configure-serverless is incompatible", + setFlag: "configure-serverless", + flagVal: "true", + wantErr: "--configure-serverless requires --host to be specified", + }, { name: "no flags set is ok", }, @@ -440,6 +486,8 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) { cmd.Flags().String("account-id", "", "") cmd.Flags().String("workspace-id", "", "") cmd.Flags().Bool("experimental-is-unified-host", false, "") + cmd.Flags().Bool("configure-cluster", false, "") + cmd.Flags().Bool("configure-serverless", false, "") if tt.setFlag != "" { require.NoError(t, cmd.Flags().Set(tt.setFlag, tt.flagVal)) @@ -462,40 +510,25 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { 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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil - } - - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - assert.Equal(t, "https://workspace.example.com", host) - assert.Equal(t, "test-token", accessToken) - return nil, errors.New("introspection failed") + }, + introspectionErr: errors.New("introspection failed"), } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) require.NoError(t, err) + assert.Equal(t, "https://workspace.example.com", dc.introspectHost) + assert.Equal(t, "test-token", dc.introspectToken) + savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) require.NoError(t, err) require.NotNil(t, savedProfile) @@ -512,35 +545,19 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { 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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil - } - - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return &auth.IntrospectionResult{ + }, + introspection: &auth.IntrospectionResult{ AccountID: "new-account-id", WorkspaceID: "12345", - }, nil + }, } // Set up a logger that captures log records to verify the warning. @@ -554,7 +571,7 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { AccountID: "old-account-id", } - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) // Verify warning about mismatched account IDs was logged. @@ -576,35 +593,19 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { 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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil - } - - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return &auth.IntrospectionResult{ + }, + introspection: &auth.IntrospectionResult{ AccountID: "same-account-id", WorkspaceID: "12345", - }, nil + }, } var logBuf logBuffer @@ -617,7 +618,7 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { AccountID: "same-account-id", } - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) // No warning should be logged when account IDs match. @@ -625,26 +626,19 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { } func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { - originalNewDiscoveryOAuthArgument := newDiscoveryOAuthArgument - originalNewDiscoveryPersistentAuth := newDiscoveryPersistentAuth - t.Cleanup(func() { - newDiscoveryOAuthArgument = originalNewDiscoveryOAuthArgument - newDiscoveryPersistentAuth = originalNewDiscoveryPersistentAuth - }) - // Return arg without calling SetDiscoveredHost, so GetDiscoveredHost returns "". - newDiscoveryOAuthArgument = func(profileName string) (*u2m.BasicDiscoveryOAuthArgument, error) { - return u2m.NewBasicDiscoveryOAuthArgument(profileName) - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil + }, } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err := discoveryLogin(ctx, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } @@ -656,32 +650,16 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { 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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil - } - - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return nil, errors.New("introspection failed") + }, + introspectionErr: errors.New("introspection failed"), } existingProfile := &profile.Profile{ @@ -692,7 +670,7 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { // No --scopes flag (empty string), should fall back to existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -709,32 +687,16 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { 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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") - newDiscoveryPersistentAuth = func(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) { - return &fakeDiscoveryPersistentAuth{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ token: &oauth2.Token{AccessToken: "test-token"}, - }, nil - } - - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return nil, errors.New("introspection failed") + }, + introspectionErr: errors.New("introspection failed"), } existingProfile := &profile.Profile{ @@ -745,7 +707,7 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { // Explicit --scopes flag should override existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -770,33 +732,17 @@ auth_type = databricks-cli 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://new-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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://new-workspace.example.com") // Introspection fails, so workspace_id should be cleared (not left stale). - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return nil, errors.New("introspection unavailable") + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspectionErr: errors.New("introspection unavailable"), } existingProfile := &profile.Profile{ @@ -808,7 +754,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -835,36 +781,20 @@ auth_type = databricks-cli 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://new-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 - } + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://new-workspace.example.com") // Introspection succeeds with a fresh workspace_id. - introspectToken = func(ctx context.Context, host, accessToken string, _ *http.Client) (*auth.IntrospectionResult, error) { - return &auth.IntrospectionResult{ + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspection: &auth.IntrospectionResult{ AccountID: "fresh-account", WorkspaceID: "222222", - }, nil + }, } existingProfile := &profile.Profile{ @@ -874,7 +804,7 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -883,12 +813,3 @@ auth_type = databricks-cli assert.Equal(t, "https://new-workspace.example.com", savedProfile.Host) assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") } - -func TestHttpClientForIntrospection(t *testing.T) { - c := httpClientForIntrospection() - require.NotNil(t, c.Transport) - transport, ok := c.Transport.(*http.Transport) - require.True(t, ok) - // Should be a clone of DefaultTransport, not the same pointer. - assert.NotSame(t, http.DefaultTransport, transport) -}