{Compute} az vmss extension: Migrate command group to aaz-based implementation

src/azure-cli/azure/cli/command_modules/vm/aaz/latest/vmss/extension/__init__.py

@@ -9,3 +9,4 @@
 # flake8: noqa
 
 from .__cmd_group import *
+from ._rolling_upgrade import *

src/azure-cli/azure/cli/command_modules/vm/aaz/latest/vmss/extension/_rolling_upgrade.py

@@ -0,0 +1,144 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+#
+# Code generated by aaz-dev-tools
+# --------------------------------------------------------------------------------------------
+
+# pylint: skip-file
+# flake8: noqa
+
+from azure.cli.core.aaz import *
+
+
+class RollingUpgrade(AAZCommand):
+    """Starts a rolling upgrade to move all extensions for all virtual machine scale set instances to the latest available extension version. Instances which are already running the latest extension versions are not affected.
+    """
+
+    _aaz_info = {
+        "version": "2024-11-01",
+        "resources": [
+            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/virtualmachinescalesets/{}/extensionrollingupgrade", "2024-11-01"],
+        ]
+    }
+
+    AZ_SUPPORT_NO_WAIT = True
+
+    def _handler(self, command_args):
+        super()._handler(command_args)
+        return self.build_lro_poller(self._execute_operations, None)
+
+    _args_schema = None
+
+    @classmethod
+    def _build_arguments_schema(cls, *args, **kwargs):
+        if cls._args_schema is not None:
+            return cls._args_schema
+        cls._args_schema = super()._build_arguments_schema(*args, **kwargs)
+
+        # define Arg Group ""
+
+        _args_schema = cls._args_schema
+        _args_schema.resource_group = AAZResourceGroupNameArg(
+            required=True,
+        )
+        _args_schema.vm_scale_set_name = AAZStrArg(
+            options=["--vm-scale-set-name"],
+            help="The name of the VM scale set.",
+            required=True,
+            id_part="name",
+        )
+        return cls._args_schema
+
+    def _execute_operations(self):
+        self.pre_operations()
+        yield self.VirtualMachineScaleSetRollingUpgradesStartExtensionUpgrade(ctx=self.ctx)()
+        self.post_operations()
+
+    @register_callback
+    def pre_operations(self):
+        pass
+
+    @register_callback
+    def post_operations(self):
+        pass
+
+    class VirtualMachineScaleSetRollingUpgradesStartExtensionUpgrade(AAZHttpOperation):
+        CLIENT_TYPE = "MgmtClient"
+
+        def __call__(self, *args, **kwargs):
+            request = self.make_request()
+            session = self.client.send_request(request=request, stream=False, **kwargs)
+            if session.http_response.status_code in [202]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "location"},
+                    path_format_arguments=self.url_parameters,
+                )
+            if session.http_response.status_code in [200]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "location"},
+                    path_format_arguments=self.url_parameters,
+                )
+
+            return self.on_error(session.http_response)
+
+        @property
+        def url(self):
+            return self.client.format_url(
+                "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Compute/virtualMachineScaleSets/{vmScaleSetName}/extensionRollingUpgrade",
+                **self.url_parameters
+            )
+
+        @property
+        def method(self):
+            return "POST"
+
+        @property
+        def error_format(self):
+            return "ODataV4Format"
+
+        @property
+        def url_parameters(self):
+            parameters = {
+                **self.serialize_url_param(
+                    "resourceGroupName", self.ctx.args.resource_group,
+                    required=True,
+                ),
+                **self.serialize_url_param(
+                    "subscriptionId", self.ctx.subscription_id,
+                    required=True,
+                ),
+                **self.serialize_url_param(
+                    "vmScaleSetName", self.ctx.args.vm_scale_set_name,
+                    required=True,
+                ),
+            }
+            return parameters
+
+        @property
+        def query_parameters(self):
+            parameters = {
+                **self.serialize_query_param(
+                    "api-version", "2024-11-01",
+                    required=True,
+                ),
+            }
+            return parameters
+
+        def on_200(self, session):
+            pass
+
+
+class _RollingUpgradeHelper:
+    """Helper class for RollingUpgrade"""
+
+
+__all__ = ["RollingUpgrade"]

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

@@ -407,12 +407,12 @@ def load_command_table(self, _):
         g.custom_command('disable', 'decrypt_vmss')
         g.custom_show_command('show', 'show_vmss_encryption_status')
 
-    with self.command_group('vmss extension', compute_vmss_sdk) as g:
+    with self.command_group('vmss extension') as g:
         g.custom_command('delete', 'delete_vmss_extension', supports_no_wait=True)
         g.custom_show_command('show', 'get_vmss_extension')
         g.custom_command('set', 'set_vmss_extension', supports_no_wait=True)
         g.custom_command('list', 'list_vmss_extensions')
-        g.custom_command('upgrade', 'upgrade_vmss_extension', min_api='2020-06-01', supports_no_wait=True)
+        g.custom_command('upgrade', 'upgrade_vmss_extension', supports_no_wait=True)
 
     with self.command_group('vmss extension image', compute_vm_extension_image_sdk) as g:
         g.custom_command('list', 'list_vm_extension_images')

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

@@ -5283,102 +5283,109 @@ def detach_disk_from_vmss(cmd, resource_group_name, vmss_name, lun, instance_id=
 
 # region VirtualMachineScaleSets Extensions
 def delete_vmss_extension(cmd, resource_group_name, vmss_name, extension_name):
-    client = _compute_client_factory(cmd.cli_ctx)
-    vmss = client.virtual_machine_scale_sets.get(resource_group_name=resource_group_name, vm_scale_set_name=vmss_name)
+    from .aaz.latest.vmss import Update as VMSSUpdate
+    from .operations.vmss import convert_show_result_to_snake_case
+
+    vmss = get_vmss_by_aaz(cmd, resource_group_name, vmss_name)
     # Avoid unnecessary permission error
-    vmss.virtual_machine_profile.storage_profile.image_reference = None
+    if vmss.get('virtualMachineProfile', {}).get('storageProfile', {}):
+        vmss['virtualMachineProfile']['storageProfile']['imageReference'] = None
     # pylint: disable=no-member
-    if not vmss.virtual_machine_profile.extension_profile:
+    if not vmss.get('virtualMachineProfile', {}).get('extensionProfile'):
         raise CLIError('Scale set has no extensions to delete')
 
-    keep_list = [e for e in vmss.virtual_machine_profile.extension_profile.extensions
-                 if e.name != extension_name]
-    if len(keep_list) == len(vmss.virtual_machine_profile.extension_profile.extensions):
+    keep_list = [e for e in vmss.get('virtualMachineProfile', {}).get('extensionProfile', {}).get('extensions', [])
+                 if e.get('name') != extension_name]
+    if len(keep_list) == len(vmss.get('virtualMachineProfile', {}).get('extensionProfile', {}).get('extensions', [])):
         raise CLIError('Extension {} not found'.format(extension_name))
 
-    vmss.virtual_machine_profile.extension_profile.extensions = keep_list
+    if not vmss.get('virtualMachineProfile'):
+        vmss['virtualMachineProfile'] = {}
+    if not vmss.get('virtualMachineProfile', {}).get('extensionProfile'):
+        vmss['virtualMachineProfile']['extensionProfile'] = {}
+
+    vmss['virtualMachineProfile']['extensionProfile']['extensions'] = keep_list
 
-    return client.virtual_machine_scale_sets.begin_create_or_update(resource_group_name=resource_group_name,
-                                                                    vm_scale_set_name=vmss_name, parameters=vmss)
+    vmss = convert_show_result_to_snake_case(vmss)
+    vmss['resource_group'] = resource_group_name
+    vmss['vm_scale_set_name'] = vmss_name
+
+    return VMSSUpdate(cli_ctx=cmd.cli_ctx)(command_args=vmss)
 
 
 # pylint: disable=inconsistent-return-statements
 def get_vmss_extension(cmd, resource_group_name, vmss_name, extension_name):
-    client = _compute_client_factory(cmd.cli_ctx)
-    vmss = client.virtual_machine_scale_sets.get(resource_group_name=resource_group_name, vm_scale_set_name=vmss_name)
-    # pylint: disable=no-member
-    if not vmss.virtual_machine_profile.extension_profile:
-        return
-    return next((e for e in vmss.virtual_machine_profile.extension_profile.extensions
-                 if e.name == extension_name), None)
+    vmss = get_vmss_by_aaz(cmd, resource_group_name, vmss_name)
+    return next((e for e in vmss.get('virtualMachineProfile', {}).get('extensionProfile', {}).get('extensions', [])
+                 if e.get('name') == extension_name), None)
 
 
 def list_vmss_extensions(cmd, resource_group_name, vmss_name):
-    client = _compute_client_factory(cmd.cli_ctx)
-    vmss = client.virtual_machine_scale_sets.get(resource_group_name=resource_group_name, vm_scale_set_name=vmss_name)
-    # pylint: disable=no-member
-    if vmss.virtual_machine_profile and vmss.virtual_machine_profile.extension_profile:
-        return vmss.virtual_machine_profile.extension_profile.extensions
+    vmss = get_vmss_by_aaz(cmd, resource_group_name, vmss_name)
+    if vmss.get('virtualMachineProfile', {}).get('extensionProfile'):
+        return vmss['virtualMachineProfile']['extensionProfile'].get('extensions', [])
     return None
 
 
 def set_vmss_extension(cmd, resource_group_name, vmss_name, extension_name, publisher, version=None,
                        settings=None, protected_settings=None, no_auto_upgrade=False, force_update=False,
                        no_wait=False, extension_instance_name=None, provision_after_extensions=None,
                        enable_auto_upgrade=None):
+    from .aaz.latest.vmss import Update as VMSSUpdate
+    from .operations.vmss import convert_show_result_to_snake_case
+
     if not extension_instance_name:
         extension_instance_name = extension_name
 
     auto_upgrade_extensions = ['CodeIntegrityAgent']
     if extension_name in auto_upgrade_extensions and enable_auto_upgrade is None:
         enable_auto_upgrade = True
 
-    client = _compute_client_factory(cmd.cli_ctx)
-    vmss = client.virtual_machine_scale_sets.get(resource_group_name=resource_group_name, vm_scale_set_name=vmss_name)
+    vmss = get_vmss_by_aaz(cmd, resource_group_name, vmss_name)
     # Avoid unnecessary permission error
-    vmss.virtual_machine_profile.storage_profile.image_reference = None
-    VirtualMachineScaleSetExtension, VirtualMachineScaleSetExtensionProfile = cmd.get_models(
-        'VirtualMachineScaleSetExtension', 'VirtualMachineScaleSetExtensionProfile')
+    if vmss.get('virtualMachineProfile', {}).get('storageProfile', {}):
+        vmss['virtualMachineProfile']['storageProfile']['imageReference'] = None
 
-    # pylint: disable=no-member
-    version = _normalize_extension_version(cmd.cli_ctx, publisher, extension_name, version, vmss.location)
-    extension_profile = vmss.virtual_machine_profile.extension_profile
+    version = _normalize_extension_version(cmd.cli_ctx, publisher, extension_name, version, vmss.get('location'))
+    extension_profile = vmss.get('virtualMachineProfile', {}).get('extensionProfile')
     if extension_profile:
-        extensions = extension_profile.extensions
+        extensions = extension_profile.get('extensions')
         if extensions:
-            extension_profile.extensions = [x for x in extensions if
-                                            x.type_properties_type.lower() != extension_name.lower() or x.publisher.lower() != publisher.lower()]  # pylint: disable=line-too-long
-
-    if cmd.supported_api_version(min_api='2019-07-01', operation_group='virtual_machine_scale_sets'):
-        ext = VirtualMachineScaleSetExtension(name=extension_instance_name,
-                                              publisher=publisher,
-                                              type_properties_type=extension_name,
-                                              protected_settings=protected_settings,
-                                              type_handler_version=version,
-                                              settings=settings,
-                                              auto_upgrade_minor_version=(not no_auto_upgrade),
-                                              provision_after_extensions=provision_after_extensions,
-                                              enable_automatic_upgrade=enable_auto_upgrade)
-    else:
-        ext = VirtualMachineScaleSetExtension(name=extension_instance_name,
-                                              publisher=publisher,
-                                              type=extension_name,
-                                              protected_settings=protected_settings,
-                                              type_handler_version=version,
-                                              settings=settings,
-                                              auto_upgrade_minor_version=(not no_auto_upgrade),
-                                              provision_after_extensions=provision_after_extensions,
-                                              enable_automatic_upgrade=enable_auto_upgrade)
+            extension_profile['extensions'] = \
+                [x for x in extensions if
+                 x.get('type_properties_type', '').lower() != extension_name.lower()
+                 or x.get('publisher').lower() != publisher.lower()]
+
+    ext = {
+        'name': extension_instance_name,
+        'auto_upgrade_minor_version': (not no_auto_upgrade),
+        'enable_automatic_upgrade': enable_auto_upgrade,
+        'protected_settings': protected_settings,
+        'provision_after_extensions': provision_after_extensions,
+        'publisher': publisher,
+        'settings': settings,
+        'type': extension_name,
+        'type_handler_version': version
+    }
 
     if force_update:
-        ext.force_update_tag = str(_gen_guid())
+        ext['force_update_tag'] = str(_gen_guid())
 
-    if not vmss.virtual_machine_profile.extension_profile:
-        vmss.virtual_machine_profile.extension_profile = VirtualMachineScaleSetExtensionProfile(extensions=[])
-    vmss.virtual_machine_profile.extension_profile.extensions.append(ext)
+    vmss = convert_show_result_to_snake_case(vmss)
 
-    return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,
-                       resource_group_name, vmss_name, vmss)
+    if not vmss.get('virtual_machine_profile'):
+        vmss['virtual_machine_profile'] = {}
+    if not vmss.get('virtual_machine_profile', {}).get('extension_profile'):
+        vmss['virtual_machine_profile']['extension_profile'] = {}
+    if not vmss.get('virtual_machine_profile', {}).get('extension_profile').get('extensions'):
+        vmss['virtual_machine_profile']['extension_profile']['extensions'] = []
+    vmss['virtual_machine_profile']['extension_profile']['extensions'].append(ext)
+
+    vmss['resource_group'] = resource_group_name
+    vmss['vm_scale_set_name'] = vmss_name
+    vmss['no_wait'] = no_wait
+
+    return VMSSUpdate(cli_ctx=cmd.cli_ctx)(command_args=vmss)
 
 
 def set_orchestration_service_state(cmd, resource_group_name, vm_scale_set_name, service_name, action, no_wait=False):
@@ -5397,9 +5404,13 @@ def set_orchestration_service_state(cmd, resource_group_name, vm_scale_set_name,
 
 
 def upgrade_vmss_extension(cmd, resource_group_name, vm_scale_set_name, no_wait=False):
-    client = _compute_client_factory(cmd.cli_ctx)
-    return sdk_no_wait(no_wait, client.virtual_machine_scale_set_rolling_upgrades.begin_start_extension_upgrade,
-                       resource_group_name, vm_scale_set_name)
+    from .aaz.latest.vmss.extension import RollingUpgrade as VMSSRollingUpgrade
+    command_args = {
+        'resource_group': resource_group_name,
+        'vm_scale_set_name': vm_scale_set_name,
+        'no_wait': no_wait
+    }
+    return VMSSRollingUpgrade(cli_ctx=cmd.cli_ctx)(command_args=command_args)
 # endregion