added support for other azure cloud environments

CHANGELOG.md

@@ -2,10 +2,12 @@
 
 ## 0.3.0
 
+* Added optional parameter azure_environment to provider config which defaults to public ([#437](https://github.com/databrickslabs/terraform-provider-databricks/pull/437)).
 * Added [databricks_service_principal](https://github.com/databrickslabs/terraform-provider-databricks/pull/386) resource.
 
 **Behavior changes**
 
+* Added optional parameter `azure_environment` to provider config which defaults to `public`.
 * Removed deprecated `library_jar`, `library_egg`, `library_whl`, `library_pypi`, `library_cran`, and `library_maven` from `databricks_cluster` and `databricks_job` in favor of more API-transparent [library](https://registry.terraform.io/providers/databrickslabs/databricks/latest/docs/resources/cluster#library-configuration-block) configuration block.
 * Removed deprecated `notebook_path` and `notebook_base_parameters` from `databricks_job` in favor of [notebook_task](https://registry.terraform.io/providers/databrickslabs/databricks/latest/docs/resources/job#notebook_task-configuration-block) configuration block.
 * Removed deprecated `jar_uri`, `jar_main_class_name`, and `jar_parameters` from `databricks_job` in favor of [spark_jar_task](https://registry.terraform.io/providers/databrickslabs/databricks/latest/docs/resources/job#spark_jar_task-configuration-block) configuration block.

common/azure_auth.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"net/http"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -32,6 +33,7 @@ type AzureAuth struct {
 	ClientSecret string
 	ClientID     string
 	TenantID     string
+	Environment  string
 
 	// temporary workaround for SP-based auth
 	PATTokenDurationSeconds string
@@ -65,6 +67,21 @@ type TokenInfo struct {
 
 var authorizerMutex sync.Mutex
 
+func (aa *AzureAuth) getAzureEnvironment() (azure.Environment, error) {
+	if aa.Environment == "" {
+		return azure.PublicCloud, nil
+	}
+
+	envName := fmt.Sprintf("AZURE%sCLOUD", strings.ToUpper(aa.Environment))
+	env, err := azure.EnvironmentFromName(envName)
+
+	if err != nil {
+		return env, err
+	}
+
+	return env, nil
+}
+
 func (aa *AzureAuth) resourceID() string {
 	if aa.ResourceID != "" {
 		if aa.SubscriptionID == "" {
@@ -151,7 +168,11 @@ func (aa *AzureAuth) addSpManagementTokenVisitor(r *http.Request, management aut
 func (aa *AzureAuth) simpleAADRequestVisitor(
 	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(azure.PublicCloud.ServiceManagementEndpoint)
+	env, err := aa.getAzureEnvironment()
+	if err != nil {
+		return nil, err
+	}
+	managementAuthorizer, err := authorizerFactory(env.ServiceManagementEndpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -191,7 +212,11 @@ func (aa *AzureAuth) acquirePAT(
 	if aa.temporaryPat != nil {
 		return aa.temporaryPat, nil
 	}
-	management, err := factory(azure.PublicCloud.ServiceManagementEndpoint)
+	env, err := aa.getAzureEnvironment()
+	if err != nil {
+		return nil, err
+	}
+	management, err := factory(env.ServiceManagementEndpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -292,6 +317,10 @@ func (aa *AzureAuth) getClientSecretAuthorizer(resource string) (autorest.Author
 		// todo: probably should be two different ones...
 		return aa.authorizer, nil
 	}
+	env, err := aa.getAzureEnvironment()
+	if err != nil {
+		return nil, err
+	}
 	if resource != AzureDatabricksResourceID {
 		es := auth.EnvironmentSettings{
 			Values: map[string]string{
@@ -300,12 +329,12 @@ func (aa *AzureAuth) getClientSecretAuthorizer(resource string) (autorest.Author
 				auth.TenantID:     aa.TenantID,
 				auth.Resource:     resource,
 			},
-			Environment: azure.PublicCloud,
+			Environment: env,
 		}
 		return es.GetAuthorizer()
 	}
 	platformTokenOAuthCfg, err := adal.NewOAuthConfigWithAPIVersion(
-		azure.PublicCloud.ActiveDirectoryEndpoint,
+		env.ActiveDirectoryEndpoint,
 		aa.TenantID,
 		nil)
 	if err != nil {

common/azure_auth_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -269,3 +270,57 @@ func TestAzureAuth_configureWithClientSecret(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Len(t, zi.Zones, 3)
 }
+
+func TestAzureEnvironment(t *testing.T) {
+	aa := AzureAuth{}
+	env, err := aa.getAzureEnvironment()
+
+	assert.Nil(t, err)
+	assert.Equal(t, azure.PublicCloud, env)
+
+	aa.Environment = "public"
+	env, err = aa.getAzureEnvironment()
+	assert.Nil(t, err)
+	assert.Equal(t, azure.PublicCloud, env)
+
+	aa.Environment = "china"
+	env, err = aa.getAzureEnvironment()
+	assert.Nil(t, err)
+	assert.Equal(t, azure.ChinaCloud, env)
+
+	aa.Environment = "german"
+	env, err = aa.getAzureEnvironment()
+	assert.Nil(t, err)
+	assert.Equal(t, azure.GermanCloud, env)
+
+	aa.Environment = "usgovernment"
+	env, err = aa.getAzureEnvironment()
+	assert.Nil(t, err)
+	assert.Equal(t, azure.USGovernmentCloud, env)
+
+	aa.Environment = "xyzdummy"
+	_, err = aa.getAzureEnvironment()
+	assert.NotNil(t, err)
+}
+
+func TestInvalidAzureEnvironment(t *testing.T) {
+	aa := AzureAuth{}
+
+	aa.Environment = "xyzdummy"
+	_, envErr := aa.getAzureEnvironment()
+	assert.NotNil(t, envErr)
+
+	mockFunc := func(resource string) (autorest.Authorizer, error) {
+		return nil, nil
+	}
+
+	_, err := aa.simpleAADRequestVisitor(mockFunc)
+
+	assert.Equal(t, envErr, err)
+
+	_, err = aa.acquirePAT(mockFunc)
+	assert.Equal(t, envErr, err)
+
+	_, err = aa.getClientSecretAuthorizer("")
+	assert.Equal(t, envErr, err)
+}

common/env.go

@@ -32,6 +32,7 @@ func NewClientFromEnvironment() *DatabricksClient {
 			ClientID:       os.Getenv("ARM_CLIENT_ID"),
 			ClientSecret:   os.Getenv("ARM_CLIENT_SECRET"),
 			TenantID:       os.Getenv("ARM_TENANT_ID"),
+			Environment:    os.Getenv("ARM_ENVIRONMENT"),
 		},
 	}
 	err := client.Configure()

docs/index.md

@@ -185,6 +185,7 @@ resource "databricks_scim_user" "my-user" {
 * `azure_client_id` - (optional) This is the Azure Enterprise Application (Service principal) client id. This service principal requires contributor access to your Azure Databricks deployment. Alternatively, you can provide this value as an environment variable `DATABRICKS_AZURE_CLIENT_ID` or `ARM_CLIENT_ID`.
 * `azure_tenant_id` - (optional) This is the Azure Active Directory Tenant id in which the Enterprise Application (Service Principal) 
 resides. Alternatively, you can provide this value as an environment variable `DATABRICKS_AZURE_TENANT_ID` or `ARM_TENANT_ID`.
+* `azure_environment` - (optional) This is the Azure Environment which defaults to the `public` cloud. Other options are `german`, `china` and `usgovernment`. Alternatively, you can provide this value as an environment variable `ARM_ENVIRONMENT`.
 * `pat_token_duration_seconds` - The current implementation of the azure auth via sp requires the provider to create a temporary personal access token within Databricks. The current AAD implementation does not cover all the APIs for Authentication. This field determines the duration in which that temporary PAT token is alive. It is measured in seconds and will default to `3600` seconds. 
 * `debug_truncate_bytes` - Applicable only when `TF_LOG=DEBUG` is set. Truncate JSON fields in HTTP requests and responses above this limit. Default is *96*.
 * `debug_headers` - Applicable only when `TF_LOG=DEBUG` is set. Debug HTTP headers of requests made by the provider. Default is *false*.
@@ -195,23 +196,24 @@ There are multiple environment variable options, the `DATABRICKS_AZURE_*` enviro
 
 The following configuration attributes can be passed via environment variables:
 
-| Argument | Environment variable |
-| --: | --- |
-| `host` | `DATABRICKS_HOST` |
-| `token` | `DATABRICKS_TOKEN` |
-| `username` | `DATABRICKS_USERNAME` |
-| `password` | `DATABRICKS_PASSWORD` |
-| `config_file` | `DATABRICKS_CONFIG_FILE` |
-| `profile` | `DATABRICKS_CONFIG_PROFILE` |
-| `azure_workspace_resource_id` | `DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID` |
-| `azure_workspace_name` | `DATABRICKS_AZURE_WORKSPACE_NAME` |
-| `azure_resource_group` | `DATABRICKS_AZURE_RESOURCE_GROUP` |
-| `azure_subscription_id` | `DATABRICKS_AZURE_SUBSCRIPTION_ID` or `ARM_SUBSCRIPTION_ID` |
-| `azure_client_secret` | `DATABRICKS_AZURE_CLIENT_SECRET` or `ARM_CLIENT_SECRET` |
-| `azure_client_id` | `DATABRICKS_AZURE_CLIENT_ID` or `ARM_CLIENT_ID` |
-| `azure_tenant_id` | `DATABRICKS_AZURE_TENANT_ID` or `ARM_TENANT_ID` |
-| `debug_truncate_bytes` | `DATABRICKS_DEBUG_TRUNCATE_BYTES` |
-| `debug_headers` | `DATABRICKS_DEBUG_HEADERS` |
+|                      Argument | Environment variable                                        |
+| ----------------------------: | ----------------------------------------------------------- |
+|                        `host` | `DATABRICKS_HOST`                                           |
+|                       `token` | `DATABRICKS_TOKEN`                                          |
+|                    `username` | `DATABRICKS_USERNAME`                                       |
+|                    `password` | `DATABRICKS_PASSWORD`                                       |
+|                 `config_file` | `DATABRICKS_CONFIG_FILE`                                    |
+|                     `profile` | `DATABRICKS_CONFIG_PROFILE`                                 |
+| `azure_workspace_resource_id` | `DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID`                    |
+|        `azure_workspace_name` | `DATABRICKS_AZURE_WORKSPACE_NAME`                           |
+|        `azure_resource_group` | `DATABRICKS_AZURE_RESOURCE_GROUP`                           |
+|       `azure_subscription_id` | `DATABRICKS_AZURE_SUBSCRIPTION_ID` or `ARM_SUBSCRIPTION_ID` |
+|         `azure_client_secret` | `DATABRICKS_AZURE_CLIENT_SECRET` or `ARM_CLIENT_SECRET`     |
+|             `azure_client_id` | `DATABRICKS_AZURE_CLIENT_ID` or `ARM_CLIENT_ID`             |
+|             `azure_tenant_id` | `DATABRICKS_AZURE_TENANT_ID` or `ARM_TENANT_ID`             |
+|           `azure_environment` | `ARM_ENVIRONMENT`                                           |
+|        `debug_truncate_bytes` | `DATABRICKS_DEBUG_TRUNCATE_BYTES`                           |
+|               `debug_headers` | `DATABRICKS_DEBUG_HEADERS`                                  |
 
 ## Empty provider block
 

provider/provider.go

@@ -189,6 +189,11 @@ func DatabricksProvider() *schema.Provider {
 				Default:     false,
 				Description: "Create ephemeral PAT tokens also for AZ CLI authenticated requests",
 			},
+			"azure_environment": {
+				Type:        schema.TypeString,
+				Optional:    true,
+				DefaultFunc: schema.EnvDefaultFunc("ARM_ENVIRONMENT", "public"),
+			},
 			"skip_verify": {
 				Type:        schema.TypeBool,
 				Description: "Skip SSL certificate verification for HTTP calls. Use at your own risk.",
@@ -287,6 +292,9 @@ func DatabricksProvider() *schema.Provider {
 			if v, ok := d.GetOk("azure_use_pat_for_cli"); ok {
 				pc.AzureAuth.UsePATForCLI = v.(bool)
 			}
+			if v, ok := d.GetOk("azure_environment"); ok {
+				pc.AzureAuth.Environment = v.(string)
+			}
 			authorizationMethodsUsed := []string{}
 			for name, used := range authsUsed {
 				if used {