azurerm_databricks_workspace - Upgrade Databricks API to the 2021-04-01-preview version

azurerm/internal/services/databricks/client/client.go

@@ -1,7 +1,7 @@
 package client
 
 import (
-	"github.com/Azure/azure-sdk-for-go/services/databricks/mgmt/2018-04-01/databricks"
+	"github.com/Azure/azure-sdk-for-go/services/preview/databricks/mgmt/2021-04-01-preview/databricks"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/common"
 )
 

azurerm/internal/services/databricks/databricks_customer_managed_key_resource.go

@@ -7,7 +7,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/Azure/azure-sdk-for-go/services/databricks/mgmt/2018-04-01/databricks"
+	"github.com/Azure/azure-sdk-for-go/services/preview/databricks/mgmt/2021-04-01-preview/databricks"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/helpers/tf"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/clients"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/locks"
@@ -148,22 +148,14 @@ func DatabricksWorkspaceCustomerManagedKeyCreateUpdate(d *pluginsdk.ResourceData
 	params := workspace.Parameters
 	params.Encryption = &databricks.WorkspaceEncryptionParameter{
 		Value: &databricks.Encryption{
-			KeySource:   databricks.MicrosoftKeyvault,
+			KeySource:   databricks.KeySourceMicrosoftKeyvault,
 			KeyName:     &key.Name,
 			KeyVersion:  &key.Version,
 			KeyVaultURI: &key.KeyVaultBaseUrl,
 		},
 	}
 
-	props := databricks.Workspace{
-		Location: workspace.Location,
-		Sku:      workspace.Sku,
-		WorkspaceProperties: &databricks.WorkspaceProperties{
-			ManagedResourceGroupID: workspace.WorkspaceProperties.ManagedResourceGroupID,
-			Parameters:             params,
-		},
-		Tags: workspace.Tags,
-	}
+	props := getProps(workspace, params)
 
 	future, err := workspaceClient.CreateOrUpdate(ctx, props, resourceID.ResourceGroup, resourceID.CustomerMangagedKeyName)
 	if err != nil {
@@ -229,7 +221,7 @@ func DatabricksWorkspaceCustomerManagedKeyRead(d *pluginsdk.ResourceData, meta i
 	// 	return fmt.Errorf("retrieving Databricks Workspace %q (Resource Group %q): `Workspace.WorkspaceProperties.Encryption.Value.KeySource` was expected to be %q, got %q", id.CustomerMangagedKeyName, id.ResourceGroup, string(databricks.MicrosoftKeyvault), keySource)
 	// }
 
-	if strings.EqualFold(keySource, string(databricks.MicrosoftKeyvault)) && (keyName == "" || keyVersion == "" || keyVaultURI == "") {
+	if strings.EqualFold(keySource, string(databricks.KeySourceMicrosoftKeyvault)) && (keyName == "" || keyVersion == "" || keyVaultURI == "") {
 		return fmt.Errorf("Databricks Workspace %q (Resource Group %q): `Workspace.WorkspaceProperties.Encryption.Value(s)` were nil", id.CustomerMangagedKeyName, id.ResourceGroup)
 	}
 
@@ -282,28 +274,40 @@ func DatabricksWorkspaceCustomerManagedKeyDelete(d *pluginsdk.ResourceData, meta
 	params := workspace.Parameters
 	params.Encryption = &databricks.WorkspaceEncryptionParameter{
 		Value: &databricks.Encryption{
-			KeySource: databricks.Default,
+			KeySource: databricks.KeySourceDefault,
 		},
 	}
 
+	props := getProps(workspace, params)
+
+	future, err := client.CreateOrUpdate(ctx, props, workspaceID.ResourceGroup, workspaceID.Name)
+	if err != nil {
+		return fmt.Errorf("creating/updating %s: %+v", workspaceID, err)
+	}
+
+	if err = future.WaitForCompletionRef(ctx, client.Client); err != nil {
+		return fmt.Errorf("waiting for create/update of %s: %+v", workspaceID, err)
+	}
+
+	return nil
+}
+
+func getProps(workspace databricks.Workspace, params *databricks.WorkspaceCustomParameters) databricks.Workspace {
 	props := databricks.Workspace{
 		Location: workspace.Location,
 		Sku:      workspace.Sku,
 		WorkspaceProperties: &databricks.WorkspaceProperties{
+			PublicNetworkAccess:    workspace.PublicNetworkAccess,
 			ManagedResourceGroupID: workspace.WorkspaceProperties.ManagedResourceGroupID,
 			Parameters:             params,
 		},
 		Tags: workspace.Tags,
 	}
 
-	future, err := client.CreateOrUpdate(ctx, props, workspaceID.ResourceGroup, workspaceID.Name)
-	if err != nil {
-		return fmt.Errorf("creating/updating %s: %+v", workspaceID, err)
+	// This is only valid if Private Link only is set
+	if workspace.PublicNetworkAccess == databricks.PublicNetworkAccessDisabled {
+		props.WorkspaceProperties.RequiredNsgRules = workspace.WorkspaceProperties.RequiredNsgRules
 	}
 
-	if err = future.WaitForCompletionRef(ctx, client.Client); err != nil {
-		return fmt.Errorf("waiting for create/update of %s: %+v", workspaceID, err)
-	}
-
-	return nil
+	return props
 }

azurerm/internal/services/databricks/databricks_customer_managed_key_resource_test.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/Azure/azure-sdk-for-go/services/databricks/mgmt/2018-04-01/databricks"
+	"github.com/Azure/azure-sdk-for-go/services/preview/databricks/mgmt/2021-04-01-preview/databricks"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/acceptance"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/acceptance/check"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/clients"
@@ -117,7 +117,7 @@ func (DatabricksWorkspaceCustomerManagedKeyResource) Exists(ctx context.Context,
 
 	// This is the only way we can tell if the CMK has actually been provisioned or not...
 	if resp.WorkspaceProperties.Parameters != nil && resp.WorkspaceProperties.Parameters.Encryption != nil {
-		if resp.WorkspaceProperties.Parameters.Encryption.Value.KeySource == databricks.MicrosoftKeyvault {
+		if resp.WorkspaceProperties.Parameters.Encryption.Value.KeySource == databricks.KeySourceMicrosoftKeyvault {
 			return utils.Bool(true), nil
 		}
 	}