Skip to content

Commit 1572b1a

Browse files
simonfaltumclaude
andauthored
Add discovery login via login.databricks.com (#4702)
## Why When users run `databricks auth login` without specifying a host, the CLI had no fallback. Users had to know their workspace URL upfront. This adds a discovery flow via login.databricks.com where users authenticate in the browser, select a workspace, and the CLI automatically discovers the workspace URL. This PR adds discovery login as the default option. A follow-up PR once this is merged, will merge the profile / new profile selector from `auth token` with `auth login` to ensure a uniform experience, and ensuring that users who already have profiles setup will be given the option to re-login as the primary option. ## Changes Before: `databricks auth login` without `--host` prompted for a host URL or failed. Now: `databricks auth login` without `--host` opens login.databricks.com in the browser. After the user authenticates and selects a workspace, the CLI saves the profile with the discovered host, account ID, and workspace ID. Implementation details: - `shouldUseDiscovery()` detects when no host is available from flags, args, or existing profile - `discoveryLogin()` orchestrates the browser-based flow using the SDK's discovery OAuth support - `IntrospectToken()` (in `libs/auth`) calls the workspace introspection endpoint to extract account/workspace IDs (best-effort, non-fatal on failure) - `splitScopes()` deduplicates scope parsing across login paths - Error messages include actionable tips (e.g., "you can specify a workspace directly with: databricks auth login --host <url>") - Discovery dependencies (`PersistentAuth`, `OAuthArgument`, `IntrospectToken`) are injectable via package-level vars for testability ### Post-review fixes - `IntrospectToken` accepts an `*http.Client` parameter instead of using `http.DefaultClient` (enterprise proxy/TLS compatibility) - Introspection HTTP client clones `http.DefaultTransport` to inherit system CA certs, proxy settings, and timeouts - Discovery relogin now reuses existing profile scopes when `--scopes` is not explicitly set (consistent with normal login path) - Explicitly clear stale routing fields (`account_id`, `workspace_id`, `experimental_is_unified_host`, `cluster_id`, `serverless_compute_id`) before saving discovery profile - `workspace_id` only set when introspection returns a fresh value - Consolidate error message construction with `discoveryErr()` helper ## Test plan - [x] Unit tests for `shouldUseDiscovery` (table-driven, 6 cases) - [x] Unit tests for `splitScopes` (empty, single, whitespace, empty entries) - [x] Unit tests for `IntrospectToken` (success, zero workspace ID, HTTP errors, malformed JSON, auth header, endpoint path, custom HTTP client) - [x] Integration test for `discoveryLogin` with introspection failure (verifies profile is still saved) - [x] Integration test for `discoveryLogin` with introspection success (verifies metadata extraction) - [x] Regression test for relogin preserving existing profile scopes - [x] Regression test for clearing stale routing fields from unified profile - [x] `go test ./cmd/auth/... ./libs/auth/...` passes --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 9aeed55 commit 1572b1a

File tree

11 files changed

+1067
-12
lines changed

11 files changed

+1067
-12
lines changed
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Simulates the login.databricks.com discovery flow for acceptance tests.
4+
5+
When the CLI opens this "browser" with the login.databricks.com URL,
6+
the script extracts the OAuth parameters from the destination_url,
7+
constructs a callback to localhost with an iss parameter pointing
8+
at the testserver, and fetches it.
9+
10+
Usage: discovery_browser.py <url>
11+
"""
12+
13+
import os
14+
import sys
15+
import urllib.parse
16+
import urllib.request
17+
18+
if len(sys.argv) < 2:
19+
sys.stderr.write("Usage: discovery_browser.py <url>\n")
20+
sys.exit(1)
21+
22+
url = sys.argv[1]
23+
parsed = urllib.parse.urlparse(url)
24+
top_params = urllib.parse.parse_qs(parsed.query)
25+
26+
destination_url = top_params.get("destination_url", [None])[0]
27+
if not destination_url:
28+
sys.stderr.write(f"No destination_url found in: {url}\n")
29+
sys.exit(1)
30+
31+
dest_parsed = urllib.parse.urlparse(destination_url)
32+
dest_params = urllib.parse.parse_qs(dest_parsed.query)
33+
34+
redirect_uri = dest_params.get("redirect_uri", [None])[0]
35+
state = dest_params.get("state", [None])[0]
36+
37+
if not redirect_uri or not state:
38+
sys.stderr.write(f"Missing redirect_uri or state in destination_url: {destination_url}\n")
39+
sys.exit(1)
40+
41+
# The testserver's host acts as the workspace issuer.
42+
testserver_host = os.environ.get("DATABRICKS_HOST", "")
43+
if not testserver_host:
44+
sys.stderr.write("DATABRICKS_HOST not set\n")
45+
sys.exit(1)
46+
47+
issuer = testserver_host.rstrip("/") + "/oidc"
48+
49+
# Build the callback URL with code, state, and iss (the workspace issuer).
50+
callback_params = urllib.parse.urlencode(
51+
{
52+
"code": "oauth-code",
53+
"state": state,
54+
"iss": issuer,
55+
}
56+
)
57+
callback_url = f"{redirect_uri}?{callback_params}"
58+
59+
try:
60+
response = urllib.request.urlopen(callback_url)
61+
if response.status != 200:
62+
sys.stderr.write(f"Callback failed: {callback_url} (status {response.status})\n")
63+
sys.exit(1)
64+
except Exception as e:
65+
sys.stderr.write(f"Callback failed: {callback_url} ({e})\n")
66+
sys.exit(1)
67+
68+
sys.exit(0)
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
2+
[DEFAULT]
3+
4+
[discovery-test]
5+
host = [DATABRICKS_URL]
6+
workspace_id = 12345
7+
auth_type = databricks-cli
8+
9+
[__settings__]
10+
default_profile = discovery-test

acceptance/cmd/auth/login/discovery/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.
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
2+
>>> [CLI] auth login --profile discovery-test
3+
Opening login.databricks.com in your browser...
4+
Profile discovery-test was successfully saved
5+
6+
>>> [CLI] auth profiles
7+
Name Host Valid
8+
discovery-test (Default) [DATABRICKS_URL] YES
9+
10+
>>> print_requests.py --get //tokens/introspect
11+
{
12+
"method": "GET",
13+
"path": "/api/2.0/tokens/introspect"
14+
}
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
sethome "./home"
2+
3+
# Use the discovery browser script that simulates login.databricks.com
4+
export BROWSER="discovery_browser.py"
5+
6+
trace $CLI auth login --profile discovery-test
7+
8+
trace $CLI auth profiles
9+
10+
# Verify the introspection endpoint was called (workspace_id in profile confirms this too).
11+
trace print_requests.py --get //tokens/introspect
12+
13+
# Track the .databrickscfg file that was created to surface changes.
14+
mv "./home/.databrickscfg" "./out.databrickscfg"
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
Ignore = [
2+
"home"
3+
]
4+
RecordRequests = true
5+
6+
# Override the introspection endpoint so we can verify it gets called.
7+
[[Server]]
8+
Pattern = "GET /api/2.0/tokens/introspect"
9+
Response.Body = '''
10+
{
11+
"principal_context": {
12+
"authentication_scope": {
13+
"account_id": "test-account-123",
14+
"workspace_id": 12345
15+
}
16+
}
17+
}
18+
'''

0 commit comments

Comments
 (0)