Handle Azure authentication when WorkspaceResourceID is provided

config/auth_azure_cli.go

@@ -95,17 +95,49 @@ type internalCliToken struct {
 	ExpiresOn    string `json:"expiresOn"`
 }
 
+func (ts *azureCliTokenSource) getSubscription() string {
+	if ts.workspaceResourceId == "" {
+		return ""
+	}
+	components := strings.Split(ts.workspaceResourceId, "/")
+	if len(components) < 3 {
+		logger.Warnf(context.Background(), "Invalid azure workspace resource ID")
+		return ""
+	}
+	return components[2]
+}
+
 func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) {
-	out, err := exec.Command("az", "account", "get-access-token", "--resource",
-		ts.resource, "--output", "json").Output()
-	if ee, ok := err.(*exec.ExitError); ok {
-		return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr))
+	subscription := ts.getSubscription()
+	var out []byte
+	args := []string{"account", "get-access-token", "--resource",
+		ts.resource, "--output", "json"}
+	if subscription != "" {
+		extendedArgs := make([]string, len(args))
+		copy(extendedArgs, args)
+		extendedArgs = append(extendedArgs, "--subscription", subscription)
+		// This will fail if the user has access to the workspace, but not to the subscription
+		// itself.
+		// In such case, we fall back to not using the subscription.
+		result, err := exec.Command("az", extendedArgs...).Output()
+		if err != nil {
+			logger.Warnf(context.Background(), "Failed to get token for subscription. Using resource only token.")
+		} else {
+			out = result
+		}
 	}
-	if err != nil {
-		return nil, fmt.Errorf("cannot get access token: %v", err)
+	if out == nil {
+		result, err := exec.Command("az", args...).Output()
+		if ee, ok := err.(*exec.ExitError); ok {
+			return nil, fmt.Errorf("cannot get access token: %s", string(ee.Stderr))
+		}
+		if err != nil {
+			return nil, fmt.Errorf("cannot get access token: %v", err)
+		}
+		out = result
 	}
 	var it internalCliToken
-	err = json.Unmarshal(out, &it)
+	err := json.Unmarshal(out, &it)
 	if err != nil {
 		return nil, fmt.Errorf("cannot unmarshal CLI result: %w", err)
 	}

config/auth_azure_cli_test.go

@@ -15,6 +15,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"}
+var azDummyWitInvalidResourceId = &Config{Host: "https://adb-xyz.c.azuredatabricks.net/", AzureResourceID: "invalidResourceId"}
 
 // testdataPath returns the PATH to use for the duration of a test.
 // It must only return absolute directories because Go refuses to run
@@ -104,6 +105,22 @@ func TestAzureCliCredentials_ValidWithAzureResourceId(t *testing.T) {
 	assert.Equal(t, azDummyWithResourceId.AzureResourceID, r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id"))
 }
 
+func TestAzureCliCredentials_Fallback(t *testing.T) {
+	env.CleanupEnvironment(t)
+	os.Setenv("PATH", testdataPath())
+	aa := AzureCliCredentials{}
+	visitor, err := aa.Configure(context.Background(), azDummyWitInvalidResourceId)
+	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, azDummyWitInvalidResourceId.AzureResourceID, r.Header.Get("X-Databricks-Azure-Workspace-Resource-Id"))
+	assert.Equal(t, "...", r.Header.Get("X-Databricks-Azure-SP-Management-Token"))
+}
+
 func TestAzureCliCredentials_AlwaysExpired(t *testing.T) {
 	env.CleanupEnvironment(t)
 	os.Setenv("PATH", testdataPath())