Set necessary headers when authenticating via Azure CLI

config/auth_azure_cli.go

@@ -24,14 +24,14 @@ func (c AzureCliCredentials) Name() string {
 // implementing azureHostResolver for ensureWorkspaceUrl to work
 func (c AzureCliCredentials) tokenSourceFor(
 	ctx context.Context, cfg *Config, env azureEnvironment, resource string) oauth2.TokenSource {
-	return &azureCliTokenSource{resource: resource}
+	return &azureCliTokenSource{loginAppId: resource}
 }
 
 func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (func(*http.Request) error, error) {
 	if !cfg.IsAzure() {
 		return nil, nil
 	}
-	ts := azureCliTokenSource{cfg.getAzureLoginAppID()}
+	ts := azureCliTokenSource{cfg.getAzureLoginAppID(), cfg.AzureResourceID}
 	_, err := ts.Token()
 	if err != nil {
 		if strings.Contains(err.Error(), "No subscription found") {
@@ -50,11 +50,12 @@ 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
+	return azureVisitor(cfg.AzureResourceID, refreshableVisitor(&ts)), nil
 }
 
 type azureCliTokenSource struct {
-	resource string
+	loginAppId          string
+	workspaceResourceId string
 }
 
 type internalCliToken struct {
@@ -66,7 +67,7 @@ type internalCliToken struct {
 
 func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) {
 	out, err := exec.Command("az", "account", "get-access-token", "--resource",
-		ts.resource, "--output", "json").Output()
+		ts.loginAppId, "--output", "json").Output()
 	if ee, ok := err.(*exec.ExitError); ok {
 		return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr))
 	}
@@ -83,7 +84,7 @@ func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) {
 		return nil, fmt.Errorf("cannot parse expiry: %w", err)
 	}
 	logger.Infof(context.Background(), "Refreshed OAuth token for %s from Azure CLI, which expires on %s",
-		ts.resource, it.ExpiresOn)
+		ts.loginAppId, it.ExpiresOn)
 
 	var extra map[string]interface{}
 	err = json.Unmarshal(out, &extra)

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,22 @@ 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"))
+}
+
+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, "X-Databricks-Azure-SP-Management-Token")), 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, "X-Databricks-Azure-SP-Management-Token")), nil
 }
 
 func (c AzureMsiCredentials) getInstanceEnvironment(ctx context.Context) (*azureEnvironment, error) {

config/oauth_visitors.go

@@ -56,3 +56,12 @@ func refreshableVisitor(inner oauth2.TokenSource) func(r *http.Request) error {
 		return nil
 	}
 }
+
+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)
+		}
+		return inner(r)
+	}
+}