From b91565bb42102f961397b1de53f4e601b0d3dae6 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Mon, 13 Apr 2026 11:11:36 +0000 Subject: [PATCH 1/4] Extract IsSPOG and ResolveConfigType into libs/auth Deduplicate SPOG detection logic between profiles.go and ToOAuthArgument. IsSPOG checks DiscoveryURL and IsUnifiedHost, ResolveConfigType wraps ConfigType with SPOG-aware overrides. --- cmd/auth/profiles.go | 34 +---------- libs/auth/arguments.go | 29 +++------- libs/auth/config_type.go | 48 ++++++++++++++++ libs/auth/config_type_test.go | 104 ++++++++++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 54 deletions(-) create mode 100644 libs/auth/config_type.go create mode 100644 libs/auth/config_type_test.go diff --git a/cmd/auth/profiles.go b/cmd/auth/profiles.go index e2f48d2086..51c397a9ea 100644 --- a/cmd/auth/profiles.go +++ b/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) } diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index 595181b89e..782fd08ee2 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -49,35 +49,22 @@ 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 if err := cfg.EnsureResolved(); err == nil { + // EnsureResolved populates cfg.DiscoveryURL from .well-known. } 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) to avoid env var / discovery + // back-fill from triggering SPOG routing for plain workspace hosts. + if IsSPOG(cfg, a.AccountID) { return u2m.NewProfileUnifiedOAuthArgument(host, cfg.AccountID, a.Profile) } diff --git a/libs/auth/config_type.go b/libs/auth/config_type.go new file mode 100644 index 0000000000..637b1cb7fb --- /dev/null +++ b/libs/auth/config_type.go @@ -0,0 +1,48 @@ +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 + } + + if cfg.WorkspaceID != "" && cfg.WorkspaceID != WorkspaceIDNone { + return config.WorkspaceConfig + } + return config.AccountConfig +} diff --git a/libs/auth/config_type_test.go b/libs/auth/config_type_test.go new file mode 100644 index 0000000000..0ce3b6d410 --- /dev/null +++ b/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) + }) + } +} From fecc25a8367fd33612eea2c295581a2cd87bac7a Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Mon, 13 Apr 2026 11:16:05 +0000 Subject: [PATCH 2/4] Expand comment on accountID override in ToOAuthArgument --- libs/auth/arguments.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index 782fd08ee2..05fcdd7ae0 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -62,8 +62,12 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) { return u2m.NewProfileAccountOAuthArgument(host, cfg.AccountID, a.Profile) } - // Pass a.AccountID (not cfg.AccountID) to avoid env var / discovery - // back-fill from triggering SPOG routing for plain workspace hosts. + // 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) } From ecd4ae5e9c227190948cb2d760b920f13e55ec09 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Tue, 14 Apr 2026 15:09:59 +0000 Subject: [PATCH 3/4] Add comment explaining hasWorkspace branch in ResolveConfigType --- libs/auth/config_type.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libs/auth/config_type.go b/libs/auth/config_type.go index 637b1cb7fb..520b6864cd 100644 --- a/libs/auth/config_type.go +++ b/libs/auth/config_type.go @@ -41,6 +41,9 @@ func ResolveConfigType(cfg *config.Config) config.ConfigType { 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 } From 27dbd852d61e2d8108e4d126b102b68ea1e80584 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Tue, 14 Apr 2026 15:16:43 +0000 Subject: [PATCH 4/4] Simplify EnsureResolved call pattern in ToOAuthArgument --- libs/auth/arguments.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index 05fcdd7ae0..4f724cc801 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -51,8 +51,9 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) { if a.DiscoveryURL != "" { cfg.DiscoveryURL = a.DiscoveryURL - } else if err := cfg.EnsureResolved(); err == nil { + } else { // EnsureResolved populates cfg.DiscoveryURL from .well-known. + _ = cfg.EnsureResolved() } host := cfg.CanonicalHostName()