Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
}
}

// 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
}

switch configType {
case config.AccountConfig:
a, err := databricks.NewAccountClient((*databricks.Config)(cfg))
if err != nil {
Expand Down
190 changes: 190 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,189 @@ 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 {
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
wantConfigCloud string
}{
{
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) {
// Classic accounts.* hosts are already correctly classified as AccountConfig
// by ConfigType(). Verify that behavior is preserved.
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)
}

// Use the test server URL but override the host to look like an accounts host.
// Since we can't actually use accounts.cloud.databricks.com in tests, we verify
// indirectly: a SPOG profile without workspace_id should be validated as account.
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