auth: support OAuth group role assumption

acceptance/bin/browser.py

@@ -6,14 +6,28 @@
 Usage: browser.py <url>
 """
 
+import os
 import sys
+import urllib.parse
 import urllib.request
 
 if len(sys.argv) < 2:
     sys.stderr.write("Usage: browser.py <url>\n")
     sys.exit(1)
 
 url = sys.argv[1]
+expected_group_id = os.environ.get("DATABRICKS_TEST_GROUP_ID")
+if expected_group_id is not None:
+    group_ids = urllib.parse.parse_qs(urllib.parse.urlparse(url).query).get(
+        "assume_group", []
+    )
+    expected_group_ids = [] if expected_group_id == "" else [expected_group_id]
+    if group_ids != expected_group_ids:
+        sys.stderr.write(
+            f"Expected assume_group values {expected_group_ids!r}, got {group_ids!r}\n"
+        )
+        sys.exit(1)
+
 try:
     response = urllib.request.urlopen(url)
     if response.status != 200:

acceptance/bin/discovery_browser.py

@@ -31,6 +31,22 @@
 dest_parsed = urllib.parse.urlparse(destination_url)
 dest_params = urllib.parse.parse_qs(dest_parsed.query)
 
+expected_group_id = os.environ.get("DATABRICKS_TEST_GROUP_ID")
+if expected_group_id is not None:
+    top_level_group_ids = top_params.get("assume_group", [])
+    if top_level_group_ids:
+        sys.stderr.write(
+            f"Expected no top-level assume_group, got {top_level_group_ids!r}\n"
+        )
+        sys.exit(1)
+    group_ids = dest_params.get("assume_group", [])
+    expected_group_ids = [] if expected_group_id == "" else [expected_group_id]
+    if group_ids != expected_group_ids:
+        sys.stderr.write(
+            f"Expected nested assume_group values {expected_group_ids!r}, got {group_ids!r}\n"
+        )
+        sys.exit(1)
+
 redirect_uri = dest_params.get("redirect_uri", [None])[0]
 state = dest_params.get("state", [None])[0]
 

acceptance/cmd/auth/describe/account-host-with-workspace-id/output.txt

@@ -4,6 +4,7 @@
 >>> [CLI] auth describe --profile acct-with-ws
 Host: [DATABRICKS_URL]
 Account ID: acct-123
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:

acceptance/cmd/auth/describe/bundle-profile-env/output.txt

@@ -4,6 +4,7 @@
 >>> [CLI] auth describe
 Host: [DATABRICKS_URL]
 User: [USERNAME]
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:

acceptance/cmd/auth/describe/default-profile/output.txt

@@ -4,6 +4,7 @@
 >>> [CLI] auth describe
 Host: [DATABRICKS_URL]
 User: [USERNAME]
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:
@@ -22,6 +23,7 @@ Current configuration:
 >>> [CLI] auth describe
 Host: [DATABRICKS_URL]
 User: [USERNAME]
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:

acceptance/cmd/auth/describe/group-id-normal/home/.databrickscfg.tmpl

@@ -0,0 +1,4 @@
+[normal]
+host = $DATABRICKS_HOST
+workspace_id = 900800700600
+auth_type = databricks-cli

acceptance/cmd/auth/describe/group-id-normal/input-token-cache.json

@@ -0,0 +1,10 @@
+{
+  "version": 1,
+  "tokens": {
+    "normal": {
+      "access_token": "existing-access-token",
+      "refresh_token": "existing-refresh-token",
+      "token_type": "Bearer"
+    }
+  }
+}

acceptance/cmd/auth/describe/group-id-normal/output.txt

@@ -0,0 +1,18 @@
+
+>>> [CLI] auth describe --profile normal
+Host: [DATABRICKS_URL]
+User: [USERNAME]
+Assumed group ID: None (normal user permissions)
+Authenticated with: databricks-cli
+Token storage: plaintext, ~/.databricks/token-cache.json (from DATABRICKS_AUTH_STORAGE environment variable)
+-----
+Current configuration:
+  ✓ host: [DATABRICKS_URL] (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ✓ workspace_id: [NUMID] (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ~ token: ******** (from DATABRICKS_TOKEN environment variable, not used for auth type databricks-cli)
+  ✓ profile: normal (from --profile flag)
+  ✓ databricks_cli_path: [CLI]
+  ✓ auth_type: databricks-cli (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ✓ rate_limit: [NUMID] (from DATABRICKS_RATE_LIMIT environment variable)
+  ✓ cloud: AWS
+  ✓ discovery_url: [DATABRICKS_URL]/oidc/.well-known/oauth-authorization-server

acceptance/cmd/auth/describe/group-id-normal/script

@@ -0,0 +1,9 @@
+sethome "./home"
+
+envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg"
+mkdir -p "./home/.databricks"
+cp "./input-token-cache.json" "./home/.databricks/token-cache.json"
+export DATABRICKS_AUTH_STORAGE=plaintext
+
+# A profile without group_id is reported as using normal user permissions.
+trace $CLI auth describe --profile normal

acceptance/cmd/auth/describe/group-id-normal/test.toml

@@ -0,0 +1 @@
+Ignore = ["home"]

acceptance/cmd/auth/describe/group-id-role/home/.databrickscfg.tmpl

@@ -0,0 +1,5 @@
+[role]
+host = $DATABRICKS_HOST
+workspace_id = 900800700600
+group_id = group-A
+auth_type = databricks-cli

acceptance/cmd/auth/describe/group-id-role/input-token-cache.json

@@ -0,0 +1,10 @@
+{
+  "version": 1,
+  "tokens": {
+    "role": {
+      "access_token": "existing-access-token",
+      "refresh_token": "existing-refresh-token",
+      "token_type": "Bearer"
+    }
+  }
+}

acceptance/cmd/auth/describe/group-id-role/output.txt

@@ -0,0 +1,19 @@
+
+>>> [CLI] auth describe --profile role
+Host: [DATABRICKS_URL]
+User: [USERNAME]
+Assumed group ID: group-A
+Authenticated with: databricks-cli
+Token storage: plaintext, ~/.databricks/token-cache.json (from DATABRICKS_AUTH_STORAGE environment variable)
+-----
+Current configuration:
+  ✓ host: [DATABRICKS_URL] (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ✓ workspace_id: [NUMID] (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ✓ group_id: group-A (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ~ token: ******** (from DATABRICKS_TOKEN environment variable, not used for auth type databricks-cli)
+  ✓ profile: role (from --profile flag)
+  ✓ databricks_cli_path: [CLI]
+  ✓ auth_type: databricks-cli (from [TEST_TMP_DIR]/home/.databrickscfg config file)
+  ✓ rate_limit: [NUMID] (from DATABRICKS_RATE_LIMIT environment variable)
+  ✓ cloud: AWS
+  ✓ discovery_url: [DATABRICKS_URL]/oidc/.well-known/oauth-authorization-server

acceptance/cmd/auth/describe/group-id-role/script

@@ -0,0 +1,9 @@
+sethome "./home"
+
+envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg"
+mkdir -p "./home/.databricks"
+cp "./input-token-cache.json" "./home/.databricks/token-cache.json"
+export DATABRICKS_AUTH_STORAGE=plaintext
+
+# The saved role is reported alongside the authenticated user identity.
+trace $CLI auth describe --profile role

acceptance/cmd/auth/describe/group-id-role/test.toml

@@ -0,0 +1 @@
+Ignore = ["home"]

acceptance/cmd/auth/describe/profile-overrides-env/output.txt

@@ -4,6 +4,7 @@
 >>> [CLI] auth describe --profile my-workspace
 Host: [DATABRICKS_URL]
 User: [USERNAME]
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:
@@ -22,6 +23,7 @@ Current configuration:
 >>> [CLI] auth describe --profile host-only
 Host: [DATABRICKS_URL]
 User: [USERNAME]
+Assumed group ID: None (normal user permissions)
 Authenticated with: pat
 -----
 Current configuration:

acceptance/cmd/auth/describe/u2m-plaintext-config/output.txt

@@ -2,6 +2,7 @@
 >>> [CLI] auth describe --profile u2m-profile
 Warn: [hostmetadata] failed to fetch host metadata for https://u2m-profile.databricks.test, will skip for 1m0s
 Unable to authenticate: error getting token: cache: no cached credentials; run `databricks auth login` to sign in
+Assumed group ID: None (normal user permissions)
 Token storage: plaintext, ~/.databricks/token-cache.json (from auth_storage in [__settings__] section of [TEST_TMP_DIR]/home/.databrickscfg)
 -----
 Current configuration:

acceptance/cmd/auth/describe/u2m-plaintext-env/output.txt

@@ -2,6 +2,7 @@
 >>> [CLI] auth describe --profile u2m-profile
 Warn: [hostmetadata] failed to fetch host metadata for https://u2m-profile.databricks.test, will skip for 1m0s
 Unable to authenticate: error getting token: cache: no cached credentials; run `databricks auth login` to sign in
+Assumed group ID: None (normal user permissions)
 Token storage: plaintext, ~/.databricks/token-cache.json (from DATABRICKS_AUTH_STORAGE environment variable)
 -----
 Current configuration:

acceptance/cmd/auth/describe/u2m-secure-default/output.txt

@@ -2,6 +2,7 @@
 >>> [CLI] auth describe --profile u2m-profile
 Warn: [hostmetadata] failed to fetch host metadata for https://u2m-profile.databricks.test, will skip for 1m0s
 Unable to authenticate: error getting token: [KEYRING_LOOKUP_ERROR]
+Assumed group ID: None (normal user permissions)
 Token storage: secure, OS keyring (service: databricks-cli) (from default)
 -----
 Current configuration:

acceptance/cmd/auth/login/discovery-group-id/out.databrickscfg

@@ -0,0 +1,12 @@
+; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
+[DEFAULT]
+
+[discovery-group]
+host         = [DATABRICKS_URL]
+account_id   = test-account-123
+workspace_id = [NUMID]
+group_id     = group-A
+auth_type    = databricks-cli
+
+[__settings__]
+default_profile = discovery-group

acceptance/cmd/auth/login/discovery-group-id/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] auth login --profile discovery-group --group-id group-A
+Opening login.databricks.com in your browser...
+Profile discovery-group was successfully saved

acceptance/cmd/auth/login/discovery-group-id/script

@@ -0,0 +1,11 @@
+sethome "./home"
+
+export BROWSER="discovery_browser.py"
+export DATABRICKS_TEST_GROUP_ID="group-A"
+
+# The browser helper verifies assume_group appears only inside destination_url.
+trace $CLI auth login --profile discovery-group --group-id group-A
+
+mv "./home/.databrickscfg" "./out.databrickscfg"
+# Consume the recorded requests so they are not treated as test output.
+print_requests.py //does-not-exist

acceptance/cmd/auth/login/discovery-group-id/test.toml

@@ -0,0 +1,17 @@
+Ignore = [
+    "home"
+]
+RecordRequests = true
+
+[[Server]]
+Pattern = "GET /api/2.0/tokens/introspect"
+Response.Body = '''
+{
+    "principal_context": {
+        "authentication_scope": {
+            "account_id": "test-account-123",
+            "workspace_id": 12345
+        }
+    }
+}
+'''

acceptance/cmd/auth/login/group-id-clear/home/.databrickscfg.tmpl

@@ -0,0 +1,4 @@
+[role]
+host = $DATABRICKS_HOST
+group_id = group-A
+auth_type = databricks-cli

acceptance/cmd/auth/login/group-id-clear/out.databrickscfg

@@ -0,0 +1,7 @@
+; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
+[DEFAULT]
+
+[role]
+host         = [DATABRICKS_URL]
+auth_type    = databricks-cli
+workspace_id = [NUMID]

acceptance/cmd/auth/login/group-id-clear/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] auth login --profile role --clear-group-id
+Profile role was successfully saved

acceptance/cmd/auth/login/group-id-clear/script

@@ -0,0 +1,11 @@
+sethome "./home"
+
+envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg"
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
+export DATABRICKS_TEST_GROUP_ID=""
+
+# Clearing omits assume_group from OAuth and removes group_id from the profile.
+trace $CLI auth login --profile role --clear-group-id
+
+mv "./home/.databrickscfg" "./out.databrickscfg"

acceptance/cmd/auth/login/group-id-clear/test.toml

@@ -0,0 +1,15 @@
+Ignore = ["home"]
+
+[[Server]]
+Pattern = "POST /oidc/v1/token"
+Response.Body = '''
+{
+  "access_token": "oauth-token",
+  "refresh_token": "oauth-refresh-token",
+  "expires_in": 3600,
+  "scope": "all-apis",
+  "token_type": "Bearer"
+}
+'''
+[Server.Response.Headers]
+"Content-Type" = ["application/json"]

acceptance/cmd/auth/login/group-id-explicit/out.databrickscfg

@@ -0,0 +1,11 @@
+; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
+[DEFAULT]
+
+[role]
+host         = [DATABRICKS_URL]
+workspace_id = [NUMID]
+group_id     = group-A
+auth_type    = databricks-cli
+
+[__settings__]
+default_profile = role

acceptance/cmd/auth/login/group-id-explicit/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] auth login --host [DATABRICKS_URL] --profile role --group-id group-A
+Profile role was successfully saved

acceptance/cmd/auth/login/group-id-explicit/script

@@ -0,0 +1,10 @@
+sethome "./home"
+
+export BROWSER="browser.py"
+export DATABRICKS_AUTH_STORAGE=plaintext
+export DATABRICKS_TEST_GROUP_ID="group-A"
+
+# The browser helper verifies the explicit group is sent exactly once.
+trace $CLI auth login --host $DATABRICKS_HOST --profile role --group-id group-A
+
+mv "./home/.databrickscfg" "./out.databrickscfg"

acceptance/cmd/auth/login/group-id-explicit/test.toml

@@ -0,0 +1,15 @@
+Ignore = ["home"]
+
+[[Server]]
+Pattern = "POST /oidc/v1/token"
+Response.Body = '''
+{
+  "access_token": "oauth-token",
+  "refresh_token": "oauth-refresh-token",
+  "expires_in": 3600,
+  "scope": "all-apis",
+  "token_type": "Bearer"
+}
+'''
+[Server.Response.Headers]
+"Content-Type" = ["application/json"]

acceptance/cmd/auth/login/group-id-preserve/home/.databrickscfg.tmpl

@@ -0,0 +1,4 @@
+[role]
+host = $DATABRICKS_HOST
+group_id = group-A
+auth_type = databricks-cli

acceptance/cmd/auth/login/group-id-preserve/out.databrickscfg

@@ -0,0 +1,8 @@
+; The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified.
+[DEFAULT]
+
+[role]
+host         = [DATABRICKS_URL]
+group_id     = group-A
+auth_type    = databricks-cli
+workspace_id = [NUMID]

acceptance/cmd/auth/login/group-id-preserve/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] auth login --profile role
+Profile role was successfully saved