add polling

Author: xmuhammad-xali

src/fleet/azext_fleet/_client_factory.py

@@ -5,6 +5,7 @@
 
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.mgmt.msi import ManagedServiceIdentityClient
+from azure.mgmt.authorization import AuthorizationManagementClient
 from azure.cli.core.profiles import (
     CustomResourceType,
     ResourceType
@@ -56,5 +57,9 @@ def get_provider_client(cli_ctx):
         cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES)
 
 
+def get_role_assignments_client(cli_ctx):
+    return get_mgmt_service_client(cli_ctx, AuthorizationManagementClient).role_assignments
+
+
 def get_msi_client(cli_ctx, subscription_id=None):
     return get_mgmt_service_client(cli_ctx, ManagedServiceIdentityClient, subscription_id=subscription_id)

src/fleet/azext_fleet/_helpers.py

@@ -16,10 +16,11 @@
 from knack.util import CLIError
 from azure.cli.command_modules.acs._roleassignments import add_role_assignment
 from azure.mgmt.core.tools import parse_resource_id
-
+from azext_fleet.constants import NETWORK_CONTRIBUTOR_ROLE_ID
 
 from azext_fleet._client_factory import get_provider_client
 from azext_fleet._client_factory import get_msi_client
+from azext_fleet._client_factory import get_role_assignments_client
 
 logger = get_logger(__name__)
 
@@ -159,13 +160,38 @@ def _load_kubernetes_configuration(filename):
 
 def assign_network_contributor_role_to_subnet(cmd, object_id, subnet_id):
     if not add_role_assignment(cmd, 'Network Contributor', object_id, scope=subnet_id):
-        logger.warning("Failed to create Network Contributor role assignment on the subnet %s.\n"
-                       "This role assignment is required for the managed identity to access the subnet.\n"
-                       "Please ensure you have sufficient permissions, or ask an administrator to run:\n"
-                       "az role assignment create --assignee-principal-type ServicePrincipal --assignee-object-id %s "
-                       "--role 'Network Contributor' --scope %s",
-                       subnet_id, object_id, subnet_id)
-    time.sleep(3)
+        logger.warning(
+            "Failed to create Network Contributor role assignment on the subnet %s.\n"
+            "This role assignment is required for the managed identity to access the subnet.\n"
+            "Please ensure you have sufficient permissions, or ask an administrator to run:\n"
+            "az role assignment create --assignee-principal-type ServicePrincipal --assignee-object-id %s "
+            "--role 'Network Contributor' --scope %s",
+            subnet_id, object_id, subnet_id)
+        return
+
+    auth_client = get_role_assignments_client(cmd.cli_ctx)
+    max_attempts = 3
+    interval = 3
+    for _ in range(max_attempts):
+        if _is_assignment_present(auth_client, subnet_id, object_id):
+            logger.warning("Role assignment for Network Contributor on subnet %s was detected.", subnet_id)
+            return
+        time.sleep(interval)
+    logger.warning(
+        "Role assignment for Network Contributor on subnet %s was not detected after %s seconds. "
+        "There may be a delay in propagation.",
+        subnet_id, max_attempts * interval)
+
+
+def _is_assignment_present(auth_client, subnet_id, object_id):
+    filter_query = f"assignedTo('{object_id}') and atScope()"
+    for assignment in auth_client.list_for_scope(subnet_id, filter=filter_query):
+        if assignment.role_definition_id.lower().endswith(NETWORK_CONTRIBUTOR_ROLE_ID) and \
+           assignment.scope.lower() == subnet_id.lower():
+            return True
+        logger.warning("No matching role assignment found for scope %s and subnet %s",
+                       assignment.scope.lower(), subnet_id.lower())
+    return False
 
 
 def get_msi_object_id(cmd, msi_resource_id):

src/fleet/azext_fleet/constants.py

@@ -7,6 +7,7 @@
 UPGRADE_TYPE_FULL = "Full"
 UPGRADE_TYPE_NODEIMAGEONLY = "NodeImageOnly"
 FLEET_1P_APP_ID = "609d2f62-527f-4451-bfd2-ac2c7850822c"
+NETWORK_CONTRIBUTOR_ROLE_ID = "4d97b98b-1d4f-4787-a291-c67834d212e7"
 
 SUPPORTED_GATE_STATES_FILTERS = ["Pending", "Skipped", "Completed"]
 SUPPORTED_GATE_STATES_PATCH = ["Completed"]

src/fleet/azext_fleet/custom.py

@@ -9,6 +9,8 @@
 
 from azure.cli.core.commands.client_factory import get_subscription_id
 from azure.cli.core.util import sdk_no_wait, get_file_json, shell_safe_json_parse
+from azure.cli.command_modules.acs._graph import resolve_object_id
+
 
 from azext_fleet._client_factory import CUSTOM_MGMT_FLEET
 from azext_fleet._helpers import is_rp_registered, print_or_merge_credentials
@@ -119,7 +121,8 @@ def create_fleet(cmd,
         if not is_rp_registered(cmd):
             raise CLIError("The Microsoft.ContainerService resource provider is not registered."
                            "Run `az provider register -n Microsoft.ContainerService --wait`.")
-        assign_network_contributor_role_to_subnet(cmd, FLEET_1P_APP_ID, agent_subnet_id)
+        object_id = resolve_object_id(cmd.cli_ctx, FLEET_1P_APP_ID)
+        assign_network_contributor_role_to_subnet(cmd, object_id, agent_subnet_id)
 
     if enable_vnet_integration and assign_identity is not None:
         object_id = get_msi_object_id(cmd, assign_identity)