[Compute] Support associating disk and snapshot with a disk-access resource (Azure#14624)

* Add support for associating disk-access with disk and snapshot

* add scenario test

* src/azure-cli/azure/cli/command_modules/vm/tests/latest/test_vm_commands.py

* fix arg_type for network_access_policy

Author: houk-ms

Diff for: src/azure-cli-core/azure/cli/core/profiles/_shared.py

@@ -137,7 +137,7 @@ def default_api_version(self):
             'disks': '2020-05-01',
             'disk_encryption_sets': '2020-05-01',
             'disk_accesses': '2020-05-01',
-            'snapshots': '2019-07-01',
+            'snapshots': '2020-05-01',
             'galleries': '2019-12-01',
             'gallery_images': '2019-12-01',
             'gallery_image_versions': '2019-12-01',

Diff for: src/azure-cli/azure/cli/command_modules/vm/_help.py

@@ -57,6 +57,9 @@
   - name: Create a disk and specify maximum number of VMs that can attach to the disk at the same time.
     text: >
         az disk create -g MyResourceGroup -n MyDisk --size-gb 256 --max-shares 2 -l centraluseuap
+  - name: Create a disk and associate it with a disk access resource.
+    text: >
+        az disk create -g MyResourceGroup -n MyDisk --size-gb 10 --network-access-policy AllowPrivate --disk-access MyDiskAccessID
 """
 
 helps['disk delete'] = """
@@ -106,6 +109,9 @@
     text: |
         az disk update --name MyManagedDisk --resource-group MyResourceGroup --size-gb 20
     crafted: true
+  - name: Update a managed disk and associate it with a disk access resource.
+    text: |
+        az disk update --name MyManagedDisk --resource-group MyResourceGroup --network-access-policy AllowPrivate --disk-access MyDiskAccessID
 """
 
 helps['disk wait'] = """
@@ -681,6 +687,8 @@
     text: az snapshot create -g MyResourceGroup -n MySnapshot2 --source MyDisk
   - name: Create a snapshot from an existing disk in another resource group.
     text: az snapshot create -g MyResourceGroup -n MySnapshot2 --source "/subscriptions/00000/resourceGroups/AnotherResourceGroup/providers/Microsoft.Compute/disks/MyDisk"
+  - name: Create a snapshot and associate it with a disk access resource.
+    text: az snapshot create -g MyResourceGroup -n MySnapshot --size-gb 10 --network-access-policy AllowPrivate --disk-access MyDiskAccessID
 """
 
 helps['snapshot grant-access'] = """
@@ -716,6 +724,9 @@
     text: |
         az snapshot update --name MySnapshot --resource-group MyResourceGroup --subscription MySubscription
     crafted: true
+  - name: Update a snapshot and associate it with a disk access resource.
+    text: |
+        az snapshot update --name MySnapshot --resource-group MyResourceGroup --network-access-policy AllowPrivate --disk-access MyDiskAccessID
 """
 
 helps['snapshot wait'] = """

Diff for: src/azure-cli/azure/cli/command_modules/vm/_params.py

@@ -118,6 +118,9 @@ def load_arguments(self, _):
                        help='Encryption type. EncryptionAtRestWithPlatformKey: Disk is encrypted with XStore managed key at rest. It is the default encryption type. EncryptionAtRestWithCustomerKey: Disk is encrypted with Customer managed key at rest.')
             c.argument('disk_encryption_set', min_api='2019-07-01', help='Name or ID of disk encryption set that is used to encrypt the disk.')
             c.argument('location', help='Location. Values from: `az account list-locations`. You can configure the default location using `az configure --defaults location=<location>`. If location is not specified and no default location specified, location will be automatically set as same as the resource group.')
+            operation_group = 'disks' if scope == 'disk' else 'snapshots'
+            c.argument('network_access_policy', min_api='2020-05-01', help='Policy for accessing the disk via network.', arg_type=get_enum_type(self.get_models('NetworkAccessPolicy', operation_group=operation_group)))
+            c.argument('disk_access', min_api='2020-05-01', help='Name or ID of the disk access resource for using private endpoints on disks.')
 
     for scope in ['disk create', 'snapshot create']:
         with self.argument_context(scope) as c:

Diff for: src/azure-cli/azure/cli/command_modules/vm/custom.py

@@ -272,7 +272,8 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
                         encryption_type=None, disk_encryption_set=None, max_shares=None,
                         disk_iops_read_only=None, disk_mbps_read_only=None,
                         image_reference=None, image_reference_lun=None,
-                        gallery_image_reference=None, gallery_image_reference_lun=None):
+                        gallery_image_reference=None, gallery_image_reference_lun=None,
+                        network_access_policy=None, disk_access=None):
     from msrestazure.tools import resource_id, is_valid_resource_id
     from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -340,6 +341,11 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
             subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
             namespace='Microsoft.Compute', type='diskEncryptionSets', name=disk_encryption_set)
 
+    if disk_access is not None and not is_valid_resource_id(disk_access):
+        disk_access = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskAccesses', name=disk_access)
+
     encryption = None
     if disk_encryption_set:
         encryption = Encryption(type=encryption_type, disk_encryption_set_id=disk_encryption_set)
@@ -362,6 +368,10 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
         disk.disk_iops_read_only = disk_iops_read_only
     if disk_mbps_read_only is not None:
         disk.disk_mbps_read_only = disk_mbps_read_only
+    if network_access_policy is not None:
+        disk.network_access_policy = network_access_policy
+    if disk_access is not None:
+        disk.disk_access_id = disk_access
 
     client = _compute_client_factory(cmd.cli_ctx)
     return sdk_no_wait(no_wait, client.disks.create_or_update, resource_group_name, disk_name, disk)
@@ -380,7 +390,8 @@ def list_managed_disks(cmd, resource_group_name=None):
 
 
 def update_managed_disk(cmd, resource_group_name, instance, size_gb=None, sku=None, disk_iops_read_write=None,
-                        disk_mbps_read_write=None, encryption_type=None, disk_encryption_set=None):
+                        disk_mbps_read_write=None, encryption_type=None, disk_encryption_set=None,
+                        network_access_policy=None, disk_access=None):
     from msrestazure.tools import resource_id, is_valid_resource_id
     from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -403,6 +414,13 @@ def update_managed_disk(cmd, resource_group_name, instance, size_gb=None, sku=No
         instance.encryption.disk_encryption_set_id = disk_encryption_set
     if encryption_type is not None:
         instance.encryption.type = encryption_type
+    if network_access_policy is not None:
+        instance.network_access_policy = network_access_policy
+    if disk_access is not None and not is_valid_resource_id(disk_access):
+        disk_access = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskAccesses', name=disk_access)
+        instance.disk_access_id = disk_access
     return instance
 # endregion
 
@@ -481,7 +499,7 @@ def create_snapshot(cmd, resource_group_name, snapshot_name, location=None, size
                     # below are generated internally from 'source'
                     source_blob_uri=None, source_disk=None, source_snapshot=None, source_storage_account_id=None,
                     hyper_v_generation=None, tags=None, no_wait=False, disk_encryption_set=None,
-                    encryption_type=None):
+                    encryption_type=None, network_access_policy=None, disk_access=None):
     from msrestazure.tools import resource_id, is_valid_resource_id
     from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -511,6 +529,11 @@ def create_snapshot(cmd, resource_group_name, snapshot_name, location=None, size
             subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
             namespace='Microsoft.Compute', type='diskEncryptionSets', name=disk_encryption_set)
 
+    if disk_access is not None and not is_valid_resource_id(disk_access):
+        disk_access = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskAccesses', name=disk_access)
+
     if disk_encryption_set is not None and encryption_type is None:
         raise CLIError('usage error: Please specify --encryption-type.')
     if encryption_type is not None:
@@ -523,6 +546,10 @@ def create_snapshot(cmd, resource_group_name, snapshot_name, location=None, size
                         encryption=encryption)
     if hyper_v_generation:
         snapshot.hyper_vgeneration = hyper_v_generation
+    if network_access_policy is not None:
+        snapshot.network_access_policy = network_access_policy
+    if disk_access is not None:
+        snapshot.disk_access_id = disk_access
 
     client = _compute_client_factory(cmd.cli_ctx)
     return sdk_no_wait(no_wait, client.snapshots.create_or_update, resource_group_name, snapshot_name, snapshot)
@@ -540,7 +567,8 @@ def list_snapshots(cmd, resource_group_name=None):
     return client.snapshots.list()
 
 
-def update_snapshot(cmd, resource_group_name, instance, sku=None, disk_encryption_set=None, encryption_type=None):
+def update_snapshot(cmd, resource_group_name, instance, sku=None, disk_encryption_set=None,
+                    encryption_type=None, network_access_policy=None, disk_access=None):
     from msrestazure.tools import resource_id, is_valid_resource_id
     from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -557,6 +585,13 @@ def update_snapshot(cmd, resource_group_name, instance, sku=None, disk_encryptio
         instance.encryption.disk_encryption_set_id = disk_encryption_set
     if encryption_type is not None:
         instance.encryption.type = encryption_type
+    if network_access_policy is not None:
+        instance.network_access_policy = network_access_policy
+    if disk_access is not None and not is_valid_resource_id(disk_access):
+        disk_access = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx), resource_group=resource_group_name,
+            namespace='Microsoft.Compute', type='diskAccesses', name=disk_access)
+        instance.disk_access_id = disk_access
     return instance
 # endregion