[Compute] az sig image-version: Add parameters `--target-edge-zone-…

…encryption` and `--target-edge-zones` to support edge zones (#24202)

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

@@ -994,6 +994,22 @@
         --target-regions westus=2=standard eastus \\
         --target-region-encryption WestUSDiskEncryptionSet1,0,WestUSDiskEncryptionSet2 \\
         EastUSDiskEncryptionSet1,0,EastUSDiskEncryptionSet2
+  - name: Add a new image version and copy it to extended locations.
+    text: |
+        az sig image-version create --resource-group MyResourceGroup \\
+        --gallery-name MyGallery --gallery-image-definition MyImage \\
+        --gallery-image-version 1.0.0 --replica-count 1 \\
+        --storage-account-type Standard_ZRS --managed-image image-name \\
+        --target-edge-zones westus=microsoftlosangeles1 eastus=microsoftlosangeles2=1 \\
+        brazilsouth=2=standard_lrs
+  - name: Add a new image version and copy it to extended locations with encryption using a disk encryption set.
+    text: |
+        az sig image-version create --resource-group MyResourceGroup \\
+        --gallery-name MyGallery --gallery-image-definition MyImage \\
+        --gallery-image-version 1.0.0 \\
+        --virtual-machine /subscriptions/00000000-0000-0000-0000-00000000xxxx/resourceGroups/imageGroups/providers/Microsoft.Compute/virtualMachines/MyVM \\
+        --target-edge-zones westus=microsoftlosangeles1 \\
+        --target-edge-zone-encryption microsoftlosangeles1,WestUSDiskEncryptionSet1,0,WestUSDiskEncryptionSet2
   - name: Add a new image version and don't wait on it. Later you can invoke "az sig image-version wait" command when ready to create a vm from the gallery image version
     text: |
         az sig image-version create --resource-group MyResourceGroup \\
@@ -1053,6 +1069,18 @@
         --gallery-name MyGallery --gallery-image-definition MyImage \\
         --gallery-image-version 1.0.0 \\
         --target-regions westcentralus=2 eastus2
+  - name: Change the replication extended locations
+    text: |
+        az sig image-version update --resource-group MyResourceGroup \\
+        --gallery-name MyGallery --gallery-image-definition MyImage \\
+        --gallery-image-version 1.0.0 \\
+        --target-edge-zones westus=microsoftlosangeles1 eastus=microsoftlosangeles2=1
+  - name: Clear the replication extended locations
+    text: |
+        az sig image-version update --resource-group MyResourceGroup \\
+        --gallery-name MyGallery --gallery-image-definition MyImage \\
+        --gallery-image-version 1.0.0 \\
+        --target-edge-zones None
   - name: Replicate to an additional region. Optional, you can set the replica count for the region and exclude this image when using the latest version of the image definition.
     text: |
         az sig image-version update --resource-group MyResourceGroup \\

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

@@ -1288,6 +1288,10 @@ def load_arguments(self, _):
         c.argument('target_region_cvm_encryption', nargs='+', min_api='2021-10-01', help='Space-separated list of customer managed key for Confidential VM encrypting the OS disk in the gallery artifact for each region. Format for each region: `<os_cvm_encryption_type>,<os_cvm_des>`. The valid values for os_cvm_encryption_type are EncryptedVMGuestStateOnlyWithPmk, EncryptedWithPmk, EncryptedWithCmk.')
         c.argument('virtual_machine', help='Resource id of VM source')
         c.argument('image_version', help='Resource id of gallery image version source')
+        c.argument('target_zone_encryption', nargs='+', min_api='2022-01-03',
+                   options_list=['--target-edge-zone-encryption', '--zone-encryption'],
+                   help='Space-separated list of customer managed keys for encrypting the OS and data disks in the gallery artifact for each region. '
+                        'Format for each edge zone: <edge zone>,<os_des>,<lun1>,<lun1_des>,<lun2>,<lun2_des>.')
 
     with self.argument_context('sig image-version list-shared') as c:
         c.argument('location', arg_type=get_location_type(self.cli_ctx), id_part='name')
@@ -1321,6 +1325,10 @@ def load_arguments(self, _):
                        help='Space-separated list of regions and their replica counts. Use `<region>[=<replica count>][=<storage account type>]` to optionally set the replica count and/or storage account type for each region. '
                             'If a replica count is not specified, the default replica count will be used. If a storage account type is not specified, the default storage account type will be used')
             c.argument('replica_count', help='The default number of replicas to be created per region. To set regional replication counts, use --target-regions', type=int)
+            c.argument('target_edge_zones', nargs='*', min_api='2022-01-03',
+                       help='Space-separated list of regions, edge zones, replica counts and storage types. Use <region>=<edge zone>[=<replica count>][=<storage account type>] to optionally set the replica count and/or storage account type for each region. '
+                            'If a replica count is not specified, the default replica count will be used. If a storage account type is not specified, the default storage account type will be used. '
+                            'If "--target-edge-zones None" is specified, the target extended locations will be cleared.')
             c.argument('allow_replicated_location_deletion', arg_type=get_three_state_flag(), min_api='2022-03-03', help='Indicate whether or not removing this gallery image version from replicated regions is allowed.')
 
     with self.argument_context('sig show-community') as c:

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

@@ -2102,11 +2102,10 @@ def process_set_applications_namespace(cmd, namespace):  # pylint: disable=unuse
 
 def process_gallery_image_version_namespace(cmd, namespace):
     from azure.cli.core.azclierror import InvalidArgumentValueError
-    TargetRegion, EncryptionImages, OSDiskImageEncryption, DataDiskImageEncryption, ConfidentialVMEncryptionType = \
-        cmd.get_models('TargetRegion', 'EncryptionImages', 'OSDiskImageEncryption', 'DataDiskImageEncryption',
-                       'ConfidentialVMEncryptionType')
-    storage_account_types_list = [item.lower() for item in ['Standard_LRS', 'Standard_ZRS', 'Premium_LRS']]
-    storage_account_types_str = ", ".join(storage_account_types_list)
+    TargetRegion, EncryptionImages, OSDiskImageEncryption, DataDiskImageEncryption, \
+        ConfidentialVMEncryptionType, GalleryTargetExtendedLocation, GalleryExtendedLocation = cmd.get_models(
+            'TargetRegion', 'EncryptionImages', 'OSDiskImageEncryption', 'DataDiskImageEncryption',
+            'ConfidentialVMEncryptionType', 'GalleryTargetExtendedLocation', 'GalleryExtendedLocation')
 
     if namespace.target_regions:
         if hasattr(namespace, 'target_region_encryption') and namespace.target_region_encryption:
@@ -2121,6 +2120,9 @@ def process_gallery_image_version_namespace(cmd, namespace):
                     'usage error: Length of --target_region_cvm_encryption should be as same as '
                     'length of target regions')
 
+        storage_account_types_list = [item.lower() for item in ['Standard_LRS', 'Standard_ZRS', 'Premium_LRS']]
+        storage_account_types_str = ", ".join(storage_account_types_list)
+
         regions_info = []
         for i, t in enumerate(namespace.target_regions):
             parts = t.split('=', 2)
@@ -2219,6 +2221,108 @@ def process_gallery_image_version_namespace(cmd, namespace):
 
         namespace.target_regions = regions_info
 
+    if hasattr(namespace, 'target_edge_zones') and namespace.target_edge_zones:
+        if len(namespace.target_edge_zones) == 1 and namespace.target_edge_zones[0].lower() == 'none':
+            namespace.target_edge_zones = []
+            return
+        if hasattr(namespace, 'target_zone_encryption') and namespace.target_zone_encryption:
+            if len(namespace.target_edge_zones) != len(namespace.target_zone_encryption):
+                raise InvalidArgumentValueError(
+                    'usage error: Length of --target-edge-zone-encryption '
+                    'should be as same as length of --target-edge-zones')
+
+        storage_account_types_list = [item.lower() for item in
+                                      ['Standard_LRS', 'Standard_ZRS', 'Premium_LRS', 'StandardSSD_LRS']]
+        storage_account_types_str = ", ".join(storage_account_types_list)
+
+        edge_zone_info = []
+        for i, t in enumerate(namespace.target_edge_zones):
+            parts = t.split('=', 3)
+            # At least the region and edge zone are specified
+            if len(parts) < 2:
+                continue
+
+            region = parts[0]
+            edge_zone = parts[1]
+            replica_count = None
+            storage_account_type = None
+
+            # Both "region" and "edge zone" are specified,
+            # but only one of "replica count" and "storage account type" is specified
+            if len(parts) == 3:
+                try:
+                    replica_count = int(parts[2])
+                except ValueError:
+                    storage_account_type = parts[2]
+                    if parts[2].lower() not in storage_account_types_list:
+                        raise ArgumentUsageError(
+                            "usage error: {} is an invalid target edge zone argument. "
+                            "The third part is neither an integer replica count or a valid storage account type. "
+                            "Storage account types must be one of {}.".format(t, storage_account_types_str))
+
+            # Not only "region" and "edge zone" are specified,
+            # but also "replica count" and "storage account type" are specified
+            elif len(parts) == 4:
+                try:
+                    replica_count = int(parts[2])  # raises ValueError if this is not a replica count, try other order.
+                    storage_account_type = parts[3]
+                    if storage_account_type not in storage_account_types_list:
+                        raise ArgumentUsageError(
+                            "usage error: {} is an invalid target edge zone argument. "
+                            "The forth part is not a valid storage account type. "
+                            "Storage account types must be one of {}.".format(t, storage_account_types_str))
+                except ValueError:
+                    raise ArgumentUsageError(
+                        "usage error: {} is an invalid target edge zone argument. "
+                        "The third part must be a valid integer replica count.".format(t))
+
+            # Parse target edge zone encryption,
+            # example: ['microsoftlosangeles1', 'des1, 0, des2, 1, des3', 'null', 'des4']
+            encryption = None
+            os_disk_image = None
+            data_disk_images = None
+            if hasattr(namespace, 'target_zone_encryption') and namespace.target_zone_encryption:
+                terms = namespace.target_zone_encryption[i].split(',')
+                if len(terms) < 2:
+                    break
+                # OS disk
+                os_disk_image = terms[1]
+                if os_disk_image == 'null':
+                    os_disk_image = None
+                else:
+                    des_id = _disk_encryption_set_format(cmd, namespace, os_disk_image)
+                    os_disk_image = OSDiskImageEncryption(disk_encryption_set_id=des_id)
+                # Data disk
+                if len(terms) > 2:
+                    data_disk_images = terms[2:]
+                    data_disk_images_len = len(data_disk_images)
+                    if data_disk_images_len % 2 != 0:
+                        raise ArgumentUsageError(
+                            'usage error: LUN and disk encryption set for data disk should appear in pair in '
+                            '--target-edge-zone-encryption. Example: 1,osdes,0,datades0,1,datades1')
+                    data_disk_image_encryption_list = []
+                    for j in range(int(data_disk_images_len / 2)):
+                        lun = data_disk_images[j * 2]
+                        des_id = data_disk_images[j * 2 + 1]
+                        des_id = _disk_encryption_set_format(cmd, namespace, des_id)
+                        data_disk_image_encryption_list.append(DataDiskImageEncryption(
+                            lun=lun, disk_encryption_set_id=des_id))
+                    data_disk_images = data_disk_image_encryption_list
+
+            if os_disk_image or data_disk_images:
+                encryption = EncryptionImages(os_disk_image=os_disk_image, data_disk_images=data_disk_images)
+
+            extended_location = GalleryExtendedLocation(name=edge_zone, type='EdgeZone')
+
+            edge_zone_info.append(
+                GalleryTargetExtendedLocation(name=region, extended_location_replica_count=replica_count,
+                                              extended_location=extended_location,
+                                              storage_account_type=storage_account_type,
+                                              encryption=encryption)
+            )
+
+        namespace.target_edge_zones = edge_zone_info
+
 
 def _disk_encryption_set_format(cmd, namespace, name):
     """

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

@@ -572,7 +572,7 @@ def load_command_table(self, _):
 
     with self.command_group('sig image-version', compute_gallery_image_versions_sdk, operation_group='gallery_image_versions', min_api='2018-06-01') as g:
         g.command('delete', 'begin_delete')
-        g.show_command('show', 'get', table_transformer='{Name:name, ResourceGroup:resourceGroup, ProvisioningState:provisioningState, TargetRegions: publishingProfile.targetRegions && join(`, `, publishingProfile.targetRegions[*].name), ReplicationState:replicationStatus.aggregatedState}')
+        g.show_command('show', 'get', table_transformer='{Name:name, ResourceGroup:resourceGroup, ProvisioningState:provisioningState, TargetRegions: publishingProfile.targetRegions && join(`, `, publishingProfile.targetRegions[*].name), EdgeZones: publishingProfile.targetExtendedLocations && join(`, `, publishingProfile.targetExtendedLocations[*].name), ReplicationState:replicationStatus.aggregatedState}')
         g.command('list', 'list_by_gallery_image')
         g.custom_command('create', 'create_image_version', supports_no_wait=True, validator=process_image_version_create_namespace)
         g.generic_update_command('update', getter_name='get_image_version_to_update', setter_arg_name='gallery_image_version', setter_name='update_image_version', setter_type=compute_custom, command_type=compute_custom, supports_no_wait=True, validator=process_image_version_update_namespace)

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

@@ -4539,7 +4539,8 @@ def create_image_version(cmd, resource_group_name, gallery_name, gallery_image_n
                          target_region_encryption=None, os_vhd_uri=None, os_vhd_storage_account=None,
                          data_vhds_uris=None, data_vhds_luns=None, data_vhds_storage_accounts=None,
                          replication_mode=None, target_region_cvm_encryption=None, virtual_machine=None,
-                         image_version=None, allow_replicated_location_deletion=None):
+                         image_version=None, target_zone_encryption=None, target_edge_zones=None,
+                         allow_replicated_location_deletion=None):
     from msrestazure.tools import resource_id, is_valid_resource_id
     from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -4568,6 +4569,10 @@ def create_image_version(cmd, resource_group_name, gallery_name, gallery_image_n
                                             end_of_life_date=end_of_life_date,
                                             target_regions=target_regions or [TargetRegion(name=location)],
                                             replica_count=replica_count, storage_account_type=storage_account_type)
+
+    if target_edge_zones:
+        profile.target_extended_locations = target_edge_zones
+
     if replication_mode is not None:
         profile.replication_mode = replication_mode
     if not cmd.supported_api_version(min_api='2022-03-03', operation_group='gallery_image_versions'):
@@ -4694,7 +4699,7 @@ def get_image_version_to_update(cmd, resource_group_name, gallery_name, gallery_
 
 def update_image_version(cmd, resource_group_name, gallery_name, gallery_image_name, gallery_image_version_name,
                          target_regions=None, replica_count=None, allow_replicated_location_deletion=None,
-                         no_wait=False, **kwargs):
+                         target_edge_zones=None, no_wait=False, **kwargs):
     image_version = kwargs['gallery_image_version']
 
     if target_regions:
@@ -4703,6 +4708,10 @@ def update_image_version(cmd, resource_group_name, gallery_name, gallery_image_n
         image_version.publishing_profile.replica_count = replica_count
     if image_version.storage_profile.source is not None:
         image_version.storage_profile.os_disk_image = image_version.storage_profile.data_disk_images = None
+    # target extended locations will be updated when --target-edge-zones is specified
+    if target_edge_zones is not None:
+        image_version.publishing_profile.target_extended_locations = target_edge_zones \
+            if len(target_edge_zones) > 0 else None
     if allow_replicated_location_deletion is not None:
         image_version.safety_profile.allow_deletion_of_replicated_locations = allow_replicated_location_deletion