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
 		}
 	}

azurerm/internal/services/databricks/databricks_workspace_private_endpoint_connection_data_source.go

@@ -0,0 +1,143 @@
+package databricks
+
+import (
+	"fmt"
+	"time"
+
+	"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/azure"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/clients"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/services/databricks/parse"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/tf/pluginsdk"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/timeouts"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
+)
+
+func dataSourceDatabricksWorkspacePrivateEndpointConnection() *pluginsdk.Resource {
+	return &pluginsdk.Resource{
+		Read: dataSourceDatabricksWorkspacePrivateEndpointConnectionRead,
+
+		Timeouts: &pluginsdk.ResourceTimeout{
+			Read: pluginsdk.DefaultTimeout(5 * time.Minute),
+		},
+
+		Schema: map[string]*pluginsdk.Schema{
+			"workspace_id": {
+				Type:         pluginsdk.TypeString,
+				Required:     true,
+				ValidateFunc: azure.ValidateResourceID,
+			},
+
+			"private_endpoint_id": {
+				Type:         pluginsdk.TypeString,
+				Required:     true,
+				ValidateFunc: azure.ValidateResourceID,
+			},
+
+			"connections": {
+				Type:     pluginsdk.TypeList,
+				Computed: true,
+				Elem: &pluginsdk.Resource{
+					Schema: map[string]*pluginsdk.Schema{
+						"workspace_private_endpoint_id": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+
+						"name": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+
+						"status": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+
+						"description": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+
+						"action_required": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
+func dataSourceDatabricksWorkspacePrivateEndpointConnectionRead(d *pluginsdk.ResourceData, meta interface{}) error {
+	client := meta.(*clients.Client).DataBricks.WorkspacesClient
+	ctx, cancel := timeouts.ForRead(meta.(*clients.Client).StopContext, d)
+	defer cancel()
+	workspace := d.Get("workspace_id").(string)
+	endpointId := d.Get("private_endpoint_id").(string)
+
+	id, err := parse.WorkspaceID(workspace)
+	if err != nil {
+		return err
+	}
+
+	resp, err := client.Get(ctx, id.ResourceGroup, id.Name)
+	if err != nil {
+		if utils.ResponseWasNotFound(resp.Response) {
+			return fmt.Errorf("%s was not found", id)
+		}
+
+		return fmt.Errorf("making Read request on %s: %+v", id, err)
+	}
+
+	d.SetId(id.ID())
+
+	d.Set("workspace_id", workspace)
+	d.Set("private_endpoint_id", endpointId)
+
+	if props := resp.WorkspaceProperties; props != nil {
+		if err := d.Set("connections", flattenPrivateEndpointConnections(props.PrivateEndpointConnections)); err != nil {
+			return fmt.Errorf("setting `connections`: %+v", err)
+		}
+	}
+
+	return nil
+}
+
+func flattenPrivateEndpointConnections(input *[]databricks.PrivateEndpointConnection) []interface{} {
+	results := make([]interface{}, 0)
+	if input == nil {
+		return results
+	}
+
+	for _, v := range *input {
+		result := make(map[string]interface{})
+
+		if name := v.Name; name != nil {
+			result["name"] = *name
+		}
+
+		if id := v.ID; id != nil {
+			result["workspace_private_endpoint_id"] = *id
+		}
+
+		if props := v.Properties; props != nil {
+			if connState := props.PrivateLinkServiceConnectionState; connState != nil {
+				if description := connState.Description; description != nil {
+					result["description"] = *description
+				}
+				if status := connState.Status; status != "" {
+					result["status"] = status
+				}
+				if actionReq := connState.ActionRequired; actionReq != nil {
+					result["action_required"] = *actionReq
+				}
+			}
+		}
+
+		results = append(results, result)
+	}
+
+	return results
+}