diff --git a/acceptance/cmd/auth/logout/default-profile/out.databrickscfg b/acceptance/cmd/auth/logout/default-profile/out.databrickscfg index 777bad84de..2508cf2868 100644 --- a/acceptance/cmd/auth/logout/default-profile/out.databrickscfg +++ b/acceptance/cmd/auth/logout/default-profile/out.databrickscfg @@ -3,5 +3,5 @@ ; Dev workspace [dev] -host = https://dev.cloud.databricks.com +host = [DATABRICKS_URL] auth_type = databricks-cli diff --git a/acceptance/cmd/auth/logout/default-profile/output.txt b/acceptance/cmd/auth/logout/default-profile/output.txt index d9a9e0f55d..3ccb937438 100644 --- a/acceptance/cmd/auth/logout/default-profile/output.txt +++ b/acceptance/cmd/auth/logout/default-profile/output.txt @@ -2,12 +2,12 @@ === Initial config ; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified. [DEFAULT] -host = https://default.cloud.databricks.com +host = [DATABRICKS_URL] auth_type = databricks-cli ; Dev workspace [dev] -host = https://dev.cloud.databricks.com +host = [DATABRICKS_URL] auth_type = databricks-cli === Delete the DEFAULT profile @@ -23,5 +23,5 @@ OK: Backup file exists ; Dev workspace [dev] -host = https://dev.cloud.databricks.com +host = [DATABRICKS_URL] auth_type = databricks-cli diff --git a/acceptance/cmd/auth/logout/default-profile/script b/acceptance/cmd/auth/logout/default-profile/script index 94aeda31fa..af571ccae8 100644 --- a/acceptance/cmd/auth/logout/default-profile/script +++ b/acceptance/cmd/auth/logout/default-profile/script @@ -1,14 +1,14 @@ sethome "./home" -cat > "./home/.databrickscfg" <<'EOF' +cat > "./home/.databrickscfg" < "./home/.databrickscfg" <<'EOF' +cat > "./home/.databrickscfg" < "./home/.databrickscfg" <<'EOF' +cat > "./home/.databrickscfg" <>> [CLI] auth profiles +Name Host Valid +logfood (Default) [DATABRICKS_URL] YES + +=== Token cache keys before logout +[ + "[DATABRICKS_URL]", + "logfood" +] + +=== Logout without --delete +>>> [CLI] auth logout --profile logfood --force +Logged out of profile "logfood". Use --delete to also remove it from the config file. + +=== Config after logout — profile should still exist +; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified. +[DEFAULT] + +[logfood] +host = [DATABRICKS_URL] +account_id = stale-account +auth_type = databricks-cli + +[__settings__] +default_profile = logfood + +=== Token cache keys after logout — both entries should be removed +[] + +=== Profiles after logout — logfood should be invalid +>>> [CLI] auth profiles +Name Host Valid +logfood (Default) [DATABRICKS_URL] NO + +=== Logged out profile should no longer return a token +>>> musterr [CLI] auth token --profile logfood +Error: cache: databricks OAuth is not configured for this host. Try logging in again with `databricks auth login --profile logfood` before retrying. If this fails, please report this issue to the Databricks CLI maintainers at https://github.com/databricks/cli/issues/new diff --git a/acceptance/cmd/auth/logout/stale-account-id-workspace-host/script b/acceptance/cmd/auth/logout/stale-account-id-workspace-host/script new file mode 100644 index 0000000000..036adfe661 --- /dev/null +++ b/acceptance/cmd/auth/logout/stale-account-id-workspace-host/script @@ -0,0 +1,52 @@ +sethome "./home" + +cat > "./home/.databrickscfg" < "./home/.databricks/token-cache.json" < "./home/.databrickscfg" <<'EOF' +cat > "./home/.databrickscfg" < "./home/.databricks/token-cache.json" <<'EOF' +cat > "./home/.databricks/token-cache.json" < "./home/.databricks/token-cache.json" <<'EOF' "access_token": "staging-cached-token", "token_type": "Bearer" }, - "https://shared.cloud.databricks.com": { + "${DATABRICKS_HOST}": { "access_token": "shared-host-token", "token_type": "Bearer" } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 6d11481f37..f6a4a434a1 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -6,6 +6,7 @@ import ( "fmt" "strings" + "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" @@ -282,19 +283,33 @@ func clearTokenCache(ctx context.Context, p profile.Profile, profiler profile.Pr return nil } -// hostCacheKeyAndMatchFn returns the token cache key and a profile match -// function for the host-based token entry. Account and unified profiles use -// host/oidc/accounts/ as the cache key and match on both host and -// account ID; workspace profiles use just the host. +// hostCacheKeyAndMatchFn returns the host-based token cache key and a profile +// match function for the given profile. The match function is used to check +// whether other profiles share the same host-based cache entry. func hostCacheKeyAndMatchFn(p profile.Profile) (string, profile.ProfileMatchFunction) { - host := (&config.Config{Host: p.Host}).CanonicalHostName() - if host == "" { + // Use ToOAuthArgument to derive the host-based cache key via the same + // routing logic the SDK used when the token was written during login. + // This includes a .well-known/databricks-config call that distinguishes + // classic workspace hosts from SPOG hosts — a distinction that cannot + // be made from the profile fields alone. + arg, err := (auth.AuthArguments{ + Host: p.Host, + AccountID: p.AccountID, + WorkspaceID: p.WorkspaceID, + IsUnifiedHost: p.IsUnifiedHost, + // Profile is deliberately empty so GetCacheKey returns the host-based + // key rather than the profile name. + // DiscoveryURL is left empty to force a fresh .well-known resolution + // so that the routing decision reflects the host's current state. + }).ToOAuthArgument() + if err != nil { return "", nil } + hostCacheKey := arg.GetCacheKey() - if p.AccountID != "" { - return host + "/oidc/accounts/" + p.AccountID, profile.WithHostAndAccountID(host, p.AccountID) + host := (&config.Config{Host: p.Host}).CanonicalHostName() + if p.AccountID != "" && strings.Contains(hostCacheKey, "/oidc/accounts/") { + return hostCacheKey, profile.WithHostAndAccountID(host, p.AccountID) } - - return host, profile.WithHost(host) + return hostCacheKey, profile.WithHost(host) } diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index 50c152deb5..652a28a67a 100644 --- a/cmd/auth/logout_test.go +++ b/cmd/auth/logout_test.go @@ -1,6 +1,9 @@ package auth import ( + "encoding/json" + "net/http" + "net/http/httptest" "os" "path/filepath" "testing" @@ -27,6 +30,11 @@ auth_type = databricks-cli host = https://my-unique-workspace.cloud.databricks.com auth_type = databricks-cli +[my-workspace-stale-account] +host = https://stale-account.cloud.databricks.com +account_id = stale-account +auth_type = databricks-cli + [my-account] host = https://accounts.cloud.databricks.com account_id = abc123 @@ -44,13 +52,15 @@ token = dev-token ` var logoutTestTokensCacheConfig = map[string]*oauth2.Token{ - "my-workspace": {AccessToken: "shared-workspace-token"}, - "shared-workspace": {AccessToken: "shared-workspace-token"}, - "my-unique-workspace": {AccessToken: "my-unique-workspace-token"}, - "my-account": {AccessToken: "my-account-token"}, - "my-unified": {AccessToken: "my-unified-token"}, + "my-workspace": {AccessToken: "shared-workspace-token"}, + "shared-workspace": {AccessToken: "shared-workspace-token"}, + "my-unique-workspace": {AccessToken: "my-unique-workspace-token"}, + "my-workspace-stale-account": {AccessToken: "stale-account-token"}, + "my-account": {AccessToken: "my-account-token"}, + "my-unified": {AccessToken: "my-unified-token"}, "https://my-workspace.cloud.databricks.com": {AccessToken: "shared-workspace-host-token"}, "https://my-unique-workspace.cloud.databricks.com": {AccessToken: "unique-workspace-host-token"}, + "https://stale-account.cloud.databricks.com": {AccessToken: "stale-account-host-token"}, "https://accounts.cloud.databricks.com/oidc/accounts/abc123": {AccessToken: "account-host-token"}, "https://unified.cloud.databricks.com/oidc/accounts/def456": {AccessToken: "unified-host-token"}, "my-m2m": {AccessToken: "m2m-service-token"}, @@ -97,6 +107,13 @@ func TestLogout(t *testing.T) { isSharedKey: false, force: true, }, + { + name: "existing workspace profile with stale account id", + profileName: "my-workspace-stale-account", + hostBasedKey: "https://stale-account.cloud.databricks.com", + isSharedKey: false, + force: true, + }, { name: "existing account profile", profileName: "my-account", @@ -329,3 +346,138 @@ default_profile = my-workspace }) } } + +// newWellKnownServer creates a mock server that serves /.well-known/databricks-config +// with the given oidc_endpoint shape. Use accountScoped=true for SPOG hosts. +func newWellKnownServer(t *testing.T, accountScoped bool, accountID string) *httptest.Server { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + w.Header().Set("Content-Type", "application/json") + 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, + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + t.Cleanup(server.Close) + return server +} + +func TestLogoutSPOGProfile(t *testing.T) { + spogServer := newWellKnownServer(t, true, "spog-acct") + + ctx := cmdio.MockDiscard(t.Context()) + configPath := writeTempConfig(t, `[DEFAULT] +[spog-profile] +host = `+spogServer.URL+` +account_id = spog-acct +workspace_id = spog-ws +auth_type = databricks-cli +`) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + hostKey := spogServer.URL + "/oidc/accounts/spog-acct" + tokenCache := &inMemoryTokenCache{ + Tokens: map[string]*oauth2.Token{ + "spog-profile": {AccessToken: "spog-profile-token"}, + hostKey: {AccessToken: "spog-host-token"}, + }, + } + + err := runLogout(ctx, logoutArgs{ + profileName: "spog-profile", + force: true, + profiler: profile.DefaultProfiler, + tokenCache: tokenCache, + configFilePath: configPath, + }) + require.NoError(t, err) + + assert.Nil(t, tokenCache.Tokens["spog-profile"]) + assert.Nil(t, tokenCache.Tokens[hostKey]) +} + +func TestHostCacheKeyAndMatchFn(t *testing.T) { + wsServer := newWellKnownServer(t, false, "ws-account") + spogServer := newWellKnownServer(t, true, "spog-account") + + cases := []struct { + name string + profile profile.Profile + wantKey string + wantKeyEmpty bool + }{ + { + name: "classic workspace", + profile: profile.Profile{ + Name: "ws", + Host: wsServer.URL, + }, + wantKey: wsServer.URL, + }, + { + name: "workspace with stale account_id", + profile: profile.Profile{ + Name: "stale", + Host: wsServer.URL, + AccountID: "stale-account", + }, + wantKey: wsServer.URL, + }, + { + name: "classic account host", + profile: profile.Profile{ + Name: "acct", + Host: "https://accounts.cloud.databricks.com", + AccountID: "abc123", + }, + wantKey: "https://accounts.cloud.databricks.com/oidc/accounts/abc123", + }, + { + name: "unified host with flag", + profile: profile.Profile{ + Name: "unified", + Host: wsServer.URL, + AccountID: "def456", + IsUnifiedHost: true, + }, + wantKey: wsServer.URL + "/oidc/accounts/def456", + }, + { + name: "SPOG profile routes to account key via discovery", + profile: profile.Profile{ + Name: "spog", + Host: spogServer.URL, + AccountID: "spog-account", + }, + wantKey: spogServer.URL + "/oidc/accounts/spog-account", + }, + { + name: "empty host returns empty", + profile: profile.Profile{ + Name: "no-host", + }, + wantKeyEmpty: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + key, matchFn := hostCacheKeyAndMatchFn(tc.profile) + if tc.wantKeyEmpty { + assert.Empty(t, key) + assert.Nil(t, matchFn) + return + } + assert.Equal(t, tc.wantKey, key) + require.NotNil(t, matchFn) + }) + } +}