{Compute} az vmss restart: Migrate command to aaz-based implementation (#32794)

Author: william051200

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

@@ -20,6 +20,7 @@
 from ._list_skus import *
 from ._patch import *
 from ._perform_maintenance import *
+from ._restart import *
 from ._reimage import *
 from ._reimageall import *
 from ._show import *

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

@@ -0,0 +1,180 @@
+# --------------------------------------------------------------------------------------------
+# 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 Restart(AAZCommand):
+    """Restarts one or more virtual machines in a VM scale set.
+    """
+
+    _aaz_info = {
+        "version": "2024-11-01",
+        "resources": [
+            ["mgmt-plane", "/subscriptions/{}/resourcegroups/{}/providers/microsoft.compute/virtualmachinescalesets/{}/restart", "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",
+        )
+
+        # define Arg Group "VmInstanceIDs"
+
+        _args_schema = cls._args_schema
+        _args_schema.instance_ids = AAZListArg(
+            options=["--instance-ids"],
+            arg_group="VmInstanceIDs",
+            help="The virtual machine scale set instance ids. Omitting the virtual machine scale set instance ids will result in the operation being performed on all virtual machines in the virtual machine scale set.",
+        )
+
+        instance_ids = cls._args_schema.instance_ids
+        instance_ids.Element = AAZStrArg()
+        return cls._args_schema
+
+    def _execute_operations(self):
+        self.pre_operations()
+        yield self.VirtualMachineScaleSetsRestart(ctx=self.ctx)()
+        self.post_operations()
+
+    @register_callback
+    def pre_operations(self):
+        pass
+
+    @register_callback
+    def post_operations(self):
+        pass
+
+    class VirtualMachineScaleSetsRestart(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}/restart",
+                **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
+
+        @property
+        def header_parameters(self):
+            parameters = {
+                **self.serialize_header_param(
+                    "Content-Type", "application/json",
+                ),
+            }
+            return parameters
+
+        @property
+        def content(self):
+            _content_value, _builder = self.new_content_builder(
+                self.ctx.args,
+                typ=AAZObjectType,
+                typ_kwargs={"flags": {"client_flatten": True}}
+            )
+            _builder.set_prop("instanceIds", AAZListType, ".instance_ids")
+
+            instance_ids = _builder.get(".instanceIds")
+            if instance_ids is not None:
+                instance_ids.set_elements(AAZStrType, ".")
+
+            return self.serialize_content(_content_value)
+
+        def on_200(self, session):
+            pass
+
+
+class _RestartHelper:
+    """Helper class for Restart"""
+
+
+__all__ = ["Restart"]

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

@@ -403,10 +403,11 @@ def load_command_table(self, _):
         g.custom_command('identity remove', 'remove_vmss_identity', validator=process_remove_identity_namespace, is_preview=True)
         g.custom_show_command('identity show', 'show_vmss_identity')
         g.custom_command('deallocate', 'deallocate_vmss', supports_no_wait=True)
-        g.custom_command('reimage', 'reimage_vmss', supports_no_wait=True)
         g.custom_command('list-instance-connection-info', 'list_vmss_instance_connection_info')
         g.custom_command('list-instance-public-ips', 'list_vmss_instance_public_ips')
         g.custom_command('list-instances', 'get_instances_list')
+        g.custom_command('reimage', 'reimage_vmss', supports_no_wait=True)
+        g.custom_command('restart', 'restart_vmss', supports_no_wait=True)
 
     with self.command_group('vmss application', operation_group='virtual_machine_scale_sets') as g:
         g.custom_command('set', 'set_vmss_applications', validator=process_set_applications_namespace)
@@ -415,7 +416,6 @@ def load_command_table(self, _):
     with self.command_group('vmss', compute_vmss_sdk, operation_group='virtual_machine_scale_sets') as g:
         g.custom_command('create', 'create_vmss', transform=DeploymentOutputLongRunningOperation(self.cli_ctx, 'Starting vmss create'), supports_no_wait=True, table_transformer=deployment_validate_table_format, validator=process_vmss_create_namespace, exception_handler=handle_template_based_exception)
         g.custom_command('get-instance-view', 'get_vmss_instance_view', table_transformer='{ProvisioningState:statuses[0].displayStatus, PowerState:statuses[1].displayStatus}')
-        g.custom_command('restart', 'restart_vmss', supports_no_wait=True)
         g.custom_command('scale', 'scale_vmss', supports_no_wait=True)
         g.custom_show_command('show', 'get_vmss', table_transformer=get_vmss_table_output_transformer(self, False))
         g.custom_command('stop', 'stop_vmss', supports_no_wait=True, validator=process_vm_vmss_stop)

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

@@ -4548,13 +4548,16 @@ def reimage_vmss(cmd, resource_group_name, vm_scale_set_name, instance_ids=None,
 
 
 def restart_vmss(cmd, resource_group_name, vm_scale_set_name, instance_ids=None, no_wait=False):
-    client = _compute_client_factory(cmd.cli_ctx)
-    VirtualMachineScaleSetVMInstanceRequiredIDs = cmd.get_models('VirtualMachineScaleSetVMInstanceRequiredIDs')
-    if instance_ids is None:
+    from .aaz.latest.vmss import Restart as VmssRestart
+    if not instance_ids:
         instance_ids = ['*']
-    instance_ids = VirtualMachineScaleSetVMInstanceRequiredIDs(instance_ids=instance_ids)
-    return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_restart, resource_group_name, vm_scale_set_name,
-                       vm_instance_i_ds=instance_ids)
+    command_args = {
+        'resource_group': resource_group_name,
+        'vm_scale_set_name': vm_scale_set_name,
+        'instance_ids': instance_ids,
+        'no_wait': no_wait
+    }
+    return VmssRestart(cli_ctx=cmd.cli_ctx)(command_args=command_args)
 
 
 # pylint: disable=inconsistent-return-statements