aaz update vm/vmss

Author: naga-nandyala

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

@@ -1681,26 +1681,6 @@ def _output(self, *args, **kwargs):
         }))
         vm = get_vm_to_update(cmd, resource_group_name, vm_name)
 
-    if add_proxy_agent_extension is not None:
-        from .aaz.latest.vm import Update as _Update
-
-        class Update(_Update):
-            pass
-
-        args = {
-            'resource_group': resource_group_name,
-            'vm_name': vm_name,
-            'no_wait': no_wait,
-            'security_profile': {
-                'proxy_agent_settings': {
-                    'add_proxy_agent_extension': add_proxy_agent_extension
-                }
-            }
-        }
-
-        LongRunningOperation(cmd.cli_ctx)(Update(cli_ctx=cmd.cli_ctx)(command_args=args))
-        vm = get_vm_to_update(cmd, resource_group_name, vm_name)
-
     disk_name = None
     if os_disk is not None:
         if is_valid_resource_id(os_disk):
@@ -1916,12 +1896,66 @@ class Update(_Update):
                 vm.scheduled_events_policy.user_initiated_reboot = {
                     "automaticallyApprove": enable_user_reboot_scheduled_events
                 }
-    client = _compute_client_factory(cmd.cli_ctx, aux_subscriptions=aux_subscriptions)
-    if wire_server_access_control_profile_reference_id is not None or \
-            imds_access_control_profile_reference_id is not None or \
-            add_proxy_agent_extension is not None:
-        kwargs['parameters'] = vm
-    return sdk_no_wait(no_wait, client.virtual_machines.begin_create_or_update, resource_group_name, vm_name, **kwargs)
+
+    # ==============================================================================
+    # AAZ Update: Use SDK for complex logic, then apply via AAZ
+    # ==============================================================================
+    # All SDK modifications to 'vm' object have already happened above (storage, security, etc.)
+    # Now we just need to apply those changes using AAZ's Update command
+    
+    class VMUpdateCustom(UpdateVM):
+        def pre_instance_update(self, instance):
+            """
+            Copy SDK VM object properties into the AAZ instance before the update.
+            This ensures AAZ sees the current state, but does NOT override user args.
+            """
+            vm_dict = vm.serialize()
+            vm_dict.pop('resources', None)  # Remove conflicting resources
+            
+            for key, value in vm_dict.items():
+                try:
+                    instance[key] = value
+                except Exception as e:
+                    logger.debug(f"Could not set {key}: {e}")
+        
+        def post_instance_update(self, instance):
+            """
+            Set AAZ-only fields after AAZ processes command_args.
+            This ensures user arguments (--enable-proxy-agent, etc.) take precedence.
+            """
+            if add_proxy_agent_extension is not None:
+                props = getattr(instance, 'properties', None)
+                sp = getattr(props, 'security_profile', None) if props else None
+                pas = getattr(sp, 'proxy_agent_settings', None) if sp else None
+                
+                if pas:
+                    pas.add_proxy_agent_extension = add_proxy_agent_extension
+                else:
+                    logger.debug("Cannot set add_proxy_agent_extension - proxy_agent_settings does not exist")
+        
+        def _output(self, *args, **kwargs):
+            """
+            Remove resources before flattening to avoid KeyError issues during output serialization.
+            """
+            self.ctx.vars.instance.resources = None
+            return self.deserialize_output(self.ctx.vars.instance, client_flatten=True)
+    
+    args = {
+        'resource_group': resource_group_name,
+        'vm_name': vm_name,
+        'no_wait': no_wait
+    }
+    
+    return VMUpdateCustom(cli_ctx=cmd.cli_ctx)(command_args=args)    
+
+
+
+    # client = _compute_client_factory(cmd.cli_ctx, aux_subscriptions=aux_subscriptions)
+    # if wire_server_access_control_profile_reference_id is not None or \
+    #         imds_access_control_profile_reference_id is not None or \
+    #         add_proxy_agent_extension is not None:
+    #     kwargs['parameters'] = vm
+    # return sdk_no_wait(no_wait, client.virtual_machines.begin_create_or_update, resource_group_name, vm_name, **kwargs)
 # endregion
 
 
@@ -4221,28 +4255,6 @@ def _output(self, *args, **kwargs):
         }))
         vmss = get_vmss_modified(cmd, resource_group_name, name, instance_id, security_type)
 
-    if add_proxy_agent_extension is not None:
-        from .aaz.latest.vmss import Update as _VMSSUpdate
-
-        class VMSSUpdate(_VMSSUpdate):
-            pass
-
-        args = {
-            'resource_group': resource_group_name,
-            'vm_scale_set_name': name,
-            'no_wait': no_wait,
-            'virtual_machine_profile': {
-                'security_profile': {
-                    'proxy_agent_settings': {
-                        'add_proxy_agent_extension': add_proxy_agent_extension
-                    }
-                }
-            }
-        }
-
-        LongRunningOperation(cmd.cli_ctx)(VMSSUpdate(cli_ctx=cmd.cli_ctx)(command_args=args))
-        vmss = get_vmss_modified(cmd, resource_group_name, name, instance_id, security_type)
-
     aux_subscriptions = None
     # pylint: disable=too-many-boolean-expressions
     if vmss and hasattr(vmss, 'virtual_machine_profile') and vmss.virtual_machine_profile and \
@@ -4586,14 +4598,58 @@ class VMSSUpdate(_VMSSUpdate):
 
     if zone_balance is not None:
         vmss.zone_balance = zone_balance
-
-    if wire_server_access_control_profile_reference_id is not None or \
-            imds_access_control_profile_reference_id is not None or \
-            add_proxy_agent_extension is not None:
-        kwargs['parameters'] = vmss
-
-    return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,
-                       resource_group_name, name, **kwargs)
+    
+    # if wire_server_access_control_profile_reference_id is not None or \
+    #         imds_access_control_profile_reference_id is not None or \
+    #         add_proxy_agent_extension is not None:
+    #     kwargs['parameters'] = vmss
+
+    # return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,
+    #                    resource_group_name, name, **kwargs)    
+
+    # ==============================================================================
+    # AAZ Update: Use SDK for complex logic, then apply via AAZ (same pattern as update_vm)
+    # ==============================================================================
+    from .aaz.latest.vmss import Update as UpdateVMSS
+    
+    class VMSSUpdateCustom(UpdateVMSS):
+        def pre_instance_update(self, instance):
+            """
+            Copy SDK VMSS object properties into the AAZ instance before the update.
+            This ensures AAZ sees the current state, but does NOT override user args.
+            """
+            vmss_dict = vmss.serialize()
+            
+            for key, value in vmss_dict.items():
+                try:
+                    instance[key] = value
+                except Exception as e:
+                    logger.debug(f"Could not set {key}: {e}")
+        
+        def post_instance_update(self, instance):
+            """
+            Set AAZ-only fields after AAZ processes command_args.
+            This ensures user arguments take precedence.
+            """
+            if add_proxy_agent_extension is not None:
+                # Navigate through nested structure safely (VMSS has extra level)
+                props = getattr(instance, 'properties', None)
+                vm_profile = getattr(props, 'virtual_machine_profile', None) if props else None
+                sp = getattr(vm_profile, 'security_profile', None) if vm_profile else None
+                pas = getattr(sp, 'proxy_agent_settings', None) if sp else None
+                
+                if pas:
+                    pas.add_proxy_agent_extension = add_proxy_agent_extension
+                else:
+                    logger.debug("Cannot set add_proxy_agent_extension - proxy_agent_settings does not exist")
+    
+    args = {
+        'resource_group': resource_group_name,
+        'vm_scale_set_name': name,
+        'no_wait': no_wait
+    }
+    
+    return VMSSUpdateCustom(cli_ctx=cmd.cli_ctx)(command_args=args)
 
 # endregion