auth profiles: probe both surfaces and OR the results

Replace ResolveConfigType's "either account or workspace" routing with a
permissive validator: probe whichever API surfaces the profile has a
signal for (host shape or field presence), and mark the profile valid if
any probe succeeds. This addresses review feedback on the previous fix —
it now also handles SPOG workspace-scoped credentials (e.g. a PAT),
which the strict "SPOG always validates as account" rule would have
falsely flagged as invalid.

Probes run sequentially against the shared cfg because the SDK's lazy
Authenticate() chain writes cfg.Host (via fixHostIfNeeded) while the
other client's construction reads cfg.Host unlocked — go test -race
flags the parallel version. Profile-level parallelism in
newProfilesCommand is unchanged, so overall auth profiles wall-clock is
still bounded by the slowest profile.

ResolveConfigType is removed; IsSPOGHost and IsClassicWorkspaceHost are
added in libs/auth as named wrappers so the host-shape check in
profiles.go reads as a flat block of named booleans.

Tests:
- TestProfileLoadSPOGWorkspaceCredential covers PAT-on-SPOG explicitly:
  workspace probe succeeds, account probe 403s, OR yields Valid=true.
- TestIsSPOGHost and TestIsClassicWorkspaceHost cover the three-way
  host classification, including the accounts-dod.* variant.

Co-authored-by: Isaac

Author: mihaimitrea-db

cmd/auth/profiles.go

@@ -57,40 +57,51 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 		return
 	}
 
-	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)
-	}
+	// Validate by probing the API surfaces this profile has a signal for.
+	// Each signal — host shape or field presence — enables its corresponding
+	// probe, and the OR of the probe results is the verdict.
+
+	// Host signals.
+	// isAccountHost:   classic accounts.* host.
+	// isSPOGHost:      unified host with account-scoped OIDC discovery.
+	// isWorkspaceHost: classic workspace host (neither of the above).
+	isAccountHost := auth.IsClassicAccountHost(cfg.CanonicalHostName())
+	isSPOGHost := auth.IsSPOGHost(cfg)
+	isWorkspaceHost := auth.IsClassicWorkspaceHost(cfg)
+
+	// Field signals.
+	// hasAccountID:       account_id is set (from file, env, or discovery back-fill).
+	// hasRealWorkspaceID: workspace_id is set to a real value.
+	hasAccountID := cfg.AccountID != ""
+	// workspace_id is "" when not present in the profile, "none" when the user picked Skip during SPOG login.
+	hasRealWorkspaceID := cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone
+
+	tryAccount := isAccountHost || isSPOGHost || hasAccountID
+	tryWorkspace := isWorkspaceHost || hasRealWorkspaceID
 
-	switch configType {
-	case config.AccountConfig:
+	var accountOK, workspaceOK bool
+	if tryAccount {
 		a, err := databricks.NewAccountClient((*databricks.Config)(cfg))
-		if err != nil {
-			return
-		}
-		_, err = a.Workspaces.List(ctx)
-		c.Host = cfg.Host
-		c.AuthType = cfg.AuthType
-		if err != nil {
-			return
+		if err == nil {
+			if _, err := a.Workspaces.List(ctx); err == nil {
+				accountOK = true
+			}
 		}
-		c.Valid = true
-	case config.WorkspaceConfig:
+	}
+	if tryWorkspace {
 		w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg))
-		if err != nil {
-			return
-		}
-		_, err = w.CurrentUser.Me(ctx)
-		c.Host = cfg.Host
-		c.AuthType = cfg.AuthType
-		if err != nil {
-			return
+		if err == nil {
+			if _, err := w.CurrentUser.Me(ctx); err == nil {
+				workspaceOK = true
+			}
 		}
-		c.Valid = true
-	case config.InvalidConfig:
-		// Invalid configuration, skip validation
-		return
 	}
+
+	// Capture Host/AuthType after the probes run: SDK Authenticate() sets
+	// cfg.AuthType lazily based on the credentials it actually exercised.
+	c.Host = cfg.Host
+	c.AuthType = cfg.AuthType
+	c.Valid = accountOK || workspaceOK
 }
 
 func newProfilesCommand() *cobra.Command {

cmd/auth/profiles_test.go

@@ -80,10 +80,11 @@ func TestProfilesDefaultMarker(t *testing.T) {
 }
 
 // newSPOGServer creates a mock SPOG server that returns account-scoped OIDC.
-// It serves only the account validation endpoint: every SPOG profile must
-// route through Workspaces.List, regardless of workspace_id (see
-// ResolveConfigType). The workspace endpoint deliberately returns 500 so any
-// regression that routes a SPOG profile to CurrentUser.Me fails the test.
+// The workspace endpoint deliberately returns 500 to mirror real SPOG hosts
+// where account-audience OAuth tokens can't load workspace OAuth config.
+// auth profiles probes both surfaces and accepts either success, so the test
+// passes when Workspaces.List succeeds even though CurrentUser.Me fails —
+// and a regression that drops the account probe surfaces as Valid=false.
 func newSPOGServer(t *testing.T, accountID string) *httptest.Server {
 	t.Helper()
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -107,8 +108,9 @@ func newSPOGServer(t *testing.T, accountID string) *httptest.Server {
 }
 
 // newWorkspaceServer creates a mock workspace server that returns workspace-scoped
-// OIDC and only serves the workspace validation endpoint. The account validation
-// endpoint returns 404 to prove the workspace path was taken.
+// OIDC and a workspace_id in discovery (mirroring real workspace hosts since
+// PR #4809). It serves CurrentUser.Me; the account endpoint returns 404 so a
+// workspace probe is the only path that produces Valid=true.
 func newWorkspaceServer(t *testing.T, accountID string) *httptest.Server {
 	t.Helper()
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -117,6 +119,7 @@ func newWorkspaceServer(t *testing.T, accountID string) *httptest.Server {
 		case "/.well-known/databricks-config":
 			_ = json.NewEncoder(w).Encode(map[string]any{
 				"account_id":    accountID,
+				"workspace_id":  "ws-from-discovery",
 				"oidc_endpoint": r.Host + "/oidc",
 			})
 		case "/api/2.0/preview/scim/v2/Me":
@@ -147,10 +150,10 @@ func TestProfileLoadSPOGConfigType(t *testing.T) {
 			wantValid: true,
 		},
 		{
-			// Regression: this case used to route to CurrentUser.Me. The
-			// SPOG mock now returns 500 on that endpoint, so anything other
-			// than the account path produces wantValid=false.
-			name:        "SPOG workspace profile validated as account",
+			// SPOG with a real workspace_id: workspace probe (CurrentUser.Me)
+			// fails on the mock, account probe (Workspaces.List) succeeds —
+			// the OR makes the profile valid.
+			name:        "SPOG workspace profile valid when account probe succeeds",
 			host:        spogServer.URL,
 			accountID:   "spog-acct",
 			workspaceID: "ws-123",
@@ -203,6 +206,52 @@ func TestProfileLoadSPOGConfigType(t *testing.T) {
 	}
 }
 
+// TestProfileLoadSPOGWorkspaceCredential covers the inverse of the
+// account-OAuth case: a workspace-scoped credential (e.g. a PAT) against a
+// SPOG host. CurrentUser.Me succeeds, Workspaces.List fails (no account-level
+// access). The OR of the two probes must still mark the profile Valid=true.
+func TestProfileLoadSPOGWorkspaceCredential(t *testing.T) {
+	const accountID = "spog-acct"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch r.URL.Path {
+		case "/.well-known/databricks-config":
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"account_id":    accountID,
+				"oidc_endpoint": r.Host + "/oidc/accounts/" + accountID,
+			})
+		case "/api/2.0/preview/scim/v2/Me":
+			_ = json.NewEncoder(w).Encode(map[string]any{"userName": "test-user"})
+		case "/api/2.0/accounts/" + accountID + "/workspaces":
+			http.Error(w, "user lacks account-level access", http.StatusForbidden)
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	t.Cleanup(server.Close)
+
+	dir := t.TempDir()
+	configFile := filepath.Join(dir, ".databrickscfg")
+	t.Setenv("HOME", dir)
+	if runtime.GOOS == "windows" {
+		t.Setenv("USERPROFILE", dir)
+	}
+
+	content := "[ws-cred-on-spog]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = " + accountID + "\nworkspace_id = ws-123\n"
+	require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600))
+
+	p := &profileMetadata{
+		Name:      "ws-cred-on-spog",
+		Host:      server.URL,
+		AccountID: accountID,
+	}
+	p.Load(t.Context(), configFile, false)
+
+	assert.True(t, p.Valid, "workspace probe alone should make the profile valid")
+	assert.NotEmpty(t, p.Host)
+	assert.NotEmpty(t, p.AuthType)
+}
+
 func TestClassicAccountsHostConfigType(t *testing.T) {
 	// Classic accounts.* hosts can't be tested through Load() because httptest
 	// generates 127.0.0.1 URLs. Verify directly that ConfigType() classifies
@@ -219,9 +268,12 @@ func TestClassicAccountsHostConfigType(t *testing.T) {
 }
 
 func TestProfileLoadNoDiscoveryStaysWorkspace(t *testing.T) {
-	// When .well-known returns 404 and the unified-host fallback is false,
-	// the SPOG override should NOT trigger even if account_id is set. The
-	// profile should stay WorkspaceConfig and validate via CurrentUser.Me.
+	// account_id can linger in a profile from a prior account login on the
+	// same profile name (e.g. user logged into accounts.cloud.databricks.com,
+	// logged out, then re-used the profile name for a workspace login). A
+	// stale account_id must not promote the profile to account validation
+	// when the host itself isn't an account/SPOG surface — the workspace
+	// probe is still the right signal.
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 		switch r.URL.Path {

libs/auth/config_type.go

@@ -27,32 +27,30 @@ func HasUnifiedHostSignal(discoveryURL string) bool {
 // of an accountID guard: SPOG routing requires an account ID to construct the
 // OAuth URL, so a nil or empty accountID always returns false.
 //
-// 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).
+// The accountID parameter is separate from cfg.AccountID so that callers
+// (currently ToOAuthArgument) can pass the caller-provided value to avoid
+// env-var contamination (DATABRICKS_ACCOUNT_ID or .well-known back-fill)
+// that would otherwise misroute plain workspace hosts as SPOG.
 func IsSPOG(cfg *config.Config, accountID string) bool {
 	if accountID == "" {
 		return false
 	}
 	return HasUnifiedHostSignal(cfg.DiscoveryURL)
 }
 
-// ResolveConfigType returns the effective ConfigType for a resolved config.
-// The SDK's ConfigType() never returns AccountConfig for SPOG hosts (they
-// don't match the accounts.* prefix and resolve to UnifiedHost). For SPOG
-// the OAuth issuer is account-scoped, so every token's audience is the
-// account and workspace-API validation is unreliable — return AccountConfig
-// regardless of workspace_id.
-//
-// 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 config.AccountConfig
+// IsSPOGHost reports whether cfg points at a unified SPOG host: account-scoped
+// OIDC discovery and NOT a classic accounts.* host. Classic accounts.* hosts
+// share the same OIDC shape, so IsSPOG alone can't tell them apart; layer
+// IsClassicAccountHost on top to disambiguate.
+func IsSPOGHost(cfg *config.Config) bool {
+	if IsClassicAccountHost(cfg.CanonicalHostName()) {
+		return false
 	}
-	return configType
+	return IsSPOG(cfg, cfg.AccountID)
+}
+
+// IsClassicWorkspaceHost reports whether cfg points at a classic workspace
+// host: neither a classic accounts.* host nor a SPOG host.
+func IsClassicWorkspaceHost(cfg *config.Config) bool {
+	return !IsClassicAccountHost(cfg.CanonicalHostName()) && !IsSPOG(cfg, cfg.AccountID)
 }

libs/auth/config_type_test.go

@@ -24,82 +24,77 @@ func TestHasUnifiedHostSignal(t *testing.T) {
 	}
 }
 
-func TestResolveConfigType(t *testing.T) {
+func TestIsSPOG(t *testing.T) {
 	cases := []struct {
-		name string
-		cfg  *config.Config
-		want config.ConfigType
+		name      string
+		cfg       *config.Config
+		accountID string
+		want      bool
 	}{
 		{
-			name: "classic accounts host stays AccountConfig",
-			cfg: &config.Config{
-				Host:      "https://accounts.cloud.databricks.com",
-				AccountID: "acct-123",
-			},
-			want: config.AccountConfig,
+			name:      "account-scoped OIDC with account_id",
+			cfg:       &config.Config{DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server"},
+			accountID: "acct-123",
+			want:      true,
 		},
 		{
-			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:      "account-scoped OIDC without account_id",
+			cfg:       &config.Config{DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server"},
+			accountID: "",
+			want:      false,
 		},
 		{
-			// SPOG OAuth is account-scoped, so even profiles with a real
-			// workspace_id route to AccountConfig — the token's audience is
-			// the account, and workspace-API validation can't authenticate it.
-			name: "SPOG with real workspace_id routes to AccountConfig",
-			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.AccountConfig,
-		},
-		{
-			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: "no discovery 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,
+			name:      "workspace-scoped OIDC with account_id back-filled",
+			cfg:       &config.Config{DiscoveryURL: "https://workspace.databricks.com/oidc/.well-known/oauth-authorization-server"},
+			accountID: "acct-123",
+			want:      false,
 		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got := ResolveConfigType(tc.cfg)
-			assert.Equal(t, tc.want, got)
+			assert.Equal(t, tc.want, IsSPOG(tc.cfg, tc.accountID))
 		})
 	}
 }
+
+// Configs used across the host-classification tests below. The three host
+// shapes are mutually exclusive: exactly one helper returns true per cfg.
+// accounts-dod.* is a second classic accounts variant — same OIDC shape
+// and classification as accounts.*, different URL prefix.
+var (
+	classicAccountCfg = &config.Config{
+		Host:         "https://accounts.cloud.databricks.com",
+		AccountID:    "acct-123",
+		DiscoveryURL: "https://accounts.cloud.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+	}
+	classicAccountDodCfg = &config.Config{
+		Host:         "https://accounts-dod.cloud.databricks.us",
+		AccountID:    "acct-123",
+		DiscoveryURL: "https://accounts-dod.cloud.databricks.us/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+	}
+	spogCfg = &config.Config{
+		Host:         "https://spog.gcp.databricks.com",
+		AccountID:    "acct-123",
+		DiscoveryURL: "https://spog.gcp.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
+	}
+	classicWorkspaceCfg = &config.Config{
+		Host:         "https://dbc-xxxx.cloud.databricks.com",
+		AccountID:    "acct-123",
+		DiscoveryURL: "https://dbc-xxxx.cloud.databricks.com/oidc/.well-known/oauth-authorization-server",
+	}
+)
+
+func TestIsSPOGHost(t *testing.T) {
+	assert.False(t, IsSPOGHost(classicAccountCfg), "classic accounts.* shares the SPOG OIDC shape but is not SPOG")
+	assert.False(t, IsSPOGHost(classicAccountDodCfg), "classic accounts-dod.* shares the SPOG OIDC shape but is not SPOG")
+	assert.True(t, IsSPOGHost(spogCfg))
+	assert.False(t, IsSPOGHost(classicWorkspaceCfg))
+}
+
+func TestIsClassicWorkspaceHost(t *testing.T) {
+	assert.False(t, IsClassicWorkspaceHost(classicAccountCfg))
+	assert.False(t, IsClassicWorkspaceHost(classicAccountDodCfg))
+	assert.False(t, IsClassicWorkspaceHost(spogCfg))
+	assert.True(t, IsClassicWorkspaceHost(classicWorkspaceCfg))
+}