Skip to content

Commit 4a4b68e

Browse files
mihaimitrea-dbmihaimitrea-db
committed
Fix auth logout failing to clear token for workspace profiles with stale account_id
logout derived the token-cache key by checking whether `account_id` was set on the profile, but a workspace profile can carry a stale `account_id` from an earlier account-level login. This caused logout to build an account-style cache key (`host/oidc/accounts/<id>`) that did not match the workspace-style key the SDK actually wrote, so the cached token was never deleted. Use `config.HostType()` instead to decide the cache-key shape, which matches the SDK's own OAuth routing logic and correctly treats workspace hosts as workspace profiles regardless of leftover account metadata. Co-authored-by: Isaac
1 parent d062420 commit 4a4b68e

5 files changed

Lines changed: 134 additions & 10 deletions

File tree

‎acceptance/cmd/auth/logout/stale-account-id-workspace-host/out.test.toml‎

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎acceptance/cmd/auth/logout/stale-account-id-workspace-host/output.txt‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
2+
=== Profiles before logout — logfood should be valid
3+
4+
>>> [CLI] auth profiles
5+
Name Host Valid
6+
logfood (Default) [DATABRICKS_URL] YES
7+
8+
=== Token cache keys before logout
9+
[
10+
"[DATABRICKS_URL]",
11+
"logfood"
12+
]
13+
14+
=== Logout without --delete
15+
>>> [CLI] auth logout --profile logfood --force
16+
Logged out of profile "logfood". Use --delete to also remove it from the config file.
17+
18+
=== Config after logout — profile should still exist
19+
; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
20+
[DEFAULT]
21+
22+
[logfood]
23+
host = [DATABRICKS_URL]
24+
account_id = stale-account
25+
auth_type = databricks-cli
26+
27+
[__settings__]
28+
default_profile = logfood
29+
30+
=== Token cache keys after logout — both entries should be removed
31+
[]
32+
33+
=== Profiles after logout — logfood should be invalid
34+
35+
>>> [CLI] auth profiles
36+
Name Host Valid
37+
logfood (Default) [DATABRICKS_URL] NO
38+
39+
=== Logged out profile should no longer return a token
40+
41+
>>> musterr [CLI] auth token --profile logfood
42+
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

‎acceptance/cmd/auth/logout/stale-account-id-workspace-host/script‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
sethome "./home"
2+
3+
cat > "./home/.databrickscfg" <<EOF
4+
; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
5+
[DEFAULT]
6+
7+
[logfood]
8+
host = ${DATABRICKS_HOST}
9+
account_id = stale-account
10+
auth_type = databricks-cli
11+
12+
[__settings__]
13+
default_profile = logfood
14+
EOF
15+
16+
mkdir -p "./home/.databricks"
17+
cat > "./home/.databricks/token-cache.json" <<EOF
18+
{
19+
"version": 1,
20+
"tokens": {
21+
"logfood": {
22+
"access_token": "logfood-cached-token",
23+
"token_type": "Bearer"
24+
},
25+
"${DATABRICKS_HOST}": {
26+
"access_token": "logfood-host-token",
27+
"token_type": "Bearer"
28+
}
29+
}
30+
}
31+
EOF
32+
33+
title "Profiles before logout — logfood should be valid\n"
34+
trace $CLI auth profiles
35+
36+
title "Token cache keys before logout\n"
37+
jq -S '.tokens | keys' "./home/.databricks/token-cache.json"
38+
39+
title "Logout without --delete"
40+
trace $CLI auth logout --profile logfood --force
41+
42+
title "Config after logout — profile should still exist\n"
43+
cat "./home/.databrickscfg"
44+
45+
title "Token cache keys after logout — both entries should be removed\n"
46+
jq -S '.tokens | keys' "./home/.databricks/token-cache.json"
47+
48+
title "Profiles after logout — logfood should be invalid\n"
49+
trace $CLI auth profiles
50+
51+
title "Logged out profile should no longer return a token\n"
52+
trace musterr $CLI auth token --profile logfood

‎cmd/auth/logout.go‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -276,18 +276,29 @@ func clearTokenCache(ctx context.Context, p profile.Profile, profiler profile.Pr
276276
// hostCacheKeyAndMatchFn returns the token cache key and a profile match
277277
// function for the host-based token entry. Account and unified profiles use
278278
// host/oidc/accounts/<account_id> as the cache key and match on both host and
279-
// account ID; workspace profiles use just the host.
279+
// account ID. Workspace profiles use just the host, even if the profile still
280+
// carries stale account metadata from an earlier login.
280281
func hostCacheKeyAndMatchFn(p profile.Profile) (string, profile.ProfileMatchFunction) {
281-
host := (&config.Config{Host: p.Host}).CanonicalHostName()
282+
cfg := &config.Config{
283+
Host: p.Host,
284+
AccountID: p.AccountID,
285+
WorkspaceID: p.WorkspaceID,
286+
Experimental_IsUnifiedHost: p.IsUnifiedHost,
287+
}
288+
host := cfg.CanonicalHostName()
282289
if host == "" {
283290
return "", nil
284291
}
285292

286-
if p.AccountID != "" {
293+
switch cfg.HostType() {
294+
case config.AccountHost, config.UnifiedHost:
295+
if p.AccountID == "" {
296+
return "", nil
297+
}
287298
return host + "/oidc/accounts/" + p.AccountID, profile.WithHostAndAccountID(host, p.AccountID)
299+
default:
300+
return host, profile.WithHost(host)
288301
}
289-
290-
return host, profile.WithHost(host)
291302
}
292303

293304
// resolveLogoutArg resolves a positional argument to a profile name. It first

‎cmd/auth/logout_test.go‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ auth_type = databricks-cli
2626
host = https://my-unique-workspace.cloud.databricks.com
2727
auth_type = databricks-cli
2828
29+
[my-workspace-stale-account]
30+
host = https://stale-account.cloud.databricks.com
31+
account_id = stale-account
32+
auth_type = databricks-cli
33+
2934
[my-account]
3035
host = https://accounts.cloud.databricks.com
3136
account_id = abc123
@@ -43,13 +48,15 @@ token = dev-token
4348
`
4449

4550
var logoutTestTokensCacheConfig = map[string]*oauth2.Token{
46-
"my-workspace": {AccessToken: "shared-workspace-token"},
47-
"shared-workspace": {AccessToken: "shared-workspace-token"},
48-
"my-unique-workspace": {AccessToken: "my-unique-workspace-token"},
49-
"my-account": {AccessToken: "my-account-token"},
50-
"my-unified": {AccessToken: "my-unified-token"},
51+
"my-workspace": {AccessToken: "shared-workspace-token"},
52+
"shared-workspace": {AccessToken: "shared-workspace-token"},
53+
"my-unique-workspace": {AccessToken: "my-unique-workspace-token"},
54+
"my-workspace-stale-account": {AccessToken: "stale-account-token"},
55+
"my-account": {AccessToken: "my-account-token"},
56+
"my-unified": {AccessToken: "my-unified-token"},
5157
"https://my-workspace.cloud.databricks.com": {AccessToken: "shared-workspace-host-token"},
5258
"https://my-unique-workspace.cloud.databricks.com": {AccessToken: "unique-workspace-host-token"},
59+
"https://stale-account.cloud.databricks.com": {AccessToken: "stale-account-host-token"},
5360
"https://accounts.cloud.databricks.com/oidc/accounts/abc123": {AccessToken: "account-host-token"},
5461
"https://unified.cloud.databricks.com/oidc/accounts/def456": {AccessToken: "unified-host-token"},
5562
"my-m2m": {AccessToken: "m2m-service-token"},
@@ -96,6 +103,13 @@ func TestLogout(t *testing.T) {
96103
isSharedKey: false,
97104
force: true,
98105
},
106+
{
107+
name: "existing workspace profile with stale account id",
108+
profileName: "my-workspace-stale-account",
109+
hostBasedKey: "https://stale-account.cloud.databricks.com",
110+
isSharedKey: false,
111+
force: true,
112+
},
99113
{
100114
name: "existing account profile",
101115
profileName: "my-account",

0 commit comments

Comments
 (0)