Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fallback to account-level auth if possible when using CLI auth #943

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

mgyucht
Copy link
Contributor

@mgyucht mgyucht commented Jun 7, 2024

Changes

Within Databricks, it is possible to authenticate to a workspace using an account-level OAuth token. However, in the CLI today, OAuth tokens are stored on per-oauth-endpoint basis: one token is stored per account and one per workspace. There is currently no way to identify the account for a given workspace via the REST API.

This change allows the SDK to attempt to load an OAuth token at both the account and workspace level if the account ID is configured in DatabricksConfig when trying to login to a workspace. The initial request remains the same (try to get the OAuth token for the workspace). If this fails, however, and the account ID is configured, the SDK then makes a second request to get the OAuth token for the account. If this exists, this token can be used to interact with the workspace.

This is useful for tools like Terraform. When users authenticate via U2M auth and then apply a Terraform template containing databricks_mws_workspaces resources, the apply step fails because the CLI doesn't have an OAuth token stored for the newly minted workspace. However, after this change, it will fallback to the account-level token and succeed.

Tests

  • Unit tests to verify that the SDK falls back to the account-scoped token
  • Manual test: removed the workspace-level token for a workspace from ~/.databricks/token-cache.json, then tried to authenticate to that workspace while also specifying its account ID in the config. It succeeded.

@codecov-commenter
Copy link

Codecov Report

Attention: Patch coverage is 50.00000% with 18 lines in your changes missing coverage. Please review.

Project coverage is 7.13%. Comparing base (094bca8) to head (7ab4ef0).

Files Patch % Lines
config/auth_databricks_cli.go 44.00% 13 Missing and 1 partial ⚠️
common/environment/environments.go 0.00% 2 Missing ⚠️
account_client.go 0.00% 1 Missing ⚠️
config/config.go 85.71% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@          Coverage Diff          @@
##            main    #943   +/-   ##
=====================================
  Coverage   7.13%   7.13%           
=====================================
  Files        281     281           
  Lines      64633   64652   +19     
=====================================
+ Hits        4609    4615    +6     
- Misses     59715   59728   +13     
  Partials     309     309           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants