{Compute} az vm identity: Migrate commands to aaz-based implementation

src/azure-cli-core/azure/cli/core/commands/arm.py

@@ -759,46 +759,50 @@ def _find_property(instance, path):
 
 
 def assign_identity(cli_ctx, getter, setter, identity_role=None, identity_scope=None):
-    import time
-    from azure.core.exceptions import HttpResponseError
-
     # get
     resource = getter()
     resource = setter(resource)
 
     # create role assignment:
     if identity_scope:
         principal_id = resource.identity.principal_id
+        create_role_assignment(cli_ctx, principal_id, identity_role, identity_scope)
 
-        identity_role_id = resolve_role_id(cli_ctx, identity_role, identity_scope)
-        assignments_client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_assignments
-        RoleAssignmentCreateParameters = get_sdk(cli_ctx, ResourceType.MGMT_AUTHORIZATION,
-                                                 'RoleAssignmentCreateParameters', mod='models',
-                                                 operation_group='role_assignments')
-        parameters = RoleAssignmentCreateParameters(role_definition_id=identity_role_id, principal_id=principal_id,
-                                                    principal_type=None)
-
-        logger.info("Creating an assignment with a role '%s' on the scope of '%s'", identity_role_id, identity_scope)
-        retry_times = 36
-        assignment_name = _gen_guid()
-        for retry_time in range(0, retry_times):
-            try:
-                assignments_client.create(scope=identity_scope, role_assignment_name=assignment_name,
-                                          parameters=parameters)
-                break
-            except HttpResponseError as ex:
-                if ex.error.code == 'RoleAssignmentExists':
-                    logger.info('Role assignment already exists')
-                    break
-                if retry_time < retry_times and ' does not exist in the directory ' in ex.message:
-                    time.sleep(5)
-                    logger.warning('Retrying role assignment creation: %s/%s', retry_time + 1,
-                                   retry_times)
-                    continue
-                raise
     return resource
 
 
+def create_role_assignment(cli_ctx, principal_id, identity_role=None, identity_scope=None):
+    import time
+    from azure.core.exceptions import HttpResponseError
+
+    identity_role_id = resolve_role_id(cli_ctx, identity_role, identity_scope)
+    assignments_client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_assignments
+    RoleAssignmentCreateParameters = get_sdk(cli_ctx, ResourceType.MGMT_AUTHORIZATION,
+                                             'RoleAssignmentCreateParameters', mod='models',
+                                             operation_group='role_assignments')
+    parameters = RoleAssignmentCreateParameters(role_definition_id=identity_role_id, principal_id=principal_id,
+                                                principal_type=None)
+
+    logger.info("Creating an assignment with a role '%s' on the scope of '%s'", identity_role_id, identity_scope)
+    retry_times = 36
+    assignment_name = _gen_guid()
+    for retry_time in range(0, retry_times):
+        try:
+            assignments_client.create(scope=identity_scope, role_assignment_name=assignment_name,
+                                      parameters=parameters)
+            break
+        except HttpResponseError as ex:
+            if ex.error.code == 'RoleAssignmentExists':
+                logger.info('Role assignment already exists')
+                break
+            if retry_time < retry_times and ' does not exist in the directory ' in ex.message:
+                time.sleep(5)
+                logger.warning('Retrying role assignment creation: %s/%s', retry_time + 1,
+                               retry_times)
+                continue
+            raise
+
+
 def resolve_role_id(cli_ctx, role, scope):
     import uuid
     client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_definitions

src/azure-cli/azure/cli/command_modules/sqlvm/tests/latest/test_sqlvm_commands.py

@@ -48,7 +48,7 @@ def __init__(self, name_prefix=sqlvm_name_prefix, location='westus',
     def create_resource(self, name, **kwargs):
         group = self._get_resource_group(**kwargs)
         template = ('az vm create -l {} -g {} -n {} --admin-username {} --admin-password {} --image {}'
-                    ' --size Standard_DS2_v2 --nsg-rule NONE')
+                    ' --size Standard_B2ms --nsg-rule NONE')
         execute(DummyCli(), template.format(self.location, group, name, self.vm_user, self.vm_password, self.image))
         return {self.parameter_name: name}
 

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

@@ -7,10 +7,13 @@
 import os
 import re
 import importlib
+from enum import Enum
 
 from urllib.parse import urlparse
 
 from azure.cli.core.commands.arm import ArmTemplateBuilder
+from azure.cli.core.commands.client_factory import get_mgmt_service_client
+from azure.cli.core.profiles import ResourceType
 
 from knack.log import get_logger
 from knack.util import CLIError
@@ -32,7 +35,7 @@ def get_target_network_api(cli_ctx):
     if cli_ctx.cloud.profile == 'latest':
         version = '2022-01-01'
     else:
-        from azure.cli.core.profiles import get_api_version, ResourceType
+        from azure.cli.core.profiles import get_api_version
         version = get_api_version(cli_ctx, ResourceType.MGMT_NETWORK)
     return version
 
@@ -46,8 +49,6 @@ def read_content_if_is_file(string_or_file):
 
 
 def _resolve_api_version(cli_ctx, provider_namespace, resource_type, parent_path):
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from azure.cli.core.profiles import ResourceType
     client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES)
     provider = client.providers.get(provider_namespace)
 
@@ -75,10 +76,8 @@ def log_pprint_template(template):
 def check_existence(cli_ctx, value, resource_group, provider_namespace, resource_type,
                     parent_name=None, parent_type=None, static_version=None):
     # check for name or ID and set the type flags
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     from azure.core.exceptions import HttpResponseError
     from azure.mgmt.core.tools import parse_resource_id
-    from azure.cli.core.profiles import ResourceType
     id_parts = parse_resource_id(value)
     resource_client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES,
                                               subscription_id=id_parts.get('subscription', None)).resources
@@ -414,8 +413,6 @@ def _update(model, lun, value):
 
 
 def get_storage_blob_uri(cli_ctx, storage):
-    from azure.cli.core.profiles._shared import ResourceType
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     if urlparse(storage).scheme:
         storage_uri = storage
     else:
@@ -757,3 +754,29 @@ def _open(filename, mode):
         f.write(public_bytes)
 
     return public_bytes.decode()
+
+
+def _gen_guid():
+    import uuid
+    return uuid.uuid4()
+
+
+def assign_identity(cli_ctx, getter, setter, identity_role=None, identity_scope=None):
+    from azure.cli.core.commands.arm import create_role_assignment
+
+    # get
+    resource = getter()
+    resource = setter(resource)
+
+    # create role assignment:
+    if identity_scope:
+        principal_id = resource.get('identity', {}).get('principalId') or resource.get('identity', {}).get('principal_id')
+        create_role_assignment(cli_ctx, principal_id, identity_role, identity_scope)
+    return resource
+
+
+class IdentityType(Enum):
+    SYSTEM_ASSIGNED = 'SystemAssigned'
+    USER_ASSIGNED = 'UserAssigned'
+    SYSTEM_ASSIGNED_USER_ASSIGNED = 'SystemAssigned, UserAssigned'
+    NONE = 'None'

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

@@ -288,11 +288,12 @@ def load_command_table(self, _):
         from .operations.snapshot import SnapshotUpdate
         self.command_table['snapshot update'] = SnapshotUpdate(loader=self)
 
-    with self.command_group('vm', compute_vm_sdk) as g:
-        g.custom_command('identity assign', 'assign_vm_identity', validator=process_assign_identity_namespace)
-        g.custom_command('identity remove', 'remove_vm_identity', validator=process_remove_identity_namespace, min_api='2017-12-01')
-        g.custom_show_command('identity show', 'show_vm_identity')
+    with self.command_group('vm identity') as g:
+        g.custom_command('assign', 'assign_vm_identity', validator=process_assign_identity_namespace)
+        g.custom_command('remove', 'remove_vm_identity', validator=process_remove_identity_namespace, min_api='2017-12-01')
+        g.custom_show_command('show', 'show_vm_identity')
 
+    with self.command_group('vm', compute_vm_sdk) as g:
         g.custom_command('application set', 'set_vm_applications', validator=process_set_applications_namespace, min_api='2021-07-01')
         g.custom_command('application list', 'list_vm_applications', min_api='2021-07-01')