{Compute} Migrate disk and snapshot commands to aaz-based (#32596)

Author: william051200

src/azure-cli/azure/cli/command_modules/vm/_actions.py

@@ -323,3 +323,36 @@ def _get_latest_image_version(cli_ctx, location, publisher, offer, sku, edge_zon
         if not top_one:
             raise InvalidArgumentValueError("Can't resolve the version of '{}:{}:{}'".format(publisher, offer, sku))
     return top_one[0].name
+
+
+def _get_latest_image_version_by_aaz(cli_ctx, location, publisher, offer, sku, edge_zone=None):
+    from azure.cli.core.azclierror import InvalidArgumentValueError
+    if edge_zone is not None:
+        from .aaz.latest.vm.image.edge_zone import List as VmImageEdgeZoneList
+        command_args = {
+            'edge_zone': edge_zone,
+            'location': location,
+            'offer': offer,
+            'publisher': publisher,
+            'sku': sku,
+            'orderby': 'name desc',
+            'top': 1,
+        }
+        top_one = VmImageEdgeZoneList(cli_ctx=cli_ctx)(command_args=command_args)
+        if not top_one:
+            raise InvalidArgumentValueError("Can't resolve the version of '{}:{}:{}:{}'"
+                                            .format(publisher, offer, sku, edge_zone))
+    else:
+        from .aaz.latest.vm.image import List as VmImageList
+        command_args = {
+            'location': location,
+            'offer': offer,
+            'publisher': publisher,
+            'sku': sku,
+            'orderby': 'name desc',
+            'top': 1,
+        }
+        top_one = VmImageList(cli_ctx=cli_ctx)(command_args=command_args)
+        if not top_one:
+            raise InvalidArgumentValueError("Can't resolve the version of '{}:{}:{}'".format(publisher, offer, sku))
+    return top_one[0].get('name')

src/azure-cli/azure/cli/command_modules/vm/_validators.py

@@ -1952,11 +1952,8 @@ def validate_vmss_disk(cmd, namespace):
         raise CLIError('usage error: --disk EXIST_DISK --instance-id ID')
 
 
-def _validate_gallery_image_reference(cmd, namespace):
-    from azure.cli.core.profiles import ResourceType
-    is_validate = 'gallery_image_reference' in namespace and namespace.gallery_image_reference is not None \
-                  and cmd.supported_api_version(resource_type=ResourceType.MGMT_COMPUTE,
-                                                operation_group='disks', min_api='2022-03-02')
+def _validate_gallery_image_reference(namespace):
+    is_validate = 'gallery_image_reference' in namespace and namespace.gallery_image_reference is not None
     if not is_validate:
         return
 
@@ -1985,9 +1982,9 @@ def process_disk_create_namespace(cmd, namespace):
     from azure.core.exceptions import HttpResponseError
     validate_tags(namespace)
     validate_edge_zone(cmd, namespace)
-    _validate_gallery_image_reference(cmd, namespace)
+    _validate_gallery_image_reference(namespace)
     _validate_security_data_uri(namespace)
-    _validate_upload_type(cmd, namespace)
+    _validate_upload_type(namespace)
     _validate_secure_vm_disk_encryption_set(namespace)
     _validate_hyper_v_generation(namespace)
     if namespace.source:
@@ -2020,20 +2017,14 @@ def _validate_security_data_uri(namespace):
             'Please specify --source when using the --security-data-uri parameter')
 
 
-def _validate_upload_type(cmd, namespace):
+def _validate_upload_type(namespace):
     if 'upload_type' not in namespace:
         return
 
     if not namespace.upload_type and namespace.for_upload:
         namespace.upload_type = 'Upload'
 
     if namespace.upload_type == 'UploadWithSecurityData':
-
-        if not cmd.supported_api_version(min_api='2021-08-01', operation_group='disks'):
-            raise ArgumentUsageError(
-                "'UploadWithSecurityData' is not supported in the current profile. "
-                "Please upgrade your profile with 'az cloud set --profile newerProfile' and try again")
-
         if not namespace.security_type:
             raise RequiredArgumentMissingError(
                 "Please specify --security-type when the value of --upload-type is 'UploadWithSecurityData'")
@@ -2069,7 +2060,7 @@ def process_snapshot_create_namespace(cmd, namespace):
     from azure.core.exceptions import HttpResponseError
     validate_tags(namespace)
     validate_edge_zone(cmd, namespace)
-    _validate_gallery_image_reference(cmd, namespace)
+    _validate_gallery_image_reference(namespace)
     if namespace.source:
         usage_error = 'usage error: --source {SNAPSHOT | DISK} | --source VHD_BLOB_URI [--source-storage-account-id ID]'
         try:

src/azure-cli/azure/cli/command_modules/vm/commands.py

@@ -6,8 +6,7 @@
 from azure.cli.command_modules.vm._client_factory import (cf_vm,
                                                           cf_vm_ext, cf_vm_ext_image,
                                                           cf_vm_image, cf_vm_image_term, cf_usage,
-                                                          cf_vmss, cf_disks, cf_snapshots,
-                                                          cf_images,
+                                                          cf_vmss, cf_images,
                                                           cf_galleries, cf_gallery_images, cf_gallery_image_versions,
                                                           cf_proximity_placement_groups,
                                                           cf_dedicated_hosts, cf_dedicated_host_groups,
@@ -63,22 +62,11 @@ def load_command_table(self, _):
         operation_group='availability_sets'
     )
 
-    compute_disk_sdk = CliCommandType(
-        operations_tmpl='azure.mgmt.compute.operations#DisksOperations.{}',
-        client_factory=cf_disks,
-        operation_group='disks'
-    )
-
     compute_image_sdk = CliCommandType(
         operations_tmpl='azure.mgmt.compute.operations#ImagesOperations.{}',
         client_factory=cf_images
     )
 
-    compute_snapshot_sdk = CliCommandType(
-        operations_tmpl='azure.mgmt.compute.operations#SnapshotsOperations.{}',
-        client_factory=cf_snapshots
-    )
-
     compute_vm_sdk = CliCommandType(
         operations_tmpl='azure.mgmt.compute.operations#VirtualMachinesOperations.{}',
         client_factory=cf_vm
@@ -208,7 +196,7 @@ def load_command_table(self, _):
         from .operations.ppg import PPGShow
         self.command_table["ppg show"] = PPGShow(loader=self)
 
-    with self.command_group('disk', compute_disk_sdk, operation_group='disks') as g:
+    with self.command_group('disk', operation_group='disks') as g:
         g.custom_command('create', 'create_managed_disk', supports_no_wait=True, table_transformer=transform_disk_create_table_output, validator=process_disk_create_namespace)
         from .operations.disk import DiskUpdate, DiskGrantAccess
         self.command_table["disk grant-access"] = DiskGrantAccess(loader=self)
@@ -283,7 +271,7 @@ def load_command_table(self, _):
         g.custom_command('remove', 'remove_template_error_handler', supports_local_cache=True)
         g.custom_show_command('show', 'show_template_error_handler', supports_local_cache=True)
 
-    with self.command_group('snapshot', compute_snapshot_sdk, operation_group='snapshots') as g:
+    with self.command_group('snapshot', operation_group='snapshots') as g:
         g.custom_command('create', 'create_snapshot', validator=process_snapshot_create_namespace, supports_no_wait=True)
         from .operations.snapshot import SnapshotUpdate
         self.command_table['snapshot update'] = SnapshotUpdate(loader=self)

src/azure-cli/azure/cli/command_modules/vm/custom.py

@@ -41,7 +41,7 @@
 from ._vm_diagnostics_templates import get_default_diag_config
 
 from ._actions import (load_images_from_aliases_doc, load_extension_images_thru_services,
-                       load_images_thru_services, _get_latest_image_version)
+                       load_images_thru_services, _get_latest_image_version, _get_latest_image_version_by_aaz)
 from ._client_factory import (_compute_client_factory, cf_vm_image_term)
 
 from .aaz.latest.vm.disk import AttachDetachDataDisk
@@ -484,8 +484,8 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
             if len(terms) == 4:  # URN
                 disk_publisher, disk_offer, disk_sku, disk_version = terms[0], terms[1], terms[2], terms[3]
                 if disk_version.lower() == 'latest':
-                    disk_version = _get_latest_image_version(cmd.cli_ctx, location, disk_publisher, disk_offer,
-                                                             disk_sku)
+                    disk_version = _get_latest_image_version_by_aaz(cmd.cli_ctx, location, disk_publisher, disk_offer,
+                                                                    disk_sku)
             else:  # error
                 raise CLIError('usage error: --image-reference should be ID or URN (publisher:offer:sku:version).')
         else:
@@ -494,14 +494,20 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
             disk_publisher, disk_offer, disk_sku, disk_version = \
                 terms['child_name_1'], terms['child_name_3'], terms['child_name_4'], terms['child_name_5']
 
-        client = _compute_client_factory(cmd.cli_ctx)
-        response = client.virtual_machine_images.get(location=location, publisher_name=disk_publisher,
-                                                     offer=disk_offer, skus=disk_sku, version=disk_version)
+        from .aaz.latest.vm.image import Show as VmImageShow
+        command_args = {
+            'location': location,
+            'offer': disk_offer,
+            'publisher': disk_publisher,
+            'sku': disk_sku,
+            'version': disk_version,
+        }
+        response = VmImageShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)
 
-        if hasattr(response, 'hyper_v_generation'):
-            if response.hyper_v_generation == 'V1':
+        if response.get('hyper_v_generation'):
+            if response.get('hyper_v_generation') == 'V1':
                 logger.warning(UPGRADE_SECURITY_HINT)
-            elif response.hyper_v_generation == 'V2':
+            elif response.get('hyper_v_generation') == 'V2':
                 # set default value of hyper_v_generation
                 if hyper_v_generation == 'V1':
                     hyper_v_generation = 'V2'
@@ -512,7 +518,7 @@ def create_managed_disk(cmd, resource_group_name, disk_name, location=None,  # p
                     logger.warning(UPGRADE_SECURITY_HINT)
 
         # image_reference is an ID now
-        image_reference = {'id': response.id}
+        image_reference = {'id': response.get('id')}
         if image_reference_lun is not None:
             image_reference['lun'] = image_reference_lun
 

src/azure-cli/azure/cli/command_modules/vm/tests/latest/recordings/test_confidential_disk_encryption_set.yaml

@@ -1125,7 +1125,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.76.0 azsdk-python-core/1.35.0 Python/3.11.9 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions?$top=1&$orderby=name%20desc&api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions?$top=1&$orderby=name%20desc&api-version=2024-07-01
   response:
     body:
       string: "[\r\n  {\r\n    \"location\": \"westus\",\r\n    \"name\": \"20348.469.220106\",\r\n
@@ -1178,7 +1178,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.76.0 azsdk-python-core/1.35.0 Python/3.11.9 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions/20348.469.220106?api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions/20348.469.220106?api-version=2024-07-01
   response:
     body:
       string: "{\r\n  \"properties\": {\r\n    \"hyperVGeneration\": \"V2\",\r\n    \"architecture\":
@@ -1534,7 +1534,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.76.0 azsdk-python-core/1.35.0 Python/3.11.9 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions?$top=1&$orderby=name%20desc&api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions?$top=1&$orderby=name%20desc&api-version=2024-07-01
   response:
     body:
       string: "[\r\n  {\r\n    \"location\": \"westus\",\r\n    \"name\": \"20348.469.220106\",\r\n
@@ -1587,7 +1587,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.76.0 azsdk-python-core/1.35.0 Python/3.11.9 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions/20348.469.220106?api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/westus/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/windows-cvm/skus/2022-datacenter-cvm/versions/20348.469.220106?api-version=2024-07-01
   response:
     body:
       string: "{\r\n  \"properties\": {\r\n    \"hyperVGeneration\": \"V2\",\r\n    \"architecture\":

src/azure-cli/azure/cli/command_modules/vm/tests/latest/recordings/test_create_image_version_with_region_cvm_encryption_pmk.yaml

@@ -502,7 +502,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.77.0 azsdk-python-core/1.35.0 Python/3.10.11 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/centraluseuap/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/WindowsServer/skus/2022-datacenter-smalldisk-g2/versions?$top=1&$orderby=name%20desc&api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/centraluseuap/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/WindowsServer/skus/2022-datacenter-smalldisk-g2/versions?$top=1&$orderby=name%20desc&api-version=2024-07-01
   response:
     body:
       string: "[\r\n  {\r\n    \"location\": \"CentralUSEUAP\",\r\n    \"name\": \"20348.4171.250903\",\r\n
@@ -554,7 +554,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.77.0 azsdk-python-core/1.35.0 Python/3.10.11 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/centraluseuap/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/WindowsServer/skus/2022-datacenter-smalldisk-g2/versions/20348.4171.250903?api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/centraluseuap/publishers/MicrosoftWindowsServer/artifacttypes/vmimage/offers/WindowsServer/skus/2022-datacenter-smalldisk-g2/versions/20348.4171.250903?api-version=2024-07-01
   response:
     body:
       string: "{\r\n  \"properties\": {\r\n    \"hyperVGeneration\": \"V2\",\r\n    \"architecture\":

src/azure-cli/azure/cli/command_modules/vm/tests/latest/recordings/test_disk_encryption_set_disk_update.yaml

@@ -831,7 +831,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.77.0 azsdk-python-core/1.35.0 Python/3.10.11 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/eastus/publishers/Debian/artifacttypes/vmimage/offers/debian-11/skus/11-backports-gen2/versions?$top=1&$orderby=name%20desc&api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/eastus/publishers/Debian/artifacttypes/vmimage/offers/debian-11/skus/11-backports-gen2/versions?$top=1&$orderby=name%20desc&api-version=2024-07-01
   response:
     body:
       string: "[\r\n  {\r\n    \"location\": \"eastus\",\r\n    \"name\": \"0.20250703.2162\",\r\n
@@ -883,7 +883,7 @@ interactions:
       User-Agent:
       - AZURECLI/2.77.0 azsdk-python-core/1.35.0 Python/3.10.11 (Windows-10-10.0.26100-SP0)
     method: GET
-    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/eastus/publishers/Debian/artifacttypes/vmimage/offers/debian-11/skus/11-backports-gen2/versions/0.20250703.2162?api-version=2024-11-01
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Compute/locations/eastus/publishers/Debian/artifacttypes/vmimage/offers/debian-11/skus/11-backports-gen2/versions/0.20250703.2162?api-version=2024-07-01
   response:
     body:
       string: "{\r\n  \"properties\": {\r\n    \"hyperVGeneration\": \"V2\",\r\n    \"architecture\":