Skip to content

{Compute} az vmss identity: Migrate commands to aaz-based implementation #32685

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
william051200 wants to merge 7 commits into
base: dev
Choose a base branch
from
+3,169 −2,952
Open

{Compute} az vmss identity: Migrate commands to aaz-based implementation #32685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
011a10f
Migrate vmss identity show command
william051200 Jan 22, 2026
4e9b4be
Migrate vmss identity assign command
william051200 Jan 22, 2026
3d62683
Migrate vmss identity remove command
william051200 Jan 22, 2026
fffa9a6
Add handling for empty variable
william051200 Jan 22, 2026
76d0688
Update test case and re-record test case
william051200 Jan 22, 2026
8a71143
Fix code
william051200 Jan 22, 2026
a2b5bcc
Fix code styling
william051200 Jan 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/azure-cli/azure/cli/command_modules/vm/_vm_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -780,3 +780,9 @@ class IdentityType(Enum):
USER_ASSIGNED = 'UserAssigned'
SYSTEM_ASSIGNED_USER_ASSIGNED = 'SystemAssigned, UserAssigned'
NONE = 'None'


class UpgradeMode(Enum):
AUTOMATIC = 'Automatic'
MANUAL = 'Manual'
ROLLING = 'Rolling'
6 changes: 4 additions & 2 deletions src/azure-cli/azure/cli/command_modules/vm/commands.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,10 +396,12 @@ def load_command_table(self, _):
g.custom_command('create', 'create_dedicated_host_group')
g.generic_update_command('update')

with self.command_group('vmss', compute_vmss_sdk, operation_group='virtual_machine_scale_sets') as g:
with self.command_group('vmss') as g:
g.custom_command('identity assign', 'assign_vmss_identity', validator=process_assign_identity_namespace)
g.custom_command('identity remove', 'remove_vmss_identity', validator=process_remove_identity_namespace, min_api='2017-12-01', is_preview=True)
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')

with self.command_group('vmss', compute_vmss_sdk, operation_group='virtual_machine_scale_sets') as g:
g.custom_command('application set', 'set_vmss_applications', validator=process_set_applications_namespace, min_api='2021-07-01')
g.custom_command('application list', 'list_vmss_applications', min_api='2021-07-01')
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)
Expand Down
209 changes: 126 additions & 83 deletions src/azure-cli/azure/cli/command_modules/vm/custom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -842,8 +842,31 @@ def show_vm_identity(cmd, resource_group_name, vm_name):


def show_vmss_identity(cmd, resource_group_name, vm_name):
client = _compute_client_factory(cmd.cli_ctx)
return client.virtual_machine_scale_sets.get(resource_group_name, vm_name).identity
vm = get_vmss_by_aaz(cmd, resource_group_name, vm_name)

identity = vm.get("identity", {}) if vm else None

if not identity:
return None

user_assigned = identity.get('userAssignedIdentities', {})

if not user_assigned:
identity['userAssignedIdentities'] = None
else:
for user_identity in user_assigned.keys():
if not user_assigned.get(user_identity).get('clientId'):
user_assigned[user_identity]['clientId'] = None
if not user_assigned.get(user_identity).get('principalId'):
user_assigned[user_identity]['principalId'] = None

if not identity.get('principalId'):
identity['principalId'] = None

if not identity.get('tenantId'):
identity['tenantId'] = None

return identity or None
Comment on lines +849 to +869
Copy link
Contributor

@yanzhudd yanzhudd Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

may I ask why we need to make these changes?

Copy link
Member Author

@william051200 william051200 Jan 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The purpose of this change is to match the same output structure before migration.

For example:
Before set "None" to the field

{
  "systemAssignedIdentity": "",
  "userAssignedIdentities": {
    "/subscriptions/fe3bc2b7-22b0-4133-af8f-9439cb831236/resourceGroups/william-rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/id1": {}
  }
}

After set "None" to the field:

{
  "systemAssignedIdentity": "",
  "userAssignedIdentities": {
    "/subscriptions/fe3bc2b7-22b0-4133-af8f-9439cb831236/resourceGroups/william-rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/id1": {
      "clientId": null,
      "principalId": null
    }
  }
}



def assign_vm_identity(cmd, resource_group_name, vm_name, assign_identity=None, identity_role=None,
Expand Down Expand Up @@ -2562,45 +2585,6 @@ def list_vm_extension_images(


# region VirtualMachines Identity
def _remove_identities(cmd, resource_group_name, name, identities, getter, setter):
from ._vm_utils import MSI_LOCAL_ID
ResourceIdentityType = cmd.get_models('ResourceIdentityType', operation_group='virtual_machines')
remove_system_assigned_identity = False
if MSI_LOCAL_ID in identities:
remove_system_assigned_identity = True
identities.remove(MSI_LOCAL_ID)
resource = getter(cmd, resource_group_name, name)
if resource.identity is None:
return None
emsis_to_remove = []
if identities:
existing_emsis = {x.lower() for x in (resource.identity.user_assigned_identities or {}).keys()}
emsis_to_remove = {x.lower() for x in identities}
non_existing = emsis_to_remove.difference(existing_emsis)
if non_existing:
raise CLIError("'{}' are not associated with '{}'".format(','.join(non_existing), name))
if not list(existing_emsis - emsis_to_remove): # if all emsis are gone, we need to update the type
if resource.identity.type == ResourceIdentityType.user_assigned:
resource.identity.type = ResourceIdentityType.none
elif resource.identity.type == ResourceIdentityType.system_assigned_user_assigned:
resource.identity.type = ResourceIdentityType.system_assigned

resource.identity.user_assigned_identities = None
if remove_system_assigned_identity:
resource.identity.type = (ResourceIdentityType.none
if resource.identity.type == ResourceIdentityType.system_assigned
else ResourceIdentityType.user_assigned)

if emsis_to_remove:
if resource.identity.type not in [ResourceIdentityType.none, ResourceIdentityType.system_assigned]:
resource.identity.user_assigned_identities = {}
for identity in emsis_to_remove:
resource.identity.user_assigned_identities[identity] = None

result = LongRunningOperation(cmd.cli_ctx)(setter(resource_group_name, name, resource))
return result.identity


def _remove_identities_by_aaz(cmd, resource_group_name, name, identities, getter, setter):
from ._vm_utils import MSI_LOCAL_ID

Expand Down Expand Up @@ -2645,6 +2629,10 @@ def _remove_identities_by_aaz(cmd, resource_group_name, name, identities, getter
existing_identity['type'] = IdentityType.NONE.value

result = LongRunningOperation(cmd.cli_ctx)(setter(resource_group_name, name, resource))

if not result:
return None

return result.get('identity') or None


Expand Down Expand Up @@ -3587,49 +3575,75 @@ def reset_linux_ssh(cmd, resource_group_name, vm_name, no_wait=False):
# region VirtualMachineScaleSets
def assign_vmss_identity(cmd, resource_group_name, vmss_name, assign_identity=None, identity_role=None,
identity_role_id=None, identity_scope=None):
VirtualMachineScaleSetIdentity, UpgradeMode, ResourceIdentityType, VirtualMachineScaleSetUpdate = cmd.get_models(
'VirtualMachineScaleSetIdentity', 'UpgradeMode', 'ResourceIdentityType', 'VirtualMachineScaleSetUpdate')
IdentityUserAssignedIdentitiesValue = cmd.get_models(
'VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue') or cmd.get_models('UserAssignedIdentitiesValue')
from azure.cli.core.commands.arm import assign_identity as assign_identity_helper
client = _compute_client_factory(cmd.cli_ctx)
_, _, external_identities, enable_local_identity = _build_identities_info(assign_identity)
identity, _, external_identities, enable_local_identity = _build_identities_info(assign_identity)
from ._vm_utils import assign_identity as assign_identity_helper, UpgradeMode

command_args = {'resource_group': resource_group_name, 'vm_scale_set_name': vmss_name}

def getter():
return client.virtual_machine_scale_sets.get(resource_group_name, vmss_name)
return get_vmss_by_aaz(cmd, resource_group_name, vmss_name)

def setter(vmss, external_identities=external_identities):

if vmss.identity and vmss.identity.type == ResourceIdentityType.system_assigned_user_assigned:
identity_types = ResourceIdentityType.system_assigned_user_assigned
elif vmss.identity and vmss.identity.type == ResourceIdentityType.system_assigned and external_identities:
identity_types = ResourceIdentityType.system_assigned_user_assigned
elif vmss.identity and vmss.identity.type == ResourceIdentityType.user_assigned and enable_local_identity:
identity_types = ResourceIdentityType.system_assigned_user_assigned
if vmss.get('identity', {}).get('type', None) == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value:
identity_types = IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value
elif vmss.get('identity', {}).get('type', None) == IdentityType.SYSTEM_ASSIGNED.value and external_identities:
identity_types = IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value
elif vmss.get('identity', {}).get('type', None) == IdentityType.USER_ASSIGNED.value and enable_local_identity:
identity_types = IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value
elif external_identities and enable_local_identity:
identity_types = ResourceIdentityType.system_assigned_user_assigned
identity_types = IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value
elif external_identities:
identity_types = ResourceIdentityType.user_assigned
identity_types = IdentityType.USER_ASSIGNED.value
else:
identity_types = ResourceIdentityType.system_assigned
vmss.identity = VirtualMachineScaleSetIdentity(type=identity_types)
if external_identities:
vmss.identity.user_assigned_identities = {}
for identity in external_identities:
vmss.identity.user_assigned_identities[identity] = IdentityUserAssignedIdentitiesValue()
vmss_patch = VirtualMachineScaleSetUpdate()
vmss_patch.identity = vmss.identity
poller = client.virtual_machine_scale_sets.begin_update(resource_group_name, vmss_name, vmss_patch)
return LongRunningOperation(cmd.cli_ctx)(poller)
identity_types = IdentityType.SYSTEM_ASSIGNED.value

if identity_types == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value:
command_args['mi_system_assigned'] = "True"
command_args['mi_user_assigned'] = []
elif identity_types == IdentityType.USER_ASSIGNED.value:
command_args['mi_user_assigned'] = []
else:
command_args['mi_system_assigned'] = "True"
command_args['mi_user_assigned'] = []

if vmss.get('identity', {}).get('userAssignedIdentities', None):
for key in vmss.get('identity').get('userAssignedIdentities').keys():
command_args['mi_user_assigned'].append(key)

if identity.get('userAssignedIdentities'):
for key in identity.get('userAssignedIdentities', {}).keys():
if key not in command_args['mi_user_assigned']:
command_args['mi_user_assigned'].append(key)

from .operations.vmss import VMSSPatch
update_vmss_identity = VMSSPatch(cli_ctx=cmd.cli_ctx)(command_args=command_args)
LongRunningOperation(cmd.cli_ctx)(update_vmss_identity)
result = update_vmss_identity.result()
return result

assign_identity_helper(cmd.cli_ctx, getter, setter, identity_role=identity_role_id, identity_scope=identity_scope)
vmss = client.virtual_machine_scale_sets.get(resource_group_name, vmss_name)
if vmss.upgrade_policy.mode == UpgradeMode.manual:

vmss = getter()
if vmss.get('upgradePolicy', {}).get('mode', '') == UpgradeMode.MANUAL.value:
logger.warning("With manual upgrade mode, you will need to run 'az vmss update-instances -g %s -n %s "
"--instance-ids *' to propagate the change", resource_group_name, vmss_name)

return _construct_identity_info(identity_scope, identity_role, vmss.identity.principal_id,
vmss.identity.user_assigned_identities)
info = _construct_identity_info(
identity_scope,
identity_role,
vmss.get('identity').get('principalId') if vmss.get('identity') else None,
vmss.get('identity').get('userAssignedIdentities') if vmss.get('identity') else None)

if not info.get('userAssignedIdentities'):
return info

for user_identity in info.get('userAssignedIdentities', {}).keys():
if not info['userAssignedIdentities'][user_identity].get('clientId'):
info['userAssignedIdentities'][user_identity]['clientId'] = None
if not info['userAssignedIdentities'][user_identity].get('principalId'):
info['userAssignedIdentities'][user_identity]['principalId'] = None

return info


# pylint: disable=too-many-locals, too-many-statements
Expand Down Expand Up @@ -4199,6 +4213,24 @@ def get_vmss(cmd, resource_group_name, name, instance_id=None, include_user_data
return client.virtual_machine_scale_sets.get(resource_group_name, name)


def get_vmss_by_aaz(cmd, resource_group_name, name, instance_id=None, include_user_data=False):
from .operations.vmss import VMSSShow
from .operations.vmss_vms import VMSSVMSShow

command_args = {
'resource_group': resource_group_name,
'vm_scale_set_name': name,
}

if include_user_data:
command_args['expand'] = 'userData'

if instance_id is not None:
Copy link

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

get_vmss_by_aaz accepts an instance_id argument but the instance_id is never added to command_args before calling VMSSVMSShow, so any future caller that passes instance_id will not get the expected per-instance result. Please include the instance_id in command_args when instance_id is not None, consistent with how get_vmss_modified_by_aaz and other helpers wire through this parameter.

Suggested change
if instance_id is not None:
if instance_id is not None:
command_args['instance_id'] = instance_id

Copilot uses AI. Check for mistakes.
command_args['instance_id'] = instance_id
return VMSSVMSShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)
return VMSSShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)


def _check_vmss_hyper_v_generation(cli_ctx, vmss):
hyper_v_generation = get_hyper_v_generation_from_vmss(
cli_ctx, vmss.virtual_machine_profile.storage_profile.image_reference, vmss.location)
Expand Down Expand Up @@ -5403,24 +5435,35 @@ def vmss_run_command_show(cmd,

# region VirtualMachineScaleSets Identity
def remove_vmss_identity(cmd, resource_group_name, vmss_name, identities=None):
client = _compute_client_factory(cmd.cli_ctx)
def setter(resource_group_name, vmss_name, vmss):
command_args = {
'resource_group': resource_group_name,
'vm_scale_set_name': vmss_name
}

def _get_vmss(_, resource_group_name, vmss_name):
return client.virtual_machine_scale_sets.get(resource_group_name, vmss_name)
if vmss.get('identity') and vmss['identity'].get('type') == IdentityType.USER_ASSIGNED.value:
# NOTE: The literal 'UserAssigned' is intentionally appended as a marker for
# VMSSIdentityRemove._format_content, which uses it to apply special handling
# for purely user-assigned identities. It is not a real identity resource ID.
command_args['mi_user_assigned'] = \
list(vmss.get('identity', {}).get('userAssignedIdentities', {}).keys()) + ['UserAssigned']
elif vmss.get('identity') and vmss['identity'].get('type') == IdentityType.SYSTEM_ASSIGNED.value:
command_args['mi_user_assigned'] = []
command_args['mi_system_assigned'] = 'True'
elif vmss.get('identity') and vmss['identity'].get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value:
command_args['mi_user_assigned'] = list(vmss.get('identity', {}).get('userAssignedIdentities', {}).keys())
command_args['mi_system_assigned'] = 'True'
else:
command_args['mi_user_assigned'] = []

def _set_vmss(resource_group_name, name, vmss_instance):
VirtualMachineScaleSetUpdate = cmd.get_models('VirtualMachineScaleSetUpdate',
operation_group='virtual_machine_scale_sets')
vmss_update = VirtualMachineScaleSetUpdate(identity=vmss_instance.identity)
return client.virtual_machine_scale_sets.begin_update(resource_group_name, vmss_name, vmss_update)
from .operations.vmss import VMSSIdentityRemove
return VMSSIdentityRemove(cli_ctx=cmd.cli_ctx)(command_args=command_args)

if identities is None:
from ._vm_utils import MSI_LOCAL_ID
identities = [MSI_LOCAL_ID]

return _remove_identities(cmd, resource_group_name, vmss_name, identities,
_get_vmss,
_set_vmss)
return _remove_identities_by_aaz(cmd, resource_group_name, vmss_name, identities, get_vmss_by_aaz, setter)
# endregion


Expand Down
Loading