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

Author: simonfaltum

cmd/auth/login.go

@@ -630,24 +630,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",
@@ -658,10 +667,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",
@@ -673,6 +682,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, profileName string,
 		Profile:     profileName,
 		Host:        discoveredHost,
 		AuthType:    authTypeDatabricksCLI,
+		AccountID:   accountID,
 		WorkspaceID: workspaceID,
 		Scopes:      scopesList,
 		ConfigFile:  configFile,

cmd/auth/login_test.go

@@ -688,11 +688,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)
 }
 
@@ -826,6 +827,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")
@@ -921,5 +999,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")
 }