diff --git a/acceptance/bin/browser.py b/acceptance/bin/browser.py index 2ed3cca8ff2..479f7ab6133 100755 --- a/acceptance/bin/browser.py +++ b/acceptance/bin/browser.py @@ -6,7 +6,9 @@ Usage: browser.py """ +import os import sys +import urllib.parse import urllib.request if len(sys.argv) < 2: @@ -14,6 +16,18 @@ 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: diff --git a/acceptance/bin/discovery_browser.py b/acceptance/bin/discovery_browser.py index 42099fa06dd..09e4d2a8375 100755 --- a/acceptance/bin/discovery_browser.py +++ b/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] diff --git a/acceptance/cmd/auth/describe/account-host-with-workspace-id/output.txt b/acceptance/cmd/auth/describe/account-host-with-workspace-id/output.txt index 505db5b55c1..10c42aa339d 100644 --- a/acceptance/cmd/auth/describe/account-host-with-workspace-id/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/bundle-profile-env/output.txt b/acceptance/cmd/auth/describe/bundle-profile-env/output.txt index 38d278a905d..957be935b76 100644 --- a/acceptance/cmd/auth/describe/bundle-profile-env/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/default-profile/output.txt b/acceptance/cmd/auth/describe/default-profile/output.txt index d6a00bdcb8c..b9d5a9a0e20 100644 --- a/acceptance/cmd/auth/describe/default-profile/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/group-id-normal/home/.databrickscfg.tmpl b/acceptance/cmd/auth/describe/group-id-normal/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..6ea0089e197 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-normal/input-token-cache.json b/acceptance/cmd/auth/describe/group-id-normal/input-token-cache.json new file mode 100644 index 00000000000..61dd6bbea9d --- /dev/null +++ b/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" + } + } +} diff --git a/acceptance/cmd/auth/describe/group-id-normal/out.test.toml b/acceptance/cmd/auth/describe/group-id-normal/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/describe/group-id-normal/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/describe/group-id-normal/output.txt b/acceptance/cmd/auth/describe/group-id-normal/output.txt new file mode 100644 index 00000000000..3684ebf6f66 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-normal/script b/acceptance/cmd/auth/describe/group-id-normal/script new file mode 100644 index 00000000000..0f752413a71 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-normal/test.toml b/acceptance/cmd/auth/describe/group-id-normal/test.toml new file mode 100644 index 00000000000..1f7fede55dc --- /dev/null +++ b/acceptance/cmd/auth/describe/group-id-normal/test.toml @@ -0,0 +1 @@ +Ignore = ["home"] diff --git a/acceptance/cmd/auth/describe/group-id-role/home/.databrickscfg.tmpl b/acceptance/cmd/auth/describe/group-id-role/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..3c6de25e0f4 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-role/input-token-cache.json b/acceptance/cmd/auth/describe/group-id-role/input-token-cache.json new file mode 100644 index 00000000000..9479bfd857e --- /dev/null +++ b/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" + } + } +} diff --git a/acceptance/cmd/auth/describe/group-id-role/out.test.toml b/acceptance/cmd/auth/describe/group-id-role/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/describe/group-id-role/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/describe/group-id-role/output.txt b/acceptance/cmd/auth/describe/group-id-role/output.txt new file mode 100644 index 00000000000..f3069668942 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-role/script b/acceptance/cmd/auth/describe/group-id-role/script new file mode 100644 index 00000000000..7aa1ad2feae --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/describe/group-id-role/test.toml b/acceptance/cmd/auth/describe/group-id-role/test.toml new file mode 100644 index 00000000000..1f7fede55dc --- /dev/null +++ b/acceptance/cmd/auth/describe/group-id-role/test.toml @@ -0,0 +1 @@ +Ignore = ["home"] diff --git a/acceptance/cmd/auth/describe/profile-overrides-env/output.txt b/acceptance/cmd/auth/describe/profile-overrides-env/output.txt index 60872b97469..91114afe12d 100644 --- a/acceptance/cmd/auth/describe/profile-overrides-env/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/u2m-plaintext-config/output.txt b/acceptance/cmd/auth/describe/u2m-plaintext-config/output.txt index 85c9edca654..f771a4562e8 100644 --- a/acceptance/cmd/auth/describe/u2m-plaintext-config/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/u2m-plaintext-env/output.txt b/acceptance/cmd/auth/describe/u2m-plaintext-env/output.txt index 10d1896b040..af2bdc82a99 100644 --- a/acceptance/cmd/auth/describe/u2m-plaintext-env/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/describe/u2m-secure-default/output.txt b/acceptance/cmd/auth/describe/u2m-secure-default/output.txt index fab13663d44..59d0d39fa1b 100644 --- a/acceptance/cmd/auth/describe/u2m-secure-default/output.txt +++ b/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: diff --git a/acceptance/cmd/auth/login/discovery-group-id/out.databrickscfg b/acceptance/cmd/auth/login/discovery-group-id/out.databrickscfg new file mode 100644 index 00000000000..8734514e826 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/discovery-group-id/out.test.toml b/acceptance/cmd/auth/login/discovery-group-id/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/discovery-group-id/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/discovery-group-id/output.txt b/acceptance/cmd/auth/login/discovery-group-id/output.txt new file mode 100644 index 00000000000..63ba3609be9 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/discovery-group-id/script b/acceptance/cmd/auth/login/discovery-group-id/script new file mode 100644 index 00000000000..ea76a9668d6 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/discovery-group-id/test.toml b/acceptance/cmd/auth/login/discovery-group-id/test.toml new file mode 100644 index 00000000000..5ef6983b71e --- /dev/null +++ b/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 + } + } +} +''' diff --git a/acceptance/cmd/auth/login/group-id-clear/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/group-id-clear/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..ab67956d8c6 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/group-id-clear/out.databrickscfg b/acceptance/cmd/auth/login/group-id-clear/out.databrickscfg new file mode 100644 index 00000000000..9b75132d4e7 --- /dev/null +++ b/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] diff --git a/acceptance/cmd/auth/login/group-id-clear/out.test.toml b/acceptance/cmd/auth/login/group-id-clear/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-clear/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-clear/output.txt b/acceptance/cmd/auth/login/group-id-clear/output.txt new file mode 100644 index 00000000000..0399908a713 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/group-id-clear/script b/acceptance/cmd/auth/login/group-id-clear/script new file mode 100644 index 00000000000..2df1ab565a5 --- /dev/null +++ b/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" diff --git a/acceptance/cmd/auth/login/group-id-clear/test.toml b/acceptance/cmd/auth/login/group-id-clear/test.toml new file mode 100644 index 00000000000..0881b275203 --- /dev/null +++ b/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"] diff --git a/acceptance/cmd/auth/login/group-id-explicit/out.databrickscfg b/acceptance/cmd/auth/login/group-id-explicit/out.databrickscfg new file mode 100644 index 00000000000..f80559ae257 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/group-id-explicit/out.test.toml b/acceptance/cmd/auth/login/group-id-explicit/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-explicit/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-explicit/output.txt b/acceptance/cmd/auth/login/group-id-explicit/output.txt new file mode 100644 index 00000000000..91722c295c7 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/group-id-explicit/script b/acceptance/cmd/auth/login/group-id-explicit/script new file mode 100644 index 00000000000..d48b76316c9 --- /dev/null +++ b/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" diff --git a/acceptance/cmd/auth/login/group-id-explicit/test.toml b/acceptance/cmd/auth/login/group-id-explicit/test.toml new file mode 100644 index 00000000000..0881b275203 --- /dev/null +++ b/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"] diff --git a/acceptance/cmd/auth/login/group-id-preserve/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/group-id-preserve/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..ab67956d8c6 --- /dev/null +++ b/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 diff --git a/acceptance/cmd/auth/login/group-id-preserve/out.databrickscfg b/acceptance/cmd/auth/login/group-id-preserve/out.databrickscfg new file mode 100644 index 00000000000..7bda5867d22 --- /dev/null +++ b/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] diff --git a/acceptance/cmd/auth/login/group-id-preserve/out.test.toml b/acceptance/cmd/auth/login/group-id-preserve/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-preserve/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-preserve/output.txt b/acceptance/cmd/auth/login/group-id-preserve/output.txt new file mode 100644 index 00000000000..eddf3c17a98 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-preserve/output.txt @@ -0,0 +1,3 @@ + +>>> [CLI] auth login --profile role +Profile role was successfully saved diff --git a/acceptance/cmd/auth/login/group-id-preserve/script b/acceptance/cmd/auth/login/group-id-preserve/script new file mode 100644 index 00000000000..83bce6a0548 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-preserve/script @@ -0,0 +1,12 @@ +sethome "./home" + +envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg" +export BROWSER="browser.py" +export DATABRICKS_AUTH_STORAGE=plaintext +export DATABRICKS_GROUP_ID="ambient-group" +export DATABRICKS_TEST_GROUP_ID="group-A" + +# Omitting the flag preserves the profile group and ignores the ambient SDK value. +trace $CLI auth login --profile role + +mv "./home/.databrickscfg" "./out.databrickscfg" diff --git a/acceptance/cmd/auth/login/group-id-preserve/test.toml b/acceptance/cmd/auth/login/group-id-preserve/test.toml new file mode 100644 index 00000000000..0881b275203 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-preserve/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"] diff --git a/acceptance/cmd/auth/login/group-id-refresh/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/group-id-refresh/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..3c6de25e0f4 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/home/.databrickscfg.tmpl @@ -0,0 +1,5 @@ +[role] +host = $DATABRICKS_HOST +workspace_id = 900800700600 +group_id = group-A +auth_type = databricks-cli diff --git a/acceptance/cmd/auth/login/group-id-refresh/input-token-cache.json b/acceptance/cmd/auth/login/group-id-refresh/input-token-cache.json new file mode 100644 index 00000000000..9479bfd857e --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/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" + } + } +} diff --git a/acceptance/cmd/auth/login/group-id-refresh/out.test.toml b/acceptance/cmd/auth/login/group-id-refresh/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-refresh/output.txt b/acceptance/cmd/auth/login/group-id-refresh/output.txt new file mode 100644 index 00000000000..88cc006a1fd --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/output.txt @@ -0,0 +1,8 @@ + +>>> [CLI] auth token --profile role --force-refresh +{ + "token_type": "Bearer" +} + +>>> jq -s map(select(.path == "/oidc/v1/token")) | length == 1 and (.[0].raw_body | contains("grant_type=refresh_token")) and (.[0].raw_body | contains("assume_group") | not) [TEST_TMP_DIR]/out.requests.txt +true diff --git a/acceptance/cmd/auth/login/group-id-refresh/script b/acceptance/cmd/auth/login/group-id-refresh/script new file mode 100644 index 00000000000..e901be2c254 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/script @@ -0,0 +1,12 @@ +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 + +trace $CLI auth token --profile role --force-refresh | jq '{token_type}' +# Refresh must use grant_type=refresh_token without resending assume_group. +trace jq -s 'map(select(.path == "/oidc/v1/token")) | length == 1 and (.[0].raw_body | contains("grant_type=refresh_token")) and (.[0].raw_body | contains("assume_group") | not)' "$TEST_TMP_DIR/out.requests.txt" +# Consume the recorded requests so they are not treated as test output. +print_requests.py //does-not-exist diff --git a/acceptance/cmd/auth/login/group-id-refresh/test.toml b/acceptance/cmd/auth/login/group-id-refresh/test.toml new file mode 100644 index 00000000000..a22d90f0161 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-refresh/test.toml @@ -0,0 +1,16 @@ +Ignore = ["home"] +RecordRequests = true + +[[Server]] +Pattern = "POST /oidc/v1/token" +Response.Body = ''' +{ + "access_token": "refreshed-access-token", + "refresh_token": "refreshed-refresh-token", + "expires_in": 3600, + "scope": "all-apis", + "token_type": "Bearer" +} +''' +[Server.Response.Headers] +"Content-Type" = ["application/json"] diff --git a/acceptance/cmd/auth/login/group-id-rejection/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/group-id-rejection/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..8264dae27ee --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/home/.databrickscfg.tmpl @@ -0,0 +1,4 @@ +[role] +host = $DATABRICKS_HOST +auth_type = databricks-cli +group_id = group-A diff --git a/acceptance/cmd/auth/login/group-id-rejection/input-token-cache.json b/acceptance/cmd/auth/login/group-id-rejection/input-token-cache.json new file mode 100644 index 00000000000..9479bfd857e --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/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" + } + } +} diff --git a/acceptance/cmd/auth/login/group-id-rejection/out.test.toml b/acceptance/cmd/auth/login/group-id-rejection/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-rejection/output.txt b/acceptance/cmd/auth/login/group-id-rejection/output.txt new file mode 100644 index 00000000000..5c17cbf1ec9 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/output.txt @@ -0,0 +1,9 @@ + +>>> errcode [CLI] auth login --profile role --group-id group-B +Error: authorize: Post "[DATABRICKS_URL]/oidc/v1/token": http 403: {"error":"access_denied","error_description":"group role rejected"} + +Exit code: 1 + +>>> cmp ./before.databrickscfg ./home/.databrickscfg + +>>> cmp ./before-token-cache.json ./home/.databricks/token-cache.json diff --git a/acceptance/cmd/auth/login/group-id-rejection/script b/acceptance/cmd/auth/login/group-id-rejection/script new file mode 100644 index 00000000000..198afb2370f --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/script @@ -0,0 +1,17 @@ +sethome "./home" + +envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg" +mkdir -p "./home/.databricks" +cp "./input-token-cache.json" "./home/.databricks/token-cache.json" +cp "./home/.databrickscfg" "./before.databrickscfg" +cp "./home/.databricks/token-cache.json" "./before-token-cache.json" + +export BROWSER="browser.py" +export DATABRICKS_AUTH_STORAGE=plaintext +export DATABRICKS_TEST_GROUP_ID="group-B" + +trace errcode $CLI auth login --profile role --group-id group-B + +# A rejected role must not change either persistent file. +trace cmp "./before.databrickscfg" "./home/.databrickscfg" +trace cmp "./before-token-cache.json" "./home/.databricks/token-cache.json" diff --git a/acceptance/cmd/auth/login/group-id-rejection/test.toml b/acceptance/cmd/auth/login/group-id-rejection/test.toml new file mode 100644 index 00000000000..25c7b5eb957 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-rejection/test.toml @@ -0,0 +1,11 @@ +Local = true +Cloud = false +Ignore = ["home", "before.databrickscfg", "before-token-cache.json"] +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] + +[[Server]] +Pattern = "POST /oidc/v1/token" +Response.StatusCode = 403 +Response.Body = '{"error":"access_denied","error_description":"group role rejected"}' +[Server.Response.Headers] +"Content-Type" = ["application/json"] diff --git a/acceptance/cmd/auth/login/group-id-switch/home/.databrickscfg.tmpl b/acceptance/cmd/auth/login/group-id-switch/home/.databrickscfg.tmpl new file mode 100644 index 00000000000..ab67956d8c6 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/home/.databrickscfg.tmpl @@ -0,0 +1,4 @@ +[role] +host = $DATABRICKS_HOST +group_id = group-A +auth_type = databricks-cli diff --git a/acceptance/cmd/auth/login/group-id-switch/out.databrickscfg b/acceptance/cmd/auth/login/group-id-switch/out.databrickscfg new file mode 100644 index 00000000000..81c97ec2a61 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/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-B +auth_type = databricks-cli +workspace_id = [NUMID] diff --git a/acceptance/cmd/auth/login/group-id-switch/out.test.toml b/acceptance/cmd/auth/login/group-id-switch/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/group-id-switch/output.txt b/acceptance/cmd/auth/login/group-id-switch/output.txt new file mode 100644 index 00000000000..e77eda70b00 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/output.txt @@ -0,0 +1,3 @@ + +>>> [CLI] auth login --profile role --group-id group-B +Profile role was successfully saved diff --git a/acceptance/cmd/auth/login/group-id-switch/script b/acceptance/cmd/auth/login/group-id-switch/script new file mode 100644 index 00000000000..3c13004d934 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/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="group-B" + +# An explicit group replaces the group already saved in the profile. +trace $CLI auth login --profile role --group-id group-B + +mv "./home/.databrickscfg" "./out.databrickscfg" diff --git a/acceptance/cmd/auth/login/group-id-switch/test.toml b/acceptance/cmd/auth/login/group-id-switch/test.toml new file mode 100644 index 00000000000..0881b275203 --- /dev/null +++ b/acceptance/cmd/auth/login/group-id-switch/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"] diff --git a/cmd/auth/describe.go b/cmd/auth/describe.go index a58e964de27..05a2633b4d9 100644 --- a/cmd/auth/describe.go +++ b/cmd/auth/describe.go @@ -27,6 +27,11 @@ var authTemplate = `{{"Host:" | bold}} {{.Status.Details.Host}} {{- if .Status.Username}} {{"User:" | bold}} {{.Status.Username}} {{- end}} +{{- if .Status.GroupID}} +{{"Assumed group ID:" | bold}} {{.Status.GroupID}} +{{- else}} +{{"Assumed group ID:" | bold}} None (normal user permissions) +{{- end}} {{"Authenticated with:" | bold}} {{.Status.Details.AuthType}} {{- if .Status.TokenStorage}} {{"Token storage:" | bold}} {{.Status.TokenStorage.Mode}}, {{.Status.TokenStorage.Location}} {{ printf "(from %s)" .Status.TokenStorage.Source | italic}} @@ -35,6 +40,11 @@ var authTemplate = `{{"Host:" | bold}} {{.Status.Details.Host}} ` + configurationTemplate var errorTemplate = `Unable to authenticate: {{.Status.Error}} +{{- if .Status.GroupID}} +{{"Assumed group ID:" | bold}} {{.Status.GroupID}} +{{- else}} +{{"Assumed group ID:" | bold}} None (normal user permissions) +{{- end}} {{- if .Status.TokenStorage}} {{"Token storage:" | bold}} {{.Status.TokenStorage.Mode}}, {{.Status.TokenStorage.Location}} {{ printf "(from %s)" .Status.TokenStorage.Source | italic}} {{- end}} @@ -198,6 +208,7 @@ func successAuthStatus(ctx context.Context, cmd *cobra.Command, cfg *config.Conf Details: details, Username: username, AccountID: accountID, + GroupID: cfg.GroupID, TokenStorage: resolveTokenStorageInfo(ctx, details.AuthType), } } @@ -211,6 +222,7 @@ func errorAuthStatus(ctx context.Context, cmd *cobra.Command, cfg *config.Config Status: "error", Error: err, Details: details, + GroupID: cfg.GroupID, TokenStorage: resolveTokenStorageInfo(ctx, details.AuthType), } } @@ -243,6 +255,7 @@ type authStatus struct { Error error `json:"error,omitempty"` Username string `json:"username,omitempty"` AccountID string `json:"account_id,omitempty"` + GroupID string `json:"-"` Details config.AuthDetails `json:"details"` TokenStorage *tokenStorageInfo `json:"token_storage,omitempty"` } diff --git a/cmd/auth/login.go b/cmd/auth/login.go index f9f2531ac74..6bfc7d273ab 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -133,6 +133,8 @@ a new profile is created. var configureServerless bool var skipWorkspace bool var scopes string + var groupID string + var clearGroupID bool cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout, "Timeout for completing login challenge in the browser") cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false, @@ -143,12 +145,23 @@ a new profile is created. "Skip workspace selection for account-level access") cmd.Flags().StringVar(&scopes, "scopes", "", "Comma-separated list of OAuth scopes to request (defaults to 'all-apis')") + cmd.Flags().StringVar(&groupID, "group-id", "", + "ID of the Databricks group whose role to assume") + cmd.Flags().BoolVar(&clearGroupID, "clear-group-id", false, + "Clear the saved group ID and log in with normal user permissions") cmd.PreRunE = profileHostConflictCheck cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() + groupIDChanged := cmd.Flag("group-id").Changed + if groupIDChanged && clearGroupID { + return errors.New("--group-id and --clear-group-id cannot be used together") + } + if groupIDChanged && groupID == "" { + return errors.New("--group-id cannot be empty") + } // Cluster and Serverless are mutually exclusive. if configureCluster && configureServerless { @@ -257,6 +270,7 @@ a new profile is created. if err != nil { return err } + effectiveGroupID := resolveGroupID(groupID, groupIDChanged, clearGroupID, existingProfile) // If no host is available from any source, use the discovery flow // via login.databricks.com. @@ -269,6 +283,8 @@ a new profile is created. profileName: profileName, timeout: loginTimeout, scopes: scopes, + groupID: effectiveGroupID, + clearGroupID: clearGroupID, existingProfile: existingProfile, browserFunc: getBrowserFunc(cmd), tokenStore: tokenStore, @@ -306,6 +322,9 @@ a new profile is created. if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) } + if effectiveGroupID != "" { + persistentAuthOpts = append(persistentAuthOpts, u2m.WithGroupID(effectiveGroupID)) + } persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...) if err != nil { return err @@ -350,6 +369,9 @@ a new profile is created. // from .well-known discovery, so stale values would be misleading). clearKeys := oauthLoginClearKeys() clearKeys = append(clearKeys, databrickscfg.ExperimentalIsUnifiedHostKey) + if clearGroupID { + clearKeys = append(clearKeys, "group_id") + } switch { case configureCluster: @@ -394,6 +416,7 @@ a new profile is created. ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), ServerlessComputeID: serverlessComputeID, Scopes: scopesList, + GroupID: effectiveGroupID, }, clearKeys...) if err != nil { return err @@ -635,6 +658,8 @@ type discoveryLoginInputs struct { profileName string timeout time.Duration scopes string + groupID string + clearGroupID bool existingProfile *profile.Profile browserFunc func(string) error tokenStore storage.Store @@ -664,6 +689,9 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } + if in.groupID != "" { + opts = append(opts, u2m.WithGroupID(in.groupID)) + } discoveryHost := env.Get(ctx, discoveryHostEnvVar) if discoveryHost != "" { opts = append(opts, u2m.WithDiscoveryHost(discoveryHost)) @@ -742,6 +770,9 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { "cluster_id", "serverless_compute_id", ) + if in.clearGroupID { + clearKeys = append(clearKeys, "group_id") + } err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: in.profileName, Host: discoveredHost, @@ -749,6 +780,7 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { AccountID: accountID, WorkspaceID: workspaceID, Scopes: scopesList, + GroupID: in.groupID, ConfigFile: configFile, }, clearKeys...) if err != nil { @@ -762,6 +794,19 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { return nil } +func resolveGroupID(flagValue string, flagChanged, clear bool, existingProfile *profile.Profile) string { + if clear { + return "" + } + if flagChanged { + return flagValue + } + if existingProfile != nil { + return existingProfile.GroupID + } + return "" +} + // splitScopes splits a comma-separated scopes string into a trimmed slice. func splitScopes(scopes string) []string { var result []string diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a8eafb4be43..e3a7b5ae06f 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -458,6 +458,45 @@ func TestSplitScopes(t *testing.T) { } } +func TestResolveGroupID(t *testing.T) { + existing := &profile.Profile{GroupID: "saved-group"} + tests := []struct { + name string + flagValue string + flagChanged bool + clear bool + existing *profile.Profile + want string + }{ + { + name: "explicit overrides profile", + flagValue: "new-group", + flagChanged: true, + existing: existing, + want: "new-group", + }, + { + name: "omitted preserves profile", + existing: existing, + want: "saved-group", + }, + { + name: "clear removes profile value", + clear: true, + existing: existing, + }, + { + name: "no existing profile", + want: "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, resolveGroupID(tt.flagValue, tt.flagChanged, tt.clear, tt.existing)) + }) + } +} + func TestRunHostDiscovery_NoHost(t *testing.T) { ctx := t.Context() args := &auth.AuthArguments{} @@ -1020,6 +1059,105 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { assert.Equal(t, "all-apis", savedProfile.Scopes) } +func TestDiscoveryLogin_GroupIDPersistence(t *testing.T) { + tests := []struct { + name string + groupID string + clearGroupID bool + existing string + want string + }{ + { + name: "explicit group overrides profile", + groupID: "new-group", + existing: "saved-group", + want: "new-group", + }, + { + name: "omitted group preserves profile", + groupID: "saved-group", + existing: "saved-group", + want: "saved-group", + }, + { + name: "clear group removes profile value", + clearGroupID: true, + existing: "saved-group", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + configPath := filepath.Join(t.TempDir(), ".databrickscfg") + contents := "[DISCOVERY]\nhost = https://old-workspace.example.com\nauth_type = databricks-cli\n" + if tt.existing != "" { + contents += "group_id = " + tt.existing + "\n" + } + require.NoError(t, os.WriteFile(configPath, []byte(contents), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + oauthArg.SetDiscoveredHost("https://workspace.example.com") + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + token: &oauth2.Token{AccessToken: "test-token"}, + }, + introspectionErr: errors.New("introspection failed"), + } + + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + groupID: tt.groupID, + clearGroupID: tt.clearGroupID, + existingProfile: &profile.Profile{GroupID: tt.existing}, + browserFunc: func(string) error { return nil }, + tokenStore: newTestStore(), + }) + require.NoError(t, err) + + savedProfile := loadTestProfile(t, ctx, "DISCOVERY") + assert.Equal(t, tt.want, savedProfile.GroupID) + }) + } +} + +func TestDiscoveryLogin_OAuthFailureDoesNotModifyProfile(t *testing.T) { + configPath := filepath.Join(t.TempDir(), ".databrickscfg") + want := []byte("[DISCOVERY]\nhost = https://old-workspace.example.com\nauth_type = databricks-cli\ngroup_id = saved-group\n") + require.NoError(t, os.WriteFile(configPath, want, 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + + oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY") + require.NoError(t, err) + dc := &fakeDiscoveryClient{ + oauthArg: oauthArg, + persistentAuth: &fakeDiscoveryPersistentAuth{ + challengeErr: errors.New("group role rejected"), + }, + } + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + groupID: "rejected-group", + browserFunc: func(string) error { return nil }, + tokenStore: newTestStore(), + }) + require.Error(t, err) + assert.ErrorContains(t, err, "login via login.databricks.com failed: group role rejected") + + got, err := os.ReadFile(configPath) + require.NoError(t, err) + assert.Equal(t, want, got) +} + func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { // Start a mock server that returns SPOG discovery metadata. server := newDiscoveryServer(t, map[string]any{ @@ -1280,3 +1418,21 @@ func TestLoginRejectsPositionalArgWithProfileFlag(t *testing.T) { err := cmd.Execute() assert.ErrorContains(t, err, `argument "https://example.com" cannot be combined with --host or --profile`) } + +func TestLoginRejectsGroupIDAndClearGroupID(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := newLoginCommand(&auth.AuthArguments{}) + cmd.Flags().String("profile", "", "") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"--group-id", "group-123", "--clear-group-id"}) + assert.EqualError(t, cmd.Execute(), "--group-id and --clear-group-id cannot be used together") +} + +func TestLoginRejectsExplicitEmptyGroupID(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := newLoginCommand(&auth.AuthArguments{}) + cmd.Flags().String("profile", "", "") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"--group-id="}) + assert.EqualError(t, cmd.Execute(), "--group-id cannot be empty") +} diff --git a/libs/databrickscfg/loader.go b/libs/databrickscfg/loader.go index 732dcf9024f..6917d65618e 100644 --- a/libs/databrickscfg/loader.go +++ b/libs/databrickscfg/loader.go @@ -106,6 +106,7 @@ const hostAttr = "host" var envAlwaysSkipAttrs = map[string]bool{ "workspace_id": true, "account_id": true, + "group_id": true, "auth_type": true, "discovery_url": true, "audience": true, diff --git a/libs/databrickscfg/loader_test.go b/libs/databrickscfg/loader_test.go index b1ab137d789..29bdff7595c 100644 --- a/libs/databrickscfg/loader_test.go +++ b/libs/databrickscfg/loader_test.go @@ -48,6 +48,7 @@ func TestResolveNonAuthFromEnvSkipsHostAndAuth(t *testing.T) { // workspace_id and account_id are routing identifiers, also skipped. t.Setenv("DATABRICKS_WORKSPACE_ID", "env-workspace") t.Setenv("DATABRICKS_ACCOUNT_ID", "env-account") + t.Setenv("DATABRICKS_GROUP_ID", "env-group") t.Setenv("DATABRICKS_CLUSTER_ID", "env-cluster") cfg := &config.Config{} @@ -63,6 +64,7 @@ func TestResolveNonAuthFromEnvSkipsHostAndAuth(t *testing.T) { assert.Empty(t, cfg.Cloud) assert.Empty(t, cfg.WorkspaceID) assert.Empty(t, cfg.AccountID) + assert.Empty(t, cfg.GroupID) // Non-auth attributes are still populated from the environment. assert.Equal(t, "env-cluster", cfg.ClusterID) } @@ -77,6 +79,7 @@ func TestProfileAuthLoadersProfileWinsOverSteeringEnv(t *testing.T) { t.Setenv("DATABRICKS_CLOUD", "azure") t.Setenv("DATABRICKS_WORKSPACE_ID", "env-workspace") t.Setenv("DATABRICKS_ACCOUNT_ID", "env-account") + t.Setenv("DATABRICKS_GROUP_ID", "env-group") cfg := config.Config{ Loaders: ProfileAuthLoaders, @@ -95,6 +98,21 @@ func TestProfileAuthLoadersProfileWinsOverSteeringEnv(t *testing.T) { assert.Empty(t, cfg.Cloud) assert.Empty(t, cfg.WorkspaceID) assert.Empty(t, cfg.AccountID) + assert.Empty(t, cfg.GroupID) +} + +func TestProfileAuthLoadersProfileGroupIDWinsOverEnv(t *testing.T) { + cfgFile := filepath.Join(t.TempDir(), ".databrickscfg") + require.NoError(t, os.WriteFile(cfgFile, []byte("[role]\nhost = https://workspace.test\ntoken = profile-token\ngroup_id = profile-group\n"), 0o600)) + t.Setenv("DATABRICKS_GROUP_ID", "env-group") + + cfg := config.Config{ + Loaders: ProfileAuthLoaders, + ConfigFile: cfgFile, + Profile: "role", + } + require.NoError(t, cfg.EnsureResolved()) + assert.Equal(t, "profile-group", cfg.GroupID) } func TestProfileAuthLoadersGapFillsHostFromEnv(t *testing.T) { diff --git a/libs/databrickscfg/ops_test.go b/libs/databrickscfg/ops_test.go index 480054bfed6..215155ffa20 100644 --- a/libs/databrickscfg/ops_test.go +++ b/libs/databrickscfg/ops_test.go @@ -570,6 +570,18 @@ func TestSaveToProfile_MergeSemantics(t *testing.T) { "scopes": "jobs,pipelines,clusters", }, }, + { + name: "writes group ID", + profile: "role", + saves: []saveOp{ + {cfg: &config.Config{Profile: "role", Host: "https://myworkspace.cloud.databricks.com", AuthType: "databricks-cli", GroupID: "group-123"}}, + }, + wantKeys: map[string]string{ + "host": "https://myworkspace.cloud.databricks.com", + "auth_type": "databricks-cli", + "group_id": "group-123", + }, + }, } for _, tc := range testCases { diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go index b7f6074c811..e970971e2b5 100644 --- a/libs/databrickscfg/profile/file.go +++ b/libs/databrickscfg/profile/file.go @@ -87,6 +87,7 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct ServerlessComputeID: all["serverless_compute_id"], HasClientCredentials: all["client_id"] != "" && all["client_secret"] != "", Scopes: all["scopes"], + GroupID: all["group_id"], AuthType: all["auth_type"], } if fn(profile) { diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go index efd358cd4e5..ca8ec9dc1d7 100644 --- a/libs/databrickscfg/profile/profile.go +++ b/libs/databrickscfg/profile/profile.go @@ -18,6 +18,7 @@ type Profile struct { ServerlessComputeID string HasClientCredentials bool Scopes string + GroupID string AuthType string }