Set necessary headers when authenticating via Azure CLI #584

mgyucht · 2023-08-17T07:52:44Z

Changes

The Go SDK request authentication logic is inconsistent between the Azure login types: for service principal & MSI auth, the SDK correctly adds the X-Databricks-Azure-Workspace-Resource-Id when configured, but this is missed for Azure CLI auth. Additionally, when logging in via Azure CLI using a service principal, the service management token must also be fetched from the CLI. This caused a regression for the Terraform provider: databricks/terraform-provider-databricks#2590.

This PR fixes this by defining the logic to attach these header in a common function that is used by all Azure-specific authentication types.

Tests

Added a unit test to ensure the header is being set for Azure CLI login
Made a test app that uses azure-cli to login and verified that the correct header was set on the request:

...
> * X-Databricks-Azure-Workspace-Resource-Id: /subscriptions/<REDACTED>/resourceGroups/<REDACTED>/pr... (63 more bytes)

make test passing
make fmt applied
relevant integration tests applied

…equests when AzureResourceId is configured

codecov-commenter · 2023-08-17T07:55:10Z

Codecov Report

Patch coverage: 83.33% and project coverage change: +0.22% 🎉

Comparison is base (d682ad8) 18.24% compared to head (8ef1fd1) 18.47%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #584      +/-   ##
==========================================
+ Coverage   18.24%   18.47%   +0.22%     
==========================================
  Files          85       85              
  Lines        9865     9870       +5     
==========================================
+ Hits         1800     1823      +23     
+ Misses       7910     7888      -22     
- Partials      155      159       +4

Files Changed	Coverage Δ
config/auth_azure_client_secret.go	`12.50% <0.00%> (+1.78%)`	⬆️
config/auth_azure_msi.go	`18.39% <0.00%> (+0.80%)`	⬆️
config/auth_azure_cli.go	`86.79% <81.81%> (-1.85%)`	⬇️
config/oauth_visitors.go	`80.64% <100.00%> (+32.49%)`	⬆️

... and 1 file with indirect coverage changes

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

alexott · 2023-08-17T08:05:08Z

config/oauth_visitors.go

+func azureVisitor(workspaceResourceId string, inner func(*http.Request) error) func(*http.Request) error {
+	return func(r *http.Request) error {
+		if workspaceResourceId != "" {
+			r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", workspaceResourceId)


Besides this header, we may need to add service management header as well: https://github.com/databricks/terraform-provider-databricks/blob/v1.9.2/common/azure_auth.go#L147

Isn't that already added by the serviceToServiceVisitor?

The header in question: X-Databricks-Azure-SP-Management-Token.

After looking at this original change, we realized that there was a secondary problem: if a user logs in via the CLI using a service principal, the CLI auth type needs to also request the management token from the CLI in addition to the Databricks token. The point is: all Azure-native auth types need to call azureVisitor, and all auth types that need to include a second token in the request need to call serviceToServiceVisitor; this now includes the CLI.

config/auth_azure_cli.go

config/oauth_visitors.go

pietern · 2023-08-17T11:15:25Z

config/oauth_visitors.go

+func azureVisitor(workspaceResourceId string, inner func(*http.Request) error) func(*http.Request) error {
+	return func(r *http.Request) error {
+		if workspaceResourceId != "" {
+			r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", workspaceResourceId)


Isn't that already added by the serviceToServiceVisitor?

The header in question: X-Databricks-Azure-SP-Management-Token.

config/auth_azure_cli.go

pietern · 2023-08-17T11:54:38Z

config/oauth_visitors.go

+// serviceToServiceVisitor returns a visitor that sets the Authorization header to the token from the auth token source
+// and the provided secondary header to the token from the secondary token source. If secondary is nil, the secondary
+// header is not set.
+func serviceToServiceVisitor(auth, secondary oauth2.TokenSource, secondaryHeader string) func(r *http.Request) error {


In the current approach I believe the secondary one is never nil anymore.

The comment still mentions nil secondary.

## Changes The Python SDK request authentication logic is inconsistent between the Azure login types: for service principal auth, the SDK correctly adds the X-Databricks-Azure-Workspace-Resource-Id when configured, but this is missed for Azure CLI auth. This PR fixes this by defining the logic to attach this header in a common function that is used by all Azure-specific authentication types. See databricks/databricks-sdk-go#584 for the same change in Go SDK. ## Tests - [x] Added a unit test to ensure the header is being set for Azure CLI login - [ ] `make test` run locally - [ ] `make fmt` applied - [ ] relevant integration tests applied

alexott

lgtm

pietern · 2023-08-17T12:32:11Z

config/oauth_visitors.go

+// serviceToServiceVisitor returns a visitor that sets the Authorization header to the token from the auth token source
+// and the provided secondary header to the token from the secondary token source. If secondary is nil, the secondary
+// header is not set.
+func serviceToServiceVisitor(auth, secondary oauth2.TokenSource, secondaryHeader string) func(r *http.Request) error {


The comment still mentions nil secondary.

## Changes The Java SDK request authentication logic is inconsistent between the Azure login types: for service principal auth, the SDK correctly adds the X-Databricks-Azure-Workspace-Resource-Id when configured, but this is missed for Azure CLI auth. Additionally, when logging in via Azure CLI using a service principal, the service management token must also be fetched from the CLI. This PR fixes this by defining the logic to attach these header in a common function that is used by all Azure-specific authentication types. See databricks/databricks-sdk-go#584 for the same change in the Go SDK. See databricks/databricks-sdk-py#290 for the same changes in the Python SDK. ## Tests - [x] Unit tests to cover the two scenarios for Azure CLI w.r.t. management endpoint token fetching, and one to verify that X-Databricks-Azure-Workspace-Resource-Id is included when using Azure CLI.

mgyucht added 2 commits August 17, 2023 09:35

Added X-Databricks-Azure-Workspace-Resource-Id header for all Azure r…

119ab99

…equests when AzureResourceId is configured

unit test

85d2f84

alexott reviewed Aug 17, 2023

View reviewed changes

mgyucht mentioned this pull request Aug 17, 2023

Set necessary headers when authenticating via Azure CLI databricks/databricks-sdk-py#290

Merged

4 tasks

mgyucht added 7 commits August 17, 2023 11:31

Always fetch service management token

8fb2b10

try to get service management token up front

1e5806d

undo rename

95c005f

better ordering

88cabdc

old order

ae6f57e

fix name of header

a2f3e35

add tests

478bf2e

mgyucht requested review from alexott and pietern and removed request for alexott August 17, 2023 10:50

pietern reviewed Aug 17, 2023

View reviewed changes

mgyucht mentioned this pull request Aug 17, 2023

Set necessary headers when authenticating via Azure CLI databricks/databricks-sdk-java#136

Merged

1 task

mgyucht changed the title ~~Added X-Databricks-Azure-Workspace-Resource-Id header for all Azure requests when AzureResourceId is configured~~ Set necessary headers when authenticating via Azure CLI Aug 17, 2023

Address comments

d23e2f4

pietern reviewed Aug 17, 2023

View reviewed changes

address last comments

8ef1fd1

alexott approved these changes Aug 17, 2023

View reviewed changes

pietern approved these changes Aug 17, 2023

View reviewed changes

last tweak to comment

d53740a

mgyucht enabled auto-merge August 17, 2023 13:12

mgyucht added this pull request to the merge queue Aug 17, 2023

Merged via the queue into main with commit 113aac1 Aug 17, 2023
3 checks passed

mgyucht deleted the issue-2590 branch August 17, 2023 13:16

mgyucht mentioned this pull request Aug 17, 2023

Release v0.17.0 #585

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Set necessary headers when authenticating via Azure CLI #584

Set necessary headers when authenticating via Azure CLI #584

mgyucht commented Aug 17, 2023 •

edited

codecov-commenter commented Aug 17, 2023 •

edited

alexott Aug 17, 2023

pietern Aug 17, 2023

mgyucht Aug 17, 2023

pietern Aug 17, 2023

pietern Aug 17, 2023

pietern Aug 17, 2023

alexott left a comment

pietern Aug 17, 2023

Set necessary headers when authenticating via Azure CLI #584

Set necessary headers when authenticating via Azure CLI #584

Conversation

mgyucht commented Aug 17, 2023 • edited

Changes

Tests

codecov-commenter commented Aug 17, 2023 • edited

Codecov Report

alexott Aug 17, 2023

Choose a reason for hiding this comment

pietern Aug 17, 2023

Choose a reason for hiding this comment

mgyucht Aug 17, 2023

Choose a reason for hiding this comment

pietern Aug 17, 2023

Choose a reason for hiding this comment

pietern Aug 17, 2023

Choose a reason for hiding this comment

pietern Aug 17, 2023

Choose a reason for hiding this comment

alexott left a comment

Choose a reason for hiding this comment

pietern Aug 17, 2023

Choose a reason for hiding this comment

mgyucht commented Aug 17, 2023 •

edited

codecov-commenter commented Aug 17, 2023 •

edited