databricks-sdk-go

databricks-sdk-go copied to clipboard

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

• [x] Unit tests to verify that the SDK falls back to the account-scoped token
• [x] 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.