Centrally determine Azure Environment

databricks · Aug 30, 2021 · 5362ab0 · 5362ab0
1 parent d0af809
commit 5362ab0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 18 deletions.
diff --git a/common/azure_auth.go b/common/azure_auth.go
@@ -91,11 +91,6 @@ func (aa *DatabricksClient) configureWithAzureClientSecret(ctx context.Context)
 	if !aa.IsAzureClientSecretSet() {
 		return nil, nil
 	}
-	azureEnvironment, err := aa.getAzureEnvironment()
-	if err != nil {
-		return nil, fmt.Errorf("cannot get azure environment: %w", err)
-	}
-	aa.AzureEnvironment = &azureEnvironment
 	log.Printf("[INFO] Using Azure Service Principal client secret authentication")
 	if aa.AzureUsePATForSPN {
 		log.Printf("[INFO] Generating PAT token Azure Service Principal client secret authentication")
@@ -120,11 +115,6 @@ func (aa *DatabricksClient) configureWithAzureManagedIdentity(ctx context.Contex
 	if !adal.MSIAvailable(ctx, aa.httpClient.HTTPClient) {
 		return nil, nil
 	}
-	azureEnvironment, err := aa.getAzureEnvironment()
-	if err != nil {
-		return nil, fmt.Errorf("cannot get azure environment: %w", err)
-	}
-	aa.AzureEnvironment = &azureEnvironment
 	log.Printf("[INFO] Using Azure Managed Identity authentication")
 	return aa.simpleAADRequestVisitor(ctx, func(resource string) (autorest.Authorizer, error) {
 		return auth.MSIConfig{
@@ -163,7 +153,12 @@ func (aa *DatabricksClient) simpleAADRequestVisitor(
 	ctx context.Context,
 	authorizerFactory func(resource string) (autorest.Authorizer, error),
 	visitors ...func(r *http.Request, ma autorest.Authorizer) error) (func(r *http.Request) error, error) {
-	managementAuthorizer, err := authorizerFactory(aa.AzureEnvironment.ServiceManagementEndpoint)
+	azureEnvironment, err := aa.getAzureEnvironment()
+	if err != nil {
+		return nil, fmt.Errorf("cannot get azure environment: %w", err)
+	}
+	aa.AzureEnvironment = &azureEnvironment
+	managementAuthorizer, err := authorizerFactory(azureEnvironment.ServiceManagementEndpoint)
 	if err != nil {
 		return nil, fmt.Errorf("cannot authorize management: %w", err)
 	}
@@ -205,7 +200,12 @@ func (aa *DatabricksClient) acquirePAT(
 	if aa.temporaryPat != nil {
 		return aa.temporaryPat, nil
 	}
-	management, err := factory(aa.AzureEnvironment.ServiceManagementEndpoint)
+	azureEnvironment, err := aa.getAzureEnvironment()
+	if err != nil {
+		return nil, fmt.Errorf("cannot get azure environment: %w", err)
+	}
+	aa.AzureEnvironment = &azureEnvironment
+	management, err := factory(azureEnvironment.ServiceManagementEndpoint)
 	if err != nil {
 		return nil, err
 	}

diff --git a/common/azure_cli_auth.go b/common/azure_cli_auth.go
@@ -100,13 +100,8 @@ func (aa *DatabricksClient) configureWithAzureCLI(ctx context.Context) (func(*ht
 	if aa.IsAzureClientSecretSet() {
 		return nil, nil
 	}
-	azureEnvironment, err := aa.getAzureEnvironment()
-	if err != nil {
-		return nil, fmt.Errorf("cannot get environment: %w", err)
-	}
-	aa.AzureEnvironment = &azureEnvironment
 	// verify that Azure CLI is authenticated
-	_, err = cli.GetTokenFromCLI(AzureDatabricksResourceID)
+	_, err := cli.GetTokenFromCLI(AzureDatabricksResourceID)
 	if err != nil {
 		if err.Error() == "Invoking Azure CLI failed with the following error: " {
 			return nil, fmt.Errorf("most likely Azure CLI is not installed. " +