{AKS} AIManager print manual grant commands when --no-wait skips role assignment

src/aimanager/HISTORY.rst

@@ -3,6 +3,12 @@
 Release History
 ===============
 
+1.5.1
+++++++
+* ``az aimanager create`` and ``az aimanager namespace add``: When the caller role grant is
+  skipped under ``--no-wait``, the warning now prints the exact ``az role assignment create``
+  commands to grant the roles manually, instead of suggesting a re-run.
+
 1.5.0
 ++++++
 * ``az aimanager create`` and ``az aimanager namespace add``: On success, grant the caller the

src/aimanager/azext_aimanager/_roleassignments.py

@@ -105,10 +105,36 @@ def _assign_role(assignments_client, params_model, subscription_id, scope, role_
     return None
 
 
+def _role_name(role_id):
+    # Fall back to the raw GUID (which `az role assignment create --role` also accepts) so the
+    # best-effort warning path never raises on a role id that is not in AIMANAGER_ROLE_NAMES.
+    return AIMANAGER_ROLE_NAMES.get(role_id, role_id)
+
+
+def _role_assignment_command(assignee, role_id, scope):
+    """Build the exact 'az role assignment create' command that grants one role to the caller."""
+    return (f'az role assignment create --assignee-object-id {assignee} '
+            f'--role "{_role_name(role_id)}" --scope {scope}')
+
+
 def _warn_assignment_failed(scope, role_id, object_id):
-    role_name = AIMANAGER_ROLE_NAMES.get(role_id, role_id)
     logger.warning(
         "Could not assign '%s' to the caller on %s. This is expected if you are not an Owner or "
-        "User Access Administrator. An administrator can grant it with:\n"
-        "  az role assignment create --assignee-object-id %s --role \"%s\" --scope %s",
-        role_name, scope, object_id, role_name, scope)
+        "User Access Administrator. An administrator can grant it with:\n  %s",
+        _role_name(role_id), scope,
+        _role_assignment_command(object_id, role_id, scope))
+
+
+def warn_roles_skipped_no_wait(cmd, scope, role_definition_ids):
+    """--no-wait skipped the automatic grant (success cannot be confirmed before the command
+    returns). Print the exact 'az role assignment create' commands so the caller can grant the
+    roles themselves."""
+    object_id, _ = _get_caller_identity(cmd.cli_ctx)
+    assignee = object_id or "<caller-object-id>"
+    commands = "\n".join(
+        "  " + _role_assignment_command(assignee, role_id, scope)
+        for role_id in role_definition_ids)
+    logger.warning(
+        "--no-wait was set, so the caller's role assignments on %s were skipped. Grant them once "
+        "the create succeeds (requires Owner or User Access Administrator; data-plane access can "
+        "take a few minutes to take effect):\n%s", scope, commands)

src/aimanager/azext_aimanager/custom.py

@@ -19,7 +19,7 @@
     parse_key_value_list,
     print_or_merge_credentials,
 )
-from azext_aimanager._roleassignments import assign_caller_roles
+from azext_aimanager._roleassignments import assign_caller_roles, warn_roles_skipped_no_wait
 from azext_aimanager.constants import AIMANAGER_CALLER_ROLE_IDS
 
 logger = get_logger(__name__)
@@ -54,9 +54,7 @@ def _grant_caller_roles_on_success(cmd, poller, no_wait, scope):
     the command returns before the operation completes and success cannot be confirmed.
     """
     if no_wait:
-        logger.warning(
-            "--no-wait was set, so the caller's role assignments on %s were skipped. Re-run "
-            "without --no-wait, or assign the roles manually.", scope)
+        warn_roles_skipped_no_wait(cmd, scope, AIMANAGER_CALLER_ROLE_IDS)
         return poller
     result = LongRunningOperation(cmd.cli_ctx)(poller)  # blocks until Succeeded; raises on failure
     try:

src/aimanager/azext_aimanager/tests/latest/test_aimanager.py

@@ -41,18 +41,21 @@ def test_create_assigns_roles_on_aimanager_scope(self, _construct, mock_assign,
         self.assertEqual(scope, AIMANAGER_SCOPE)
         self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
 
-    @patch.object(custom, "logger")
+    @patch.object(custom, "warn_roles_skipped_no_wait")
     @patch.object(custom, "LongRunningOperation")
     @patch(SUB_PATCH, return_value="sub")
     @patch.object(custom, "assign_caller_roles")
     @patch.object(custom, "_construct_aimanager", return_value=object())
-    def test_create_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_logger):
+    def test_create_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_warn):
         custom.create_aimanager(
             self.cmd, self.client, "rg", "aim", location="eastus2", no_wait=True)
 
         mock_assign.assert_not_called()
         mock_lro.assert_not_called()
-        mock_logger.warning.assert_called_once()  # warns that the grant was skipped under --no-wait
+        mock_warn.assert_called_once()  # prints the manual-grant remediation under --no-wait
+        _cmd, scope, roles = mock_warn.call_args.args
+        self.assertEqual(scope, AIMANAGER_SCOPE)
+        self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
 
     @patch.object(custom, "LongRunningOperation")
     @patch(SUB_PATCH, return_value="sub")
@@ -69,18 +72,21 @@ def test_namespace_add_assigns_roles_on_namespace_scope(self, _construct, mock_a
         self.assertEqual(scope, NAMESPACE_SCOPE)
         self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
 
-    @patch.object(custom, "logger")
+    @patch.object(custom, "warn_roles_skipped_no_wait")
     @patch.object(custom, "LongRunningOperation")
     @patch(SUB_PATCH, return_value="sub")
     @patch.object(custom, "assign_caller_roles")
     @patch.object(custom, "_construct_namespace", return_value=object())
-    def test_namespace_add_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_logger):
+    def test_namespace_add_skips_roles_with_no_wait(self, _construct, mock_assign, _sub, mock_lro, mock_warn):
         custom.add_aimanager_namespace(
             self.cmd, self.client, "rg", "aim", "team-alpha", no_wait=True)
 
         mock_assign.assert_not_called()
         mock_lro.assert_not_called()
-        mock_logger.warning.assert_called_once()  # warns that the grant was skipped under --no-wait
+        mock_warn.assert_called_once()  # prints the manual-grant remediation under --no-wait
+        _cmd, scope, roles = mock_warn.call_args.args
+        self.assertEqual(scope, NAMESPACE_SCOPE)
+        self.assertEqual(roles, AIMANAGER_CALLER_ROLE_IDS)
 
     @patch.object(custom, "LongRunningOperation")
     @patch(SUB_PATCH, return_value="sub")

src/aimanager/azext_aimanager/tests/latest/test_roleassignments.py

@@ -180,6 +180,34 @@ def test_get_caller_identity_malformed_token_returns_none(self, mock_profile):
         self.assertIsNone(object_id)
         self.assertIsNone(principal_type)
 
+    @patch.object(ra, "logger")
+    @patch.object(ra, "_get_caller_identity", return_value=("oid-1", "User"))
+    def test_warn_roles_skipped_no_wait_prints_exact_commands(self, _ident, mock_logger):
+        ra.warn_roles_skipped_no_wait(self.cmd, SCOPE, [ROLE_A, ROLE_B])
+
+        mock_logger.warning.assert_called_once()
+        fmt, *args = mock_logger.warning.call_args.args
+        msg = fmt % tuple(args)
+        # One runnable command per role, targeting the caller's object id and the scope.
+        self.assertEqual(msg.count("az role assignment create --assignee-object-id oid-1"), 2)
+        self.assertIn(f"--scope {SCOPE}", msg)
+        self.assertIn(f'--role "{ra.AIMANAGER_ROLE_NAMES[ROLE_A]}"', msg)
+        self.assertIn(f'--role "{ra.AIMANAGER_ROLE_NAMES[ROLE_B]}"', msg)
+        # The old, broken remediation must not reappear.
+        self.assertNotIn("Re-run without --no-wait", msg)
+        # Sets expectations that the manual grant needs elevated permissions.
+        self.assertIn("Owner or User Access Administrator", msg)
+
+    @patch.object(ra, "logger")
+    @patch.object(ra, "_get_caller_identity", return_value=(None, None))
+    def test_warn_roles_skipped_no_wait_falls_back_when_oid_unknown(self, _ident, mock_logger):
+        ra.warn_roles_skipped_no_wait(self.cmd, SCOPE, [ROLE_A])
+
+        mock_logger.warning.assert_called_once()
+        fmt, *args = mock_logger.warning.call_args.args
+        msg = fmt % tuple(args)
+        self.assertIn("--assignee-object-id <caller-object-id>", msg)
+
 
 if __name__ == '__main__':
     unittest.main()

src/aimanager/setup.py

@@ -14,7 +14,7 @@
     from distutils import log as logger
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
-VERSION = '1.5.0'
+VERSION = '1.5.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers