azurerm_databricks_workspace: Add support for CMK for disk encryption, multiple other updates

internal/services/databricks/databricks_workspace_data_source.go

@@ -46,6 +46,31 @@ func dataSourceDatabricksWorkspace() *pluginsdk.Resource {
 				Computed: true,
 			},
 
+			"managed_disk_identity": {
+				Type:     pluginsdk.TypeList,
+				Computed: true,
+				Elem: &pluginsdk.Resource{
+					Schema: map[string]*pluginsdk.Schema{
+						"principal_id": {
+							Type:      pluginsdk.TypeString,
+							Sensitive: true,
+							Computed:  true,
+						},
+
+						"tenant_id": {
+							Type:      pluginsdk.TypeString,
+							Sensitive: true,
+							Computed:  true,
+						},
+
+						"type": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+					},
+				},
+			},
+
 			"storage_account_identity": {
 				Type:     pluginsdk.TypeList,
 				Computed: true,
@@ -102,9 +127,12 @@ func dataSourceDatabricksWorkspaceRead(d *pluginsdk.ResourceData, meta interface
 			d.Set("sku", sku.Name)
 		}
 		d.Set("workspace_id", model.Properties.WorkspaceId)
-		if err := d.Set("storage_account_identity", flattenWorkspaceStorageAccountIdentity(model.Properties.StorageAccountIdentity)); err != nil {
+		if err := d.Set("storage_account_identity", flattenWorkspaceManagedIdentity(model.Properties.StorageAccountIdentity)); err != nil {
 			return fmt.Errorf("setting `storage_account_identity`: %+v", err)
 		}
+		if err := d.Set("managed_disk_identity", flattenWorkspaceManagedIdentity(model.Properties.StorageAccountIdentity)); err != nil {
+			return fmt.Errorf("setting `managed_disk_identity`: %+v", err)
+		}
 		d.Set("workspace_url", model.Properties.WorkspaceUrl)
 		d.Set("location", model.Location)
 

internal/services/databricks/databricks_workspace_resource.go

@@ -79,18 +79,53 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 
 			"customer_managed_key_enabled": {
 				Type:     pluginsdk.TypeBool,
-				ForceNew: true,
 				Optional: true,
 				Default:  false,
 			},
 
 			"managed_services_cmk_key_vault_key_id": {
 				Type:         pluginsdk.TypeString,
 				Optional:     true,
-				ForceNew:     true,
 				ValidateFunc: keyVaultValidate.KeyVaultChildID,
 			},
 
+			"managed_disk_cmk_key_vault_key_id": {
+				Type:         pluginsdk.TypeString,
+				Optional:     true,
+				ValidateFunc: keyVaultValidate.KeyVaultChildID,
+			},
+
+			"managed_disk_identity": {
+				Type:     pluginsdk.TypeList,
+				Computed: true,
+				Elem: &pluginsdk.Resource{
+					Schema: map[string]*pluginsdk.Schema{
+						"principal_id": {
+							Type:      pluginsdk.TypeString,
+							Sensitive: true,
+							Computed:  true,
+						},
+
+						"tenant_id": {
+							Type:      pluginsdk.TypeString,
+							Sensitive: true,
+							Computed:  true,
+						},
+
+						"type": {
+							Type:     pluginsdk.TypeString,
+							Computed: true,
+						},
+					},
+				},
+			},
+
+			"managed_disk_cmk_rotation_to_latest_version_enabled": {
+				Type:         pluginsdk.TypeBool,
+				Optional:     true,
+				RequiredWith: []string{"managed_disk_cmk_key_vault_key_id"},
+			},
+
 			"infrastructure_encryption_enabled": {
 				Type:     pluginsdk.TypeBool,
 				ForceNew: true,
@@ -100,15 +135,13 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 
 			"public_network_access_enabled": {
 				Type:     pluginsdk.TypeBool,
-				ForceNew: true,
 				Optional: true,
 				Default:  true,
 			},
 
 			"network_security_group_rules_required": {
 				Type:     pluginsdk.TypeString,
 				Optional: true,
-				ForceNew: true,
 				Computed: true,
 				ValidateFunc: validation.StringInSlice([]string{
 					string(workspaces.RequiredNsgRulesAllRules),
@@ -149,7 +182,6 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 
 						"no_public_ip": {
 							Type:         pluginsdk.TypeBool,
-							ForceNew:     true,
 							Optional:     true,
 							Computed:     true,
 							AtLeastOneOf: workspaceCustomParametersString(),
@@ -242,6 +274,11 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 				Computed: true,
 			},
 
+			"disk_encryption_set_id": {
+				Type:     pluginsdk.TypeString,
+				Computed: true,
+			},
+
 			"storage_account_identity": {
 				Type:     pluginsdk.TypeList,
 				Computed: true,
@@ -277,6 +314,7 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 			_, requireNsgRules := d.GetChange("network_security_group_rules_required")
 			_, backendPool := d.GetChange("load_balancer_backend_address_pool_id")
 			_, managedServicesCMK := d.GetChange("managed_services_cmk_key_vault_key_id")
+			_, managedDiskCMK := d.GetChange("managed_disk_cmk_key_vault_key_id")
 
 			oldSku, newSku := d.GetChange("sku")
 
@@ -302,8 +340,8 @@ func resourceDatabricksWorkspace() *pluginsdk.Resource {
 				}
 			}
 
-			if (customerEncryptionEnabled.(bool) || infrastructureEncryptionEnabled.(bool) || managedServicesCMK.(string) != "") && !strings.EqualFold("premium", newSku.(string)) {
-				return fmt.Errorf("'customer_managed_key_enabled', 'infrastructure_encryption_enabled' and 'managed_services_cmk_key_vault_key_id' are only available with a 'premium' workspace 'sku', got %q", newSku)
+			if (customerEncryptionEnabled.(bool) || infrastructureEncryptionEnabled.(bool) || managedServicesCMK.(string) != "" || managedDiskCMK.(string) != "") && !strings.EqualFold("premium", newSku.(string)) {
+				return fmt.Errorf("'customer_managed_key_enabled', 'infrastructure_encryption_enabled', 'managed_disk_cmk_key_vault_key_id' and 'managed_services_cmk_key_vault_key_id' are only available with a 'premium' workspace 'sku', got %q", newSku)
 			}
 
 			return nil
@@ -408,30 +446,59 @@ func resourceDatabricksWorkspaceCreateUpdate(d *pluginsdk.ResourceData, meta int
 	// Set up customer-managed keys for managed services encryption (e.g. notebook)
 	setEncrypt := false
 	encrypt := &workspaces.WorkspacePropertiesEncryption{}
-	keyIdRaw := d.Get("managed_services_cmk_key_vault_key_id").(string)
-	if keyIdRaw != "" {
+	encrypt.Entities = workspaces.EncryptionEntitiesDefinition{}
+	servicesKeyIdRaw := d.Get("managed_services_cmk_key_vault_key_id").(string)
+	if servicesKeyIdRaw != "" {
 		setEncrypt = true
-		key, err := keyVaultParse.ParseNestedItemID(keyIdRaw)
+		key, err := keyVaultParse.ParseNestedItemID(servicesKeyIdRaw)
 		if err != nil {
 			return err
 		}
 
-		encrypt.Entities = workspaces.EncryptionEntitiesDefinition{
-			ManagedServices: &workspaces.EncryptionV2{
-				// There is only one valid source for this field at this point in time so I have hardcoded the value
-				KeySource: workspaces.EncryptionKeySourceMicrosoftPointKeyvault,
-				KeyVaultProperties: &workspaces.EncryptionV2KeyVaultProperties{
-					KeyName:     key.Name,
-					KeyVersion:  key.Version,
-					KeyVaultUri: key.KeyVaultBaseUrl,
-				},
+		// make sure the key vault exists
+		keyVaultIdRaw, err := keyVaultsClient.KeyVaultIDFromBaseUrl(ctx, meta.(*clients.Client).Resource, key.KeyVaultBaseUrl)
+		if err != nil || keyVaultIdRaw == nil {
+			return fmt.Errorf("retrieving the Resource ID for the customer-managed keys for managed services Key Vault at URL %q: %+v", key.KeyVaultBaseUrl, err)
+		}
+
+		encrypt.Entities.ManagedServices = &workspaces.EncryptionV2{
+			// There is only one valid source for this field at this point in time so I have hardcoded the value
+			KeySource: workspaces.EncryptionKeySourceMicrosoftPointKeyvault,
+			KeyVaultProperties: &workspaces.EncryptionV2KeyVaultProperties{
+				KeyName:     key.Name,
+				KeyVersion:  key.Version,
+				KeyVaultUri: key.KeyVaultBaseUrl,
 			},
 		}
+	}
+
+	diskKeyIdRaw := d.Get("managed_disk_cmk_key_vault_key_id").(string)
+	if diskKeyIdRaw != "" {
+		setEncrypt = true
+		key, err := keyVaultParse.ParseNestedItemID(diskKeyIdRaw)
+		if err != nil {
+			return err
+		}
 
 		// make sure the key vault exists
 		keyVaultIdRaw, err := keyVaultsClient.KeyVaultIDFromBaseUrl(ctx, meta.(*clients.Client).Resource, key.KeyVaultBaseUrl)
 		if err != nil || keyVaultIdRaw == nil {
-			return fmt.Errorf("retrieving the Resource ID for the customer-managed keys for managed services Key Vault at URL %q: %+v", key.KeyVaultBaseUrl, err)
+			return fmt.Errorf("retrieving the Resource ID for the customer-managed keys for managed disk Key Vault at URL %q: %+v", key.KeyVaultBaseUrl, err)
+		}
+
+		encrypt.Entities.ManagedDisk = &workspaces.ManagedDiskEncryption{
+			// There is only one valid source for this field at this point in time so I have hardcoded the value
+			KeySource: workspaces.EncryptionKeySourceMicrosoftPointKeyvault,
+			KeyVaultProperties: workspaces.ManagedDiskEncryptionKeyVaultProperties{
+				KeyName:     key.Name,
+				KeyVersion:  key.Version,
+				KeyVaultUri: key.KeyVaultBaseUrl,
+			},
+		}
+
+		rotationEnabled := d.Get("managed_disk_cmk_rotation_to_latest_version_enabled").(bool)
+		if rotationEnabled {
+			encrypt.Entities.ManagedDisk.RotationToLatestKeyVersionEnabled = utils.Bool(rotationEnabled)
 		}
 	}
 
@@ -490,8 +557,12 @@ func resourceDatabricksWorkspaceCreateUpdate(d *pluginsdk.ResourceData, meta int
 		return fmt.Errorf("setting `custom_parameters`: %+v", err)
 	}
 
-	if encrypt != nil && keyIdRaw != "" {
-		d.Set("managed_services_cmk_key_vault_key_id", keyIdRaw)
+	if encrypt != nil && servicesKeyIdRaw != "" {
+		d.Set("managed_services_cmk_key_vault_key_id", servicesKeyIdRaw)
+	}
+
+	if encrypt != nil && diskKeyIdRaw != "" {
+		d.Set("managed_disk_cmk_key_vault_key_id", diskKeyIdRaw)
 	}
 
 	return resourceDatabricksWorkspaceRead(d, meta)
@@ -570,10 +641,14 @@ func resourceDatabricksWorkspaceRead(d *pluginsdk.ResourceData, meta interface{}
 			d.Set("load_balancer_backend_address_pool_id", backendPoolReadId)
 		}
 
-		if err := d.Set("storage_account_identity", flattenWorkspaceStorageAccountIdentity(model.Properties.StorageAccountIdentity)); err != nil {
+		if err := d.Set("storage_account_identity", flattenWorkspaceManagedIdentity(model.Properties.StorageAccountIdentity)); err != nil {
 			return fmt.Errorf("setting `storage_account_identity`: %+v", err)
 		}
 
+		if err := d.Set("managed_disk_identity", flattenWorkspaceManagedIdentity(model.Properties.ManagedDiskIdentity)); err != nil {
+			return fmt.Errorf("setting `managed_disk_identity`: %+v", err)
+		}
+
 		if model.Properties.WorkspaceUrl != nil {
 			d.Set("workspace_url", model.Properties.WorkspaceUrl)
 		}
@@ -601,6 +676,31 @@ func resourceDatabricksWorkspaceRead(d *pluginsdk.ResourceData, meta interface{}
 				d.Set("managed_services_cmk_key_vault_key_id", key.ID())
 			}
 		}
+		// customer managed key for managed disk
+		encryptDiskKeyName := ""
+		encryptDiskKeyVersion := ""
+		encryptDiskKeyVaultURI := ""
+		encryptDiskRotationEnabled := false
+
+		if encryption := model.Properties.Encryption; encryption != nil {
+			if encryptionProps := encryption.Entities.ManagedDisk; encryptionProps != nil {
+				encryptDiskKeyName = encryptionProps.KeyVaultProperties.KeyName
+				encryptDiskKeyVersion = encryptionProps.KeyVaultProperties.KeyVersion
+				encryptDiskKeyVaultURI = encryptionProps.KeyVaultProperties.KeyVaultUri
+				encryptDiskRotationEnabled = *encryptionProps.RotationToLatestKeyVersionEnabled
+			}
+
+		}
+
+		if encryptDiskKeyVaultURI != "" {
+			key, err := keyVaultParse.NewNestedItemID(encryptDiskKeyVaultURI, "keys", encryptDiskKeyName, encryptDiskKeyVersion)
+			if err == nil {
+				d.Set("managed_disk_cmk_key_vault_key_id", key.ID())
+			}
+			d.Set("managed_disk_cmk_rotation_to_latest_version_enabled", encryptDiskRotationEnabled)
+			d.Set("disk_encryption_set_id", model.Properties.DiskEncryptionSetId)
+		}
+
 		return tags.FlattenAndSet(d, model.Tags)
 	}
 
@@ -624,7 +724,7 @@ func resourceDatabricksWorkspaceDelete(d *pluginsdk.ResourceData, meta interface
 	return nil
 }
 
-func flattenWorkspaceStorageAccountIdentity(input *workspaces.ManagedIdentityConfiguration) []interface{} {
+func flattenWorkspaceManagedIdentity(input *workspaces.ManagedIdentityConfiguration) []interface{} {
 	if input == nil {
 		return nil
 	}