From f92e22672650e90599aa83ce736c21e588f6ae1e Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 24 Mar 2026 15:29:26 +0100 Subject: [PATCH 1/2] Add host discovery to login.databricks.com flow The discovery login flow via login.databricks.com now calls runHostDiscovery() on the discovered host to populate account_id, workspace_id, and DiscoveryURL from .well-known/databricks-config. This ensures profiles created via login.databricks.com have the same SPOG metadata as profiles created via the regular --host login path. Previously, the discovery flow relied solely on token introspection for workspace_id and deliberately skipped saving account_id. With the SPOG discovery changes now in place, account_id can be safely saved to profiles. Token introspection remains as a fallback when host metadata discovery is unavailable (e.g. classic workspace hosts). Co-authored-by: Isaac --- cmd/auth/login.go | 34 +++++++++++------- cmd/auth/login_test.go | 81 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 102 insertions(+), 13 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 29732066ce..874f654449 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -613,24 +613,33 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, return discoveryErr("login succeeded but no workspace host was discovered", nil) } - // Get the token for introspection + // Run host metadata discovery on the discovered host to detect SPOG hosts + // and populate account_id/workspace_id. This ensures profiles created via + // login.databricks.com have the same metadata as profiles created via the + // regular --host login path. + hostArgs := &auth.AuthArguments{Host: discoveredHost} + runHostDiscovery(ctx, hostArgs) + accountID := hostArgs.AccountID + workspaceID := hostArgs.WorkspaceID + + // Best-effort introspection as a fallback for workspace_id when host + // metadata discovery didn't return it (e.g. classic workspace hosts). tok, err := persistentAuth.Token() if err != nil { return fmt.Errorf("retrieving token after login: %w", err) } - // Best-effort introspection for metadata. - var workspaceID string introspection, err := dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { - // 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 + if workspaceID == "" { + workspaceID = introspection.WorkspaceID + } + if accountID == "" { + accountID = introspection.AccountID + } - // 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", @@ -641,10 +650,10 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE") clearKeys := oauthLoginClearKeys() // 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. + // discovered host. Any previous routing metadata (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. account_id + // and workspace_id are re-added from discovery/introspection results. clearKeys = append(clearKeys, "account_id", "workspace_id", @@ -656,6 +665,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string, Profile: profileName, Host: discoveredHost, AuthType: authTypeDatabricksCLI, + AccountID: accountID, WorkspaceID: workspaceID, Scopes: scopesList, ConfigFile: configFile, diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index eee4a5c38f..9b55ec5fe6 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -678,11 +678,12 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { 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). + // Account ID from introspection is now saved to the profile. 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, "new-account-id", savedProfile.AccountID) assert.Equal(t, "12345", savedProfile.WorkspaceID) } @@ -816,6 +817,83 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { assert.Equal(t, "all-apis", savedProfile.Scopes) } +func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { + // Start a mock server that returns SPOG discovery metadata. + server := newDiscoveryServer(t, map[string]any{ + "account_id": "discovered-account", + "workspace_id": "discovered-ws", + "oidc_endpoint": "https://spog.example.com/oidc/accounts/discovered-account", + }) + + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, ".databrickscfg") + err := os.WriteFile(configPath, []byte(""), 0o600) + require.NoError(t, err) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost(server.URL) + + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + // Introspection returns different values to verify discovery takes precedence. + introspection: &auth.IntrospectionResult{ + AccountID: "introspection-account", + WorkspaceID: "introspection-ws", + }, + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, 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, server.URL, savedProfile.Host) + assert.Equal(t, "discovered-account", savedProfile.AccountID, "account_id should come from host discovery") + assert.Equal(t, "discovered-ws", savedProfile.WorkspaceID, "workspace_id should come from host discovery") +} + +func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(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) + + // Use a host that won't respond to .well-known/databricks-config. + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") + + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspection: &auth.IntrospectionResult{ + AccountID: "introspection-account", + WorkspaceID: "introspection-ws", + }, + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, dc, "DISCOVERY", time.Second, "", nil, 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, "introspection-account", savedProfile.AccountID, "account_id should fall back to introspection") + assert.Equal(t, "introspection-ws", savedProfile.WorkspaceID, "workspace_id should fall back to introspection") +} + func TestDiscoveryLogin_ClearsStaleRoutingFieldsFromUnifiedProfile(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") @@ -911,5 +989,6 @@ auth_type = databricks-cli require.NoError(t, err) require.NotNil(t, savedProfile) assert.Equal(t, "https://new-workspace.example.com", savedProfile.Host) + assert.Equal(t, "fresh-account", savedProfile.AccountID, "account_id should be saved from introspection") assert.Equal(t, "222222", savedProfile.WorkspaceID, "workspace_id should be updated to fresh introspection value") } From ca5ac25abf40e2faa571248dc9da516de1dbadfd Mon Sep 17 00:00:00 2001 From: simon Date: Thu, 26 Mar 2026 23:23:36 +0100 Subject: [PATCH 2/2] Update discovery acceptance test for host discovery in login.databricks.com The login.databricks.com flow now calls runHostDiscovery() on the discovered host, which populates workspace_id from .well-known and account_id from introspection as a fallback. Update the expected output to reflect both fields being saved to the profile. Co-authored-by: Isaac --- acceptance/cmd/auth/login/discovery/out.databrickscfg | 3 ++- acceptance/cmd/auth/login/discovery/test.toml | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/acceptance/cmd/auth/login/discovery/out.databrickscfg b/acceptance/cmd/auth/login/discovery/out.databrickscfg index d6e17b7595..56763df71c 100644 --- a/acceptance/cmd/auth/login/discovery/out.databrickscfg +++ b/acceptance/cmd/auth/login/discovery/out.databrickscfg @@ -3,7 +3,8 @@ [discovery-test] host = [DATABRICKS_URL] -workspace_id = 12345 +account_id = test-account-123 +workspace_id = [NUMID] auth_type = databricks-cli [__settings__] diff --git a/acceptance/cmd/auth/login/discovery/test.toml b/acceptance/cmd/auth/login/discovery/test.toml index d430a86e03..074f73fe64 100644 --- a/acceptance/cmd/auth/login/discovery/test.toml +++ b/acceptance/cmd/auth/login/discovery/test.toml @@ -4,6 +4,9 @@ Ignore = [ RecordRequests = true # Override the introspection endpoint so we can verify it gets called. +# Host discovery (via .well-known/databricks-config) provides workspace_id. +# Introspection provides account_id as a fallback since the default +# discovery handler doesn't return it. [[Server]] Pattern = "GET /api/2.0/tokens/introspect" Response.Body = '''