Set necessary headers when authenticating via Azure CLI

config/auth_azure_cli.go

@@ -14,6 +14,9 @@ import (
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 
+// The header used to pass the workspace resource ID to the Databricks backend.
+const XDatabricksAzureSpManagementToken = "X-Databricks-Azure-SP-Management-Token"
+
 type AzureCliCredentials struct {
 }
 
@@ -31,7 +34,8 @@ func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 	if !cfg.IsAzure() {
 		return nil, nil
 	}
-	ts := azureCliTokenSource{cfg.getAzureLoginAppID()}
+	// Eagerly get a token to fail fast in case the user is not logged in with the Azure CLI.
+	ts := &azureCliTokenSource{cfg.getAzureLoginAppID(), cfg.AzureResourceID}
 	_, err := ts.Token()
 	if err != nil {
 		if strings.Contains(err.Error(), "No subscription found") {
@@ -50,11 +54,32 @@ func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 		return nil, fmt.Errorf("resolve host: %w", err)
 	}
 	logger.Infof(ctx, "Using Azure CLI authentication with AAD tokens")
-	return refreshableVisitor(&ts), nil
+
+	// There are three scenarios:
+	//
+	//  1. The user has logged in with the Azure CLI as a user and has access to the service management endpoint.
+	//  2. The user has logged in with the Azure CLI as a user and does not have access to the service management endpoint.
+	//  3. The user has logged in with the Azure CLI as a service principal, and must have access to the service management
+	//     endpoint to authenticate.
+	//
+	// If the user can't access the service management endpoint, we assume they are in case 2 and do not pass the service
+	// management token. Otherwise, we always pass the service management token.
+	var managementTs oauth2.TokenSource
+	env, err := cfg.GetAzureEnvironment()
+	if err != nil {
+		return nil, err
+	}
+	managementTs = &azureCliTokenSource{env.ServiceManagementEndpoint, ""}
+	_, err = managementTs.Token()
+	if err != nil {
+		managementTs = nil
+	}
+	return azureVisitor(cfg.AzureResourceID, serviceToServiceVisitor(ts, managementTs, XDatabricksAzureSpManagementToken, true)), nil
 }
 
 type azureCliTokenSource struct {
-	resource string
+	resource            string
+	workspaceResourceId string
 }
 
 type internalCliToken struct {

config/auth_azure_cli_test.go

@@ -14,6 +14,7 @@ import (
 )
 
 var azDummy = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/"}
+var azDummyWithResourceId = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/", AzureResourceID: "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123"}
 
 // testdataPath returns the PATH to use for the duration of a test.
 // It must only return absolute directories because Go refuses to run
@@ -67,6 +68,40 @@ func TestAzureCliCredentials_Valid(t *testing.T) {
 	assert.NoError(t, err)
 
 	assert.Equal(t, "Bearer ...", r.Header.Get("Authorization"))
+	assert.Equal(t, "", r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id"))
+	assert.Equal(t, "...", r.Header.Get("X-Databricks-Azure-SP-Management-Token"))
+}
+
+func TestAzureCliCredentials_ValidNoManagementAccess(t *testing.T) {
+	env.CleanupEnvironment(t)
+	os.Setenv("PATH", testdataPath())
+	os.Setenv("FAIL_IF", "https://management.core.windows.net/")
+	aa := AzureCliCredentials{}
+	visitor, err := aa.Configure(context.Background(), azDummy)
+	assert.NoError(t, err)
+
+	r := &http.Request{Header: http.Header{}}
+	err = visitor(r)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "Bearer ...", r.Header.Get("Authorization"))
+	assert.Equal(t, "", r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id"))
+	assert.Equal(t, "", r.Header.Get("X-Databricks-Azure-SP-Management-Token"))
+}
+
+func TestAzureCliCredentials_ValidWithAzureResourceId(t *testing.T) {
+	env.CleanupEnvironment(t)
+	os.Setenv("PATH", testdataPath())
+	aa := AzureCliCredentials{}
+	visitor, err := aa.Configure(context.Background(), azDummyWithResourceId)
+	assert.NoError(t, err)
+
+	r := &http.Request{Header: http.Header{}}
+	err = visitor(r)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "Bearer ...", r.Header.Get("Authorization"))
+	assert.Equal(t, azDummyWithResourceId.AzureResourceID, r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id"))
 }
 
 func TestAzureCliCredentials_AlwaysExpired(t *testing.T) {

config/auth_azure_client_secret.go

@@ -54,11 +54,5 @@ func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config
 	refreshCtx := context.Background()
 	inner := c.tokenSourceFor(refreshCtx, cfg, env, cfg.getAzureLoginAppID())
 	platform := c.tokenSourceFor(refreshCtx, cfg, env, env.ServiceManagementEndpoint)
-	return func(r *http.Request) error {
-		if cfg.AzureResourceID != "" {
-			r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", cfg.AzureResourceID)
-		}
-		return serviceToServiceVisitor(inner, platform,
-			"X-Databricks-Azure-SP-Management-Token")(r)
-	}, nil
+	return azureVisitor(cfg.AzureResourceID, serviceToServiceVisitor(inner, platform, XDatabricksAzureSpManagementToken, false)), nil
 }

config/auth_azure_msi.go

@@ -49,13 +49,7 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 		resource: env.ServiceManagementEndpoint,
 		clientId: cfg.AzureClientID,
 	}
-	return func(r *http.Request) error {
-		if !cfg.IsAccountClient() {
-			r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", cfg.AzureResourceID)
-		}
-		return serviceToServiceVisitor(inner, platform,
-			"X-Databricks-Azure-SP-Management-Token")(r)
-	}, nil
+	return azureVisitor(cfg.AzureResourceID, serviceToServiceVisitor(inner, platform, XDatabricksAzureSpManagementToken, false)), nil
 }
 
 func (c AzureMsiCredentials) getInstanceEnvironment(ctx context.Context) (*azureEnvironment, error) {

config/auth_gcp_google_credentials.go

@@ -42,7 +42,7 @@ func (c GoogleCredentials) Configure(ctx context.Context, cfg *Config) (func(*ht
 		return nil, fmt.Errorf("could not obtain OAuth2 token from JSON: %w", err)
 	}
 	logger.Infof(ctx, "Using Google Credentials")
-	return serviceToServiceVisitor(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token"), nil
+	return serviceToServiceVisitor(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token", false), nil
 }
 
 // Reads credentials as JSON. Credentials can be either a path to JSON file,

config/auth_gcp_google_id.go

@@ -44,7 +44,7 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (f
 		return nil, err
 	}
 	logger.Infof(ctx, "Using Google Default Application Credentials for Accounts API")
-	return serviceToServiceVisitor(inner, platform, "X-Databricks-GCP-SA-Access-Token"), nil
+	return serviceToServiceVisitor(inner, platform, "X-Databricks-GCP-SA-Access-Token", false), nil
 }
 
 func (c GoogleDefaultCredentials) idTokenSource(ctx context.Context, host, serviceAccount string,

config/oauth_visitors.go

@@ -27,32 +27,44 @@ func retriableTokenSource(ctx context.Context, ts oauth2.TokenSource) (*oauth2.T
 	})
 }
 
-func serviceToServiceVisitor(inner, cloud oauth2.TokenSource, header string) func(r *http.Request) error {
-	refreshableInner := oauth2.ReuseTokenSource(nil, inner)
-	refreshableCloud := oauth2.ReuseTokenSource(nil, cloud)
+// 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. If tolerateSecondaryFailure is true, the visitor will not return an error if the cloud token source fails.
+func serviceToServiceVisitor(auth, secondary oauth2.TokenSource, secondaryHeader string, tolerateSecondaryFailure bool) func(r *http.Request) error {
+	refreshableAuth := oauth2.ReuseTokenSource(nil, auth)
+	var refreshableSecondary oauth2.TokenSource
+	if secondary != nil {
+		refreshableSecondary = oauth2.ReuseTokenSource(nil, secondary)
+	}
 	return func(r *http.Request) error {
-		inner, err := retriableTokenSource(r.Context(), refreshableInner)
+		inner, err := retriableTokenSource(r.Context(), refreshableAuth)
 		if err != nil {
 			return fmt.Errorf("inner token: %w", err)
 		}
 		inner.SetAuthHeader(r)
-		cloud, err := retriableTokenSource(r.Context(), refreshableCloud)
-		if err != nil {
+
+		if refreshableSecondary == nil {
+			return nil
+		}
+		cloud, err := retriableTokenSource(r.Context(), refreshableSecondary)
+		if err == nil {
+			r.Header.Set(secondaryHeader, cloud.AccessToken)
+		} else if !tolerateSecondaryFailure {
 			return fmt.Errorf("cloud token: %w", err)
 		}
-		r.Header.Set(header, cloud.AccessToken)
 		return nil
 	}
 }
 
 func refreshableVisitor(inner oauth2.TokenSource) func(r *http.Request) error {
-	refreshableInner := oauth2.ReuseTokenSource(nil, inner)
+	return serviceToServiceVisitor(inner, nil, "", false)
+}
+
+func azureVisitor(workspaceResourceId string, inner func(*http.Request) error) func(*http.Request) error {
 	return func(r *http.Request) error {
-		inner, err := retriableTokenSource(r.Context(), refreshableInner)
-		if err != nil {
-			return fmt.Errorf("inner token: %w", err)
+		if workspaceResourceId != "" {
+			r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", workspaceResourceId)
 		}
-		inner.SetAuthHeader(r)
-		return nil
+		return inner(r)
 	}
 }

config/testdata/az

@@ -15,6 +15,13 @@ if [ "corrupt" == "$FAIL" ]; then
     exit
 fi
 
+for arg in "$@"; do
+    if [[ "$arg" == "$FAIL_IF" ]]; then
+        echo "Failed"
+        exit 1
+    fi
+done
+
 # Macos
 EXP="$(/bin/date -v+${EXPIRE:=10S} +'%F %T' 2>/dev/null)"
 if [ -z "${EXP}" ]; then