Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions acceptance/cmd/auth/profiles/spog-account/out.test.toml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

16 changes: 16 additions & 0 deletions acceptance/cmd/auth/profiles/spog-account/output.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@

=== SPOG account profile should be valid
>>> [CLI] auth profiles --output json
{
"profiles": [
{
"name":"spog-account",
"host":"[DATABRICKS_URL]",
"account_id":"spog-acct-123",
"workspace_id":"none",
"cloud":"aws",
"auth_type":"pat",
"valid":true
}
]
}
15 changes: 15 additions & 0 deletions acceptance/cmd/auth/profiles/spog-account/script
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
sethome "./home"

# Create a SPOG account profile: non-accounts.* host with account_id, no workspace_id.
# Before the fix, this was misclassified as WorkspaceConfig and validated with
# CurrentUser.Me, which fails on account-scoped SPOG hosts.
cat > "./home/.databrickscfg" <<EOF
[spog-account]
host = ${DATABRICKS_HOST}
account_id = spog-acct-123
workspace_id = none
token = test-token
EOF

title "SPOG account profile should be valid"
trace $CLI auth profiles --output json
20 changes: 20 additions & 0 deletions acceptance/cmd/auth/profiles/spog-account/test.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
Ignore = [
"home"
]

# Override .well-known to return account-scoped OIDC, simulating a SPOG host.
# The oidc_endpoint contains /oidc/accounts/ which signals account-scoped OIDC.
[[Server]]
Pattern = "GET /.well-known/databricks-config"
Response.Body = '''
{
"account_id": "spog-acct-123",
"oidc_endpoint": "localhost/oidc/accounts/spog-acct-123"
}
'''

# Provide a handler for the account workspaces list endpoint used by
# auth profiles validation for account-type configs.
[[Server]]
Pattern = "GET /api/2.0/accounts/spog-acct-123/workspaces"
Response.Body = '[]'
27 changes: 26 additions & 1 deletion cmd/auth/profiles.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,11 @@ import (
"errors"
"fmt"
"io/fs"
"strings"
"sync"
"time"

"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/databrickscfg"
"github.com/databricks/cli/libs/databrickscfg/profile"
Expand Down Expand Up @@ -56,7 +58,30 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
return
}

switch cfg.ConfigType() {
// 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()
isAccountScopedOIDC := cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/")
if configType != config.AccountConfig && cfg.AccountID != "" && isAccountScopedOIDC {
if cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone {
configType = config.WorkspaceConfig
} else {
configType = config.AccountConfig
}
}
Comment on lines +67 to +77
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (follow-up): The SPOG detection heuristic (DiscoveryURL contains /oidc/accounts/ + IsUnifiedHost fallback) is now duplicated between here and libs/auth/arguments.go:69-82. The cross-referencing comments help, but if the routing logic changes in one place the other can drift.

Consider extracting a shared helper in libs/auth/ in a follow-up, e.g.:

func ResolveConfigType(cfg *config.Config) config.ConfigType { ... }

Not blocking for this PR since the logic is simple and well-documented.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be good to have. Bonus points that it moves some of the logic form the SDK to the CLI.


// Legacy backward compat: profiles with Experimental_IsUnifiedHost where
// .well-known is unreachable (so DiscoveryURL is empty). Matches the
// fallback in auth.AuthArguments.ToOAuthArgument().
if configType == config.InvalidConfig && cfg.Experimental_IsUnifiedHost && cfg.AccountID != "" {
configType = config.AccountConfig
}
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: This legacy fallback always forces AccountConfig regardless of workspace_id. In arguments.go:80-81, ignoring workspace_id is intentional because unified OAuth routing is always account-scoped. But here you're choosing a validation strategy, not an OAuth flow. MatchWorkspaceProfiles in profiler.go treats any non-none workspace_id as workspace-scoped.

If a legacy experimental_is_unified_host profile has a real workspace_id and .well-known is unreachable, this would validate it with Workspaces.List instead of CurrentUser.Me. That said, the combination of Experimental_IsUnifiedHost + real workspace_id is unlikely in practice since workspace-scoped SPOG came after the legacy flag era.

Consider either:

  • Adding the same workspace_id branching here as in the discovery block above, or
  • Adding a comment explaining why it's intentionally omitted


switch configType {
case config.AccountConfig:
a, err := databricks.NewAccountClient((*databricks.Config)(cfg))
if err != nil {
Expand Down
186 changes: 186 additions & 0 deletions cmd/auth/profiles_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
package auth

import (
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"runtime"
"testing"
Expand Down Expand Up @@ -74,3 +78,185 @@ func TestProfilesDefaultMarker(t *testing.T) {
require.NoError(t, err)
assert.Equal(t, "profile-a", defaultProfile)
}

// newProfileTestServer creates a mock server for profile validation tests.
// It serves /.well-known/databricks-config with the given OIDC shape and
// responds to the workspace/account validation API endpoints.
func newProfileTestServer(t *testing.T, accountScoped bool, accountID string) *httptest.Server {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitpick: Both the workspace and account validation endpoints (/api/2.0/preview/scim/v2/Me and /api/2.0/accounts/{id}/workspaces) return success, so Valid == true doesn't prove which branch was taken. Consider making the "wrong" endpoint return an error for each test case to prove routing correctness.

suggestion: Also consider adding a negative test case where .well-known returns 404, account_id is set, Experimental_IsUnifiedHost is false, to verify the profile stays as WorkspaceConfig (proving the override doesn't trigger accidentally).

t.Helper()
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":
oidcEndpoint := r.Host + "/oidc"
if accountScoped {
oidcEndpoint = r.Host + "/oidc/accounts/" + accountID
}
_ = json.NewEncoder(w).Encode(map[string]any{
"account_id": accountID,
"oidc_endpoint": oidcEndpoint,
})
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":
_ = json.NewEncoder(w).Encode([]map[string]any{})
default:
w.WriteHeader(http.StatusNotFound)
}
}))
t.Cleanup(server.Close)
return server
}

func TestProfileLoadSPOGConfigType(t *testing.T) {
spogServer := newProfileTestServer(t, true, "spog-acct")
wsServer := newProfileTestServer(t, false, "ws-acct")

cases := []struct {
name string
host string
accountID string
workspaceID string
wantValid bool
}{
{
name: "SPOG account profile validated as account",
host: spogServer.URL,
accountID: "spog-acct",
wantValid: true,
},
{
name: "SPOG workspace profile validated as workspace",
host: spogServer.URL,
accountID: "spog-acct",
workspaceID: "ws-123",
wantValid: true,
},
{
name: "SPOG profile with workspace_id=none validated as account",
host: spogServer.URL,
accountID: "spog-acct",
workspaceID: "none",
wantValid: true,
},
{
name: "classic workspace with account_id from discovery stays workspace",
host: wsServer.URL,
accountID: "ws-acct",
wantValid: true,
},
}

for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
dir := t.TempDir()
configFile := filepath.Join(dir, ".databrickscfg")
t.Setenv("HOME", dir)
if runtime.GOOS == "windows" {
t.Setenv("USERPROFILE", dir)
}

content := "[test-profile]\nhost = " + tc.host + "\ntoken = test-token\n"
if tc.accountID != "" {
content += "account_id = " + tc.accountID + "\n"
}
if tc.workspaceID != "" {
content += "workspace_id = " + tc.workspaceID + "\n"
}
require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600))

p := &profileMetadata{
Name: "test-profile",
Host: tc.host,
AccountID: tc.accountID,
}
p.Load(t.Context(), configFile, false)

assert.Equal(t, tc.wantValid, p.Valid, "Valid mismatch")
assert.NotEmpty(t, p.Host, "Host should be set")
assert.NotEmpty(t, p.AuthType, "AuthType should be set")
})
}
}

func TestProfileLoadUnifiedHostFallback(t *testing.T) {
// When Experimental_IsUnifiedHost is set but .well-known is unreachable,
// ConfigType() returns InvalidConfig. The fallback should reclassify as
// AccountConfig so the profile is still validated.
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":
w.WriteHeader(http.StatusNotFound)
case "/api/2.0/accounts/unified-acct/workspaces":
_ = json.NewEncoder(w).Encode([]map[string]any{})
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 := "[unified-profile]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = unified-acct\nexperimental_is_unified_host = true\n"
require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600))

p := &profileMetadata{
Name: "unified-profile",
Host: server.URL,
AccountID: "unified-acct",
}
p.Load(t.Context(), configFile, false)

assert.True(t, p.Valid, "unified host profile should be valid via fallback")
assert.NotEmpty(t, p.Host)
assert.NotEmpty(t, p.AuthType)
}

func TestProfileLoadClassicAccountHost(t *testing.T) {
// Verify that a host with account-scoped OIDC from discovery is validated
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: This test name says "classic account host" but the httptest.Server URL (e.g. http://127.0.0.1:PORT) doesn't match the accounts.* prefix, so ConfigType() won't return AccountConfig for the classic-host reason. This actually exercises the SPOG discovery override path (same as TestProfileLoadSPOGConfigType).

Consider renaming to something like TestProfileLoadSPOGAccountWithDiscovery, or updating the comment to acknowledge that classic accounts.* behavior can't easily be tested with httptest and this serves as a supplementary SPOG case.

// as an account config (via Workspaces.List, not CurrentUser.Me).
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": "classic-acct",
"oidc_endpoint": r.Host + "/oidc/accounts/classic-acct",
})
case "/api/2.0/accounts/classic-acct/workspaces":
_ = json.NewEncoder(w).Encode([]map[string]any{})
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 := "[acct-profile]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = classic-acct\n"
require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600))

p := &profileMetadata{
Name: "acct-profile",
Host: server.URL,
AccountID: "classic-acct",
}
p.Load(t.Context(), configFile, false)

assert.True(t, p.Valid, "classic account profile should be valid")
assert.NotEmpty(t, p.Host)
assert.NotEmpty(t, p.AuthType)
}
Loading