[Role] az role assignment create/update: auto complete assignee_principal_type

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

@@ -146,6 +146,19 @@ def create_role_assignment(cmd, role, assignee=None, assignee_object_id=None, re
     if assignee_principal_type and not assignee_object_id:
         raise CLIError('usage error: --assignee-object-id GUID [--assignee-principal-type]')
 
+    # If assignee object id is provided without assignee principal type, try to get principle type from graph
+    if assignee_object_id and not assignee_principal_type:
+        graph_client = _graph_client_factory(cmd.cli_ctx)
+        try:
+            assignee_object_result = _get_object_stubs(graph_client, [assignee_object_id])
+            if assignee_object_result:
+                assignee_principal_type = assignee_object_result[0].object_type
+        except CloudError as ex:
+            logger.warning('Failed to query --assignee-principal-type for %s by invoking Graph API. Err: %s\n'
+                           'RBAC server might reject creating role assignment by --assignee-object-id without '
+                           '--assignee-principal-type in the future. Better to specify --assignee-principal-type '
+                           'manually.', assignee_object_id, str(ex))
+
     # If condition is set and condition-version is empty, condition-version defaults to "2.0".
     if condition and not condition_version:
         condition_version = "2.0"
@@ -281,6 +294,9 @@ def update_role_assignment(cmd, role_assignment):
             original_assignment.condition_version.startswith('2.') and assignment.condition_version.startswith('1.')):
         raise CLIError("Condition version cannot be downgraded to '1.X'.")
 
+    if not assignment.principal_type:
+        assignment.principal_type = original_assignment.principal_type
+
     return assignments_client.create(scope, name, parameters=assignment)
 
 

src/azure-cli/azure/cli/command_modules/role/tests/latest/test_role.py

@@ -357,19 +357,44 @@ def test_role_assignment_create_using_principal_type(self, resource_group):
 
         self.kwargs['rg'] = resource_group
 
+        def _test_role_assignment(assignee_object_id, assignee_principal_type=None):
+            self.kwargs['object_id'] = assignee_object_id
+            self.kwargs['principal_type'] = assignee_principal_type
+            # test role assignment on subscription level
+            if assignee_principal_type:
+                self.cmd(
+                    'role assignment create --assignee-object-id {object_id} --assignee-principal-type {principal_type} --role Reader -g {rg}')
+            else:
+                self.cmd('role assignment create --assignee-object-id {object_id} --role Reader -g {rg}')
+            self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 1))
+            self.cmd('role assignment delete -g {rg}')
+            self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 0))
+
+        def _test_role_assignment_graph_call(assignee_object_id, assignee_principal_type):
+            # test role assignment with principal type which won't trigger graph call
+            _test_role_assignment(assignee_object_id, assignee_principal_type)
+
+            # test role assignment without principal type which will trigger graph call to complete principal type
+            _test_role_assignment(assignee_object_id)
+
+            # test role assignment without principal type and graph call fail
+            from msrestazure.azure_exceptions import CloudError
+            import requests
+            mock_response = requests.Response()
+            mock_response.status_code = 403
+            mock_response.reason = 'Forbidden for url: https://graph.windows.net/.../getObjectsByObjectIds?api-version=1.6'
+            with mock.patch('azure.graphrbac.operations.ObjectsOperations.get_objects_by_object_ids',
+                            side_effect=CloudError(mock_response)):
+                _test_role_assignment(assignee_object_id)
+
         with mock.patch('azure.cli.command_modules.role.custom._gen_guid', side_effect=self.create_guid):
             # User
             user = self.create_random_name('testuser', 15)
             self.kwargs['upn'] = user + '@azuresdkteam.onmicrosoft.com'
 
             result = self.cmd('ad user create --display-name tester123 --password Test123456789 --user-principal-name {upn}').get_output_in_json()
-            self.kwargs['object_id'] = result['objectId']
             try:
-                # test role assignment on subscription level
-                self.cmd('role assignment create --assignee-object-id {object_id} --assignee-principal-type User --role Reader -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 1))
-                self.cmd('role assignment delete -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 0))
+                _test_role_assignment_graph_call(result['objectId'], 'User')
             finally:
                 try:
                     self.cmd('ad user delete --upn-or-object-id {upn}')
@@ -380,14 +405,9 @@ def test_role_assignment_create_using_principal_type(self, resource_group):
             self.kwargs['group_name'] = self.create_random_name('testgroup', 15)
             result = self.cmd(
                 'ad group create --display-name {group_name} --mail-nickname {group_name}').get_output_in_json()
-            self.kwargs['object_id'] = result['objectId']
             time.sleep(10)
             try:
-                # test role assignment on subscription level
-                self.cmd('role assignment create --assignee-object-id {object_id} --assignee-principal-type Group --role Reader -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 1))
-                self.cmd('role assignment delete -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 0))
+                _test_role_assignment_graph_call(result['objectId'], 'Group')
             finally:
                 try:
                     self.cmd('ad group delete --group {object_id}')
@@ -399,13 +419,8 @@ def test_role_assignment_create_using_principal_type(self, resource_group):
             result = self.cmd('ad sp create-for-rbac --skip-assignment --name {sp_name}').get_output_in_json()
             self.kwargs['app_id'] = result['appId']
             result = self.cmd('ad sp show --id {app_id}').get_output_in_json()
-            self.kwargs['object_id'] = result['objectId']
             try:
-                # test role assignment on subscription level
-                self.cmd('role assignment create --assignee-object-id {object_id} --assignee-principal-type ServicePrincipal --role Reader -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 1))
-                self.cmd('role assignment delete -g {rg}')
-                self.cmd('role assignment list -g {rg}', checks=self.check("length([])", 0))
+                _test_role_assignment_graph_call(result['objectId'], 'ServicePrincipal')
             finally:
                 try:
                     self.cmd('ad sp delete --id {object_id}')