Skip to content

Commit 06f4c01

Browse files
mihaimitrea-dbmihaimitrea-db
authored
Extract IsSPOG and ResolveConfigType into libs/auth (#4939)
## Summary - Extract the SPOG detection heuristic into a shared `IsSPOG(cfg, accountID)` predicate in `libs/auth/config_type.go`, replacing inline checks in both `profiles.go` and `ToOAuthArgument`. - Extract `ResolveConfigType(cfg)` that wraps `ConfigType()` with SPOG-aware overrides, used by `profiles.go`. - `ToOAuthArgument` now calls `IsSPOG` instead of duplicating the `DiscoveryURL` + `IsUnifiedHost` checks inline. Follow-up to #4929 as suggested in review comment #2 (logic duplication). ### Note on IsUnifiedHost fallback scope The `IsSPOG` predicate checks `Experimental_IsUnifiedHost` unconditionally (not gated on `configType == InvalidConfig`). This is broader than the previous inline check in `profiles.go`, which only fired for `InvalidConfig`. This is safe because the SDK currently returns `InvalidConfig` for all unified hosts (the `UnifiedHost` case was removed from `ConfigType()` in v0.126.0). It is also more robust against future SDK changes that might reclassify unified hosts differently. ## Test plan - [x] `TestResolveConfigType` — 9-case table-driven unit test in `libs/auth/config_type_test.go`. - [x] All existing `TestProfileLoad*` and `TestToOAuthArgument*` tests pass. - [x] `go test ./libs/auth/ ./cmd/auth/` passes.
1 parent f86c804 commit 06f4c01

4 files changed

Lines changed: 169 additions & 54 deletions

File tree

‎cmd/auth/profiles.go‎

Lines changed: 1 addition & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"errors"
66
"fmt"
77
"io/fs"
8-
"strings"
98
"sync"
109
"time"
1110

@@ -58,38 +57,7 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
5857
return
5958
}
6059

61-
// ConfigType() classifies based on the host URL prefix (accounts.* →
62-
// AccountConfig, everything else → WorkspaceConfig). SPOG hosts don't
63-
// match the accounts.* prefix so they're misclassified as WorkspaceConfig.
64-
// Use the resolved DiscoveryURL (from .well-known/databricks-config) to
65-
// detect SPOG hosts with account-scoped OIDC, matching the routing logic
66-
// in auth.AuthArguments.ToOAuthArgument().
67-
configType := cfg.ConfigType()
68-
hasWorkspace := cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone
69-
70-
isAccountScopedOIDC := cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/")
71-
if configType != config.AccountConfig && cfg.AccountID != "" && isAccountScopedOIDC {
72-
if hasWorkspace {
73-
configType = config.WorkspaceConfig
74-
} else {
75-
configType = config.AccountConfig
76-
}
77-
}
78-
79-
// Legacy backward compat: SDK v0.126.0 removed the UnifiedHost case from
80-
// ConfigType(), so profiles with Experimental_IsUnifiedHost now get
81-
// InvalidConfig instead of being routed to account/workspace validation.
82-
// When .well-known is also unreachable (DiscoveryURL empty), the override
83-
// above can't help. Fall back to workspace_id to choose the validation
84-
// strategy, matching the IsUnifiedHost fallback in ToOAuthArgument().
85-
if configType == config.InvalidConfig && cfg.Experimental_IsUnifiedHost && cfg.AccountID != "" {
86-
if hasWorkspace {
87-
configType = config.WorkspaceConfig
88-
} else {
89-
configType = config.AccountConfig
90-
}
91-
}
92-
60+
configType := auth.ResolveConfigType(cfg)
9361
if configType != cfg.ConfigType() {
9462
log.Debugf(ctx, "Profile %q: overrode config type from %s to %s (SPOG host)", c.Name, cfg.ConfigType(), configType)
9563
}

‎libs/auth/arguments.go‎

Lines changed: 13 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -49,35 +49,27 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) {
4949
Loaders: []config.Loader{config.ConfigAttributes},
5050
}
5151

52-
discoveryURL := a.DiscoveryURL
53-
if discoveryURL == "" {
54-
// No cached discovery, resolve fresh.
55-
if err := cfg.EnsureResolved(); err == nil {
56-
discoveryURL = cfg.DiscoveryURL
57-
}
52+
if a.DiscoveryURL != "" {
53+
cfg.DiscoveryURL = a.DiscoveryURL
54+
} else {
55+
// EnsureResolved populates cfg.DiscoveryURL from .well-known.
56+
_ = cfg.EnsureResolved()
5857
}
5958

6059
host := cfg.CanonicalHostName()
6160

62-
// Classic accounts.* hosts always use account OAuth, even if discovery
63-
// returned data. SPOG/unified hosts are handled below via discovery or
64-
// the IsUnifiedHost flag.
61+
// Classic accounts.* hosts always use account OAuth.
6562
if strings.HasPrefix(host, "https://accounts.") || strings.HasPrefix(host, "https://accounts-dod.") {
6663
return u2m.NewProfileAccountOAuthArgument(host, cfg.AccountID, a.Profile)
6764
}
6865

69-
// Route based on discovery data: a non-accounts host with an account-scoped
70-
// OIDC endpoint is a SPOG/unified host. We check a.AccountID (the caller-
71-
// provided value) rather than cfg.AccountID to avoid env var contamination
72-
// (e.g. DATABRICKS_ACCOUNT_ID set in the environment). We also require the
73-
// DiscoveryURL to contain "/oidc/accounts/" to distinguish SPOG hosts from
74-
// classic workspace hosts that may also return discovery metadata.
75-
if a.AccountID != "" && discoveryURL != "" && strings.Contains(discoveryURL, "/oidc/accounts/") {
76-
return u2m.NewProfileUnifiedOAuthArgument(host, cfg.AccountID, a.Profile)
77-
}
78-
79-
// Legacy backward compat: existing profiles with IsUnifiedHost flag.
80-
if a.IsUnifiedHost && a.AccountID != "" {
66+
// Pass a.AccountID (not cfg.AccountID) because EnsureResolved can
67+
// back-fill cfg.AccountID from two sources: the DATABRICKS_ACCOUNT_ID
68+
// env var (via ConfigAttributes) and .well-known/databricks-config
69+
// discovery (which returns account_id for every host since PR #4809).
70+
// Using cfg.AccountID would cause IsSPOG to misroute plain workspace
71+
// hosts as SPOG simply because their metadata includes an account_id.
72+
if IsSPOG(cfg, a.AccountID) {
8173
return u2m.NewProfileUnifiedOAuthArgument(host, cfg.AccountID, a.Profile)
8274
}
8375

‎libs/auth/config_type.go‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package auth
2+
3+
import (
4+
"strings"
5+
6+
"github.com/databricks/databricks-sdk-go/config"
7+
)
8+
9+
// IsSPOG returns true if the config represents a SPOG (Single Pane of Glass)
10+
// host with account-scoped OIDC. Detection is based on:
11+
// 1. The resolved DiscoveryURL containing /oidc/accounts/ (from .well-known).
12+
// 2. The Experimental_IsUnifiedHost flag as a legacy fallback.
13+
//
14+
// The accountID parameter is separate from cfg.AccountID so that callers can
15+
// control the source: ResolveConfigType passes cfg.AccountID (from config file),
16+
// while ToOAuthArgument passes the caller-provided value to avoid env var
17+
// contamination (DATABRICKS_ACCOUNT_ID or .well-known back-fill).
18+
func IsSPOG(cfg *config.Config, accountID string) bool {
19+
if accountID == "" {
20+
return false
21+
}
22+
if cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/") {
23+
return true
24+
}
25+
return cfg.Experimental_IsUnifiedHost
26+
}
27+
28+
// ResolveConfigType determines the effective ConfigType for a resolved config.
29+
// The SDK's ConfigType() classifies based on the host URL prefix alone, which
30+
// misclassifies SPOG hosts (they don't match the accounts.* prefix). This
31+
// function additionally uses IsSPOG to detect SPOG hosts.
32+
//
33+
// The cfg must already be resolved (via EnsureResolved) before calling this.
34+
func ResolveConfigType(cfg *config.Config) config.ConfigType {
35+
configType := cfg.ConfigType()
36+
if configType == config.AccountConfig {
37+
return configType
38+
}
39+
40+
if !IsSPOG(cfg, cfg.AccountID) {
41+
return configType
42+
}
43+
44+
// The WorkspaceConfig return is a no-op when configType is already
45+
// WorkspaceConfig, but is needed for InvalidConfig (legacy IsUnifiedHost
46+
// profiles where the SDK dropped the UnifiedHost case in v0.126.0).
47+
if cfg.WorkspaceID != "" && cfg.WorkspaceID != WorkspaceIDNone {
48+
return config.WorkspaceConfig
49+
}
50+
return config.AccountConfig
51+
}

‎libs/auth/config_type_test.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
package auth
2+
3+
import (
4+
"testing"
5+
6+
"github.com/databricks/databricks-sdk-go/config"
7+
"github.com/stretchr/testify/assert"
8+
)
9+
10+
func TestResolveConfigType(t *testing.T) {
11+
cases := []struct {
12+
name string
13+
cfg *config.Config
14+
want config.ConfigType
15+
}{
16+
{
17+
name: "classic accounts host stays AccountConfig",
18+
cfg: &config.Config{
19+
Host: "https://accounts.cloud.databricks.com",
20+
AccountID: "acct-123",
21+
},
22+
want: config.AccountConfig,
23+
},
24+
{
25+
name: "SPOG account-scoped OIDC without workspace routes to AccountConfig",
26+
cfg: &config.Config{
27+
Host: "https://spog.databricks.com",
28+
AccountID: "acct-123",
29+
DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
30+
},
31+
want: config.AccountConfig,
32+
},
33+
{
34+
name: "SPOG account-scoped OIDC with workspace routes to WorkspaceConfig",
35+
cfg: &config.Config{
36+
Host: "https://spog.databricks.com",
37+
AccountID: "acct-123",
38+
WorkspaceID: "ws-456",
39+
DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
40+
},
41+
want: config.WorkspaceConfig,
42+
},
43+
{
44+
name: "SPOG account-scoped OIDC with workspace_id=none routes to AccountConfig",
45+
cfg: &config.Config{
46+
Host: "https://spog.databricks.com",
47+
AccountID: "acct-123",
48+
WorkspaceID: "none",
49+
DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server",
50+
},
51+
want: config.AccountConfig,
52+
},
53+
{
54+
name: "workspace-scoped OIDC with account_id stays WorkspaceConfig",
55+
cfg: &config.Config{
56+
Host: "https://workspace.databricks.com",
57+
AccountID: "acct-123",
58+
DiscoveryURL: "https://workspace.databricks.com/oidc/.well-known/oauth-authorization-server",
59+
},
60+
want: config.WorkspaceConfig,
61+
},
62+
{
63+
name: "IsUnifiedHost fallback without discovery routes to AccountConfig",
64+
cfg: &config.Config{
65+
Host: "https://spog.databricks.com",
66+
AccountID: "acct-123",
67+
Experimental_IsUnifiedHost: true,
68+
},
69+
want: config.AccountConfig,
70+
},
71+
{
72+
name: "IsUnifiedHost fallback with workspace routes to WorkspaceConfig",
73+
cfg: &config.Config{
74+
Host: "https://spog.databricks.com",
75+
AccountID: "acct-123",
76+
WorkspaceID: "ws-456",
77+
Experimental_IsUnifiedHost: true,
78+
},
79+
want: config.WorkspaceConfig,
80+
},
81+
{
82+
name: "no discovery and no IsUnifiedHost stays WorkspaceConfig",
83+
cfg: &config.Config{
84+
Host: "https://workspace.databricks.com",
85+
AccountID: "acct-123",
86+
},
87+
want: config.WorkspaceConfig,
88+
},
89+
{
90+
name: "plain workspace without account_id",
91+
cfg: &config.Config{
92+
Host: "https://workspace.databricks.com",
93+
},
94+
want: config.WorkspaceConfig,
95+
},
96+
}
97+
98+
for _, tc := range cases {
99+
t.Run(tc.name, func(t *testing.T) {
100+
got := ResolveConfigType(tc.cfg)
101+
assert.Equal(t, tc.want, got)
102+
})
103+
}
104+
}

0 commit comments

Comments
 (0)