Make Azure MSI auth account compatible (#544)

tanmay-db · web-flow · commit e80f6e16ffa1 · 2023-07-10T13:14:20.000+02:00
## Changes
&lt;!-- Summary of your changes that are easy to understand --&gt;
Currently it is not possible to use azure msi authentication for
accounts. This PR adds that functionality

## Tests
&lt;!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
--&gt;
Existing integration tests passing for azure-prod and azure-prod-acct
Need to test manually for MSI accounts auth
- [x] `make test` passing
- [x] `make fmt` applied
- [ ] relevant integration tests applied
diff --git a/config/auth_azure_msi.go b/config/auth_azure_msi.go
@@ -27,16 +27,18 @@ func (c AzureMsiCredentials) Name() string {
 }
 
 func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (func(*http.Request) error, error) {
-	if !cfg.IsAzure() || !cfg.AzureUseMSI || cfg.AzureResourceID == "" {
+	if !cfg.IsAzure() || !cfg.AzureUseMSI || (cfg.AzureResourceID == "" && !cfg.IsAccountClient()) {
 		return nil, nil
 	}
 	env, err := c.getInstanceEnvironment(ctx)
 	if err != nil {
 		return nil, err
 	}
-	err = cfg.azureEnsureWorkspaceUrl(ctx, c)
-	if err != nil {
-		return nil, fmt.Errorf("resolve host: %w", err)
+	if !cfg.IsAccountClient() {
+		err = cfg.azureEnsureWorkspaceUrl(ctx, c)
+		if err != nil {
+			return nil, fmt.Errorf("resolve host: %w", err)
+		}
 	}
 	logger.Debugf(ctx, "Generating AAD token via Azure MSI")
 	inner := azureMsiTokenSource{
@@ -48,7 +50,9 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 		clientId: cfg.AzureClientID,
 	}
 	return func(r *http.Request) error {
-		r.Header.Set("X-Databricks-Azure-Workspace-Resource-Id", cfg.AzureResourceID)
+		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
