Extract IsSPOG and ResolveConfigType into libs/auth

cmd/auth/profiles.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
-	"strings"
 	"sync"
 	"time"
 
@@ -58,38 +57,7 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 		return
 	}
 
-	// ConfigType() classifies based on the host URL prefix (accounts.* →
-	// AccountConfig, everything else → WorkspaceConfig). SPOG hosts don't
-	// match the accounts.* prefix so they're misclassified as WorkspaceConfig.
-	// Use the resolved DiscoveryURL (from .well-known/databricks-config) to
-	// detect SPOG hosts with account-scoped OIDC, matching the routing logic
-	// in auth.AuthArguments.ToOAuthArgument().
-	configType := cfg.ConfigType()
-	hasWorkspace := cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone
-
-	isAccountScopedOIDC := cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/")
-	if configType != config.AccountConfig && cfg.AccountID != "" && isAccountScopedOIDC {
-		if hasWorkspace {
-			configType = config.WorkspaceConfig
-		} else {
-			configType = config.AccountConfig
-		}
-	}
-
-	// Legacy backward compat: SDK v0.126.0 removed the UnifiedHost case from
-	// ConfigType(), so profiles with Experimental_IsUnifiedHost now get
-	// InvalidConfig instead of being routed to account/workspace validation.
-	// When .well-known is also unreachable (DiscoveryURL empty), the override
-	// above can't help. Fall back to workspace_id to choose the validation
-	// strategy, matching the IsUnifiedHost fallback in ToOAuthArgument().
-	if configType == config.InvalidConfig && cfg.Experimental_IsUnifiedHost && cfg.AccountID != "" {
-		if hasWorkspace {
-			configType = config.WorkspaceConfig
-		} else {
-			configType = config.AccountConfig
-		}
-	}
-
+	configType := auth.ResolveConfigType(cfg)
 	if configType != cfg.ConfigType() {
 		log.Debugf(ctx, "Profile %q: overrode config type from %s to %s (SPOG host)", c.Name, cfg.ConfigType(), configType)
 	}

libs/auth/arguments.go

@@ -49,35 +49,27 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) {
 		Loaders: []config.Loader{config.ConfigAttributes},
 	}
 
-	discoveryURL := a.DiscoveryURL
-	if discoveryURL == "" {
-		// No cached discovery, resolve fresh.
-		if err := cfg.EnsureResolved(); err == nil {
-			discoveryURL = cfg.DiscoveryURL
-		}
+	if a.DiscoveryURL != "" {
+		cfg.DiscoveryURL = a.DiscoveryURL
+	} else {
+		// EnsureResolved populates cfg.DiscoveryURL from .well-known.
+		_ = cfg.EnsureResolved()
 	}
 
 	host := cfg.CanonicalHostName()
 
-	// Classic accounts.* hosts always use account OAuth, even if discovery
-	// returned data. SPOG/unified hosts are handled below via discovery or
-	// the IsUnifiedHost flag.
+	// Classic accounts.* hosts always use account OAuth.
 	if strings.HasPrefix(host, "https://accounts.") || strings.HasPrefix(host, "https://accounts-dod.") {
 		return u2m.NewProfileAccountOAuthArgument(host, cfg.AccountID, a.Profile)
 	}
 
-	// Route based on discovery data: a non-accounts host with an account-scoped
-	// OIDC endpoint is a SPOG/unified host. We check a.AccountID (the caller-
-	// provided value) rather than cfg.AccountID to avoid env var contamination
-	// (e.g. DATABRICKS_ACCOUNT_ID set in the environment). We also require the
-	// DiscoveryURL to contain "/oidc/accounts/" to distinguish SPOG hosts from
-	// classic workspace hosts that may also return discovery metadata.
-	if a.AccountID != "" && discoveryURL != "" && strings.Contains(discoveryURL, "/oidc/accounts/") {
-		return u2m.NewProfileUnifiedOAuthArgument(host, cfg.AccountID, a.Profile)
-	}
-
-	// Legacy backward compat: existing profiles with IsUnifiedHost flag.
-	if a.IsUnifiedHost && a.AccountID != "" {
+	// Pass a.AccountID (not cfg.AccountID) because EnsureResolved can
+	// back-fill cfg.AccountID from two sources: the DATABRICKS_ACCOUNT_ID
+	// env var (via ConfigAttributes) and .well-known/databricks-config
+	// discovery (which returns account_id for every host since PR #4809).
+	// Using cfg.AccountID would cause IsSPOG to misroute plain workspace
+	// hosts as SPOG simply because their metadata includes an account_id.
+	if IsSPOG(cfg, a.AccountID) {
 		return u2m.NewProfileUnifiedOAuthArgument(host, cfg.AccountID, a.Profile)
 	}
 

libs/auth/config_type.go

@@ -0,0 +1,51 @@
+package auth
+
+import (
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go/config"
+)
+
+// IsSPOG returns true if the config represents a SPOG (Single Pane of Glass)
+// host with account-scoped OIDC. Detection is based on:
+//  1. The resolved DiscoveryURL containing /oidc/accounts/ (from .well-known).
+//  2. The Experimental_IsUnifiedHost flag as a legacy fallback.
+//
+// The accountID parameter is separate from cfg.AccountID so that callers can
+// control the source: ResolveConfigType passes cfg.AccountID (from config file),
+// while ToOAuthArgument passes the caller-provided value to avoid env var
+// contamination (DATABRICKS_ACCOUNT_ID or .well-known back-fill).
+func IsSPOG(cfg *config.Config, accountID string) bool {
+	if accountID == "" {
+		return false
+	}
+	if cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/") {
+		return true
+	}
+	return cfg.Experimental_IsUnifiedHost
+}
+
+// ResolveConfigType determines the effective ConfigType for a resolved config.
+// The SDK's ConfigType() classifies based on the host URL prefix alone, which
+// misclassifies SPOG hosts (they don't match the accounts.* prefix). This
+// function additionally uses IsSPOG to detect SPOG hosts.
+//
+// The cfg must already be resolved (via EnsureResolved) before calling this.
+func ResolveConfigType(cfg *config.Config) config.ConfigType {
+	configType := cfg.ConfigType()
+	if configType == config.AccountConfig {
+		return configType
+	}
+
+	if !IsSPOG(cfg, cfg.AccountID) {
+		return configType
+	}
+
+	// The WorkspaceConfig return is a no-op when configType is already
+	// WorkspaceConfig, but is needed for InvalidConfig (legacy IsUnifiedHost
+	// profiles where the SDK dropped the UnifiedHost case in v0.126.0).
+	if cfg.WorkspaceID != "" && cfg.WorkspaceID != WorkspaceIDNone {
+		return config.WorkspaceConfig
+	}
+	return config.AccountConfig
+}

libs/auth/config_type_test.go

@@ -0,0 +1,104 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestResolveConfigType(t *testing.T) {
+	cases := []struct {
+		name string
+		cfg  *config.Config
+		want config.ConfigType
+	}{
+		{
+			name: "classic accounts host stays AccountConfig",
+			cfg: &config.Config{
+				Host:      "https://accounts.cloud.databricks.com",
+				AccountID: "acct-123",
+			},
+			want: config.AccountConfig,
+		},
+		{
+			name: "SPOG account-scoped OIDC without workspace routes to AccountConfig",
+			cfg: &config.Config{
+				Host:         "https://spog.databricks.com",
+				AccountID:    "acct-123",
+				DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+			},
+			want: config.AccountConfig,
+		},
+		{
+			name: "SPOG account-scoped OIDC with workspace routes to WorkspaceConfig",
+			cfg: &config.Config{
+				Host:         "https://spog.databricks.com",
+				AccountID:    "acct-123",
+				WorkspaceID:  "ws-456",
+				DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+			},
+			want: config.WorkspaceConfig,
+		},
+		{
+			name: "SPOG account-scoped OIDC with workspace_id=none routes to AccountConfig",
+			cfg: &config.Config{
+				Host:         "https://spog.databricks.com",
+				AccountID:    "acct-123",
+				WorkspaceID:  "none",
+				DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+			},
+			want: config.AccountConfig,
+		},
+		{
+			name: "workspace-scoped OIDC with account_id stays WorkspaceConfig",
+			cfg: &config.Config{
+				Host:         "https://workspace.databricks.com",
+				AccountID:    "acct-123",
+				DiscoveryURL: "https://workspace.databricks.com/oidc/.well-known/oauth-authorization-server",
+			},
+			want: config.WorkspaceConfig,
+		},
+		{
+			name: "IsUnifiedHost fallback without discovery routes to AccountConfig",
+			cfg: &config.Config{
+				Host:                       "https://spog.databricks.com",
+				AccountID:                  "acct-123",
+				Experimental_IsUnifiedHost: true,
+			},
+			want: config.AccountConfig,
+		},
+		{
+			name: "IsUnifiedHost fallback with workspace routes to WorkspaceConfig",
+			cfg: &config.Config{
+				Host:                       "https://spog.databricks.com",
+				AccountID:                  "acct-123",
+				WorkspaceID:                "ws-456",
+				Experimental_IsUnifiedHost: true,
+			},
+			want: config.WorkspaceConfig,
+		},
+		{
+			name: "no discovery and no IsUnifiedHost stays WorkspaceConfig",
+			cfg: &config.Config{
+				Host:      "https://workspace.databricks.com",
+				AccountID: "acct-123",
+			},
+			want: config.WorkspaceConfig,
+		},
+		{
+			name: "plain workspace without account_id",
+			cfg: &config.Config{
+				Host: "https://workspace.databricks.com",
+			},
+			want: config.WorkspaceConfig,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := ResolveConfigType(tc.cfg)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}