[Containerapp] az containerapp sessionpool: Support managed identity when create session pool with --mi-system-assigned --mi-user-assigned (#8223)

Author: najian

src/containerapp/HISTORY.rst

@@ -5,8 +5,9 @@ Release History
 upcoming
 ++++++
 * 'az containerapp create': Fix Role assignment error when the default Azure Container Registry could not be found
-* Upgrade api-version to 2024-10-02-preview
+*  Upgrade api-version to 2024-10-02-preview
 * 'az containerapp create/update': `--yaml` support property pollingInterval and cooldownPeriod
+* 'az containerapp sessionpool create': Support managed identity when create session pool with --mi-system-assigned --mi-user-assigned
 
 1.0.0b4
 ++++++

src/containerapp/azext_containerapp/_clients.py

@@ -1018,7 +1018,7 @@ def list(cls, cmd, resource_group_name, environment_name):
 
 
 class SessionPoolPreviewClient():
-    api_version = SESSION_API_VERSION
+    api_version = PREVIEW_API_VERSION
 
     @classmethod
     def create(cls, cmd, resource_group_name, name, session_pool_envelope, no_wait=False):

src/containerapp/azext_containerapp/_help.py

@@ -1897,6 +1897,20 @@
               --cpu 0.5 --memory 1Gi --target-port 80 --registry-server myregistry.azurecr.io \\
               --registry-username myregistry --registry-password $REGISTRY_PASSWORD \\
               --location eastasia
+    - name: Create or update a Session Pool with container type CustomContainer and Managed Identity to authenticate Azure container registry
+      text: |
+          az containerapp sessionpool create -n mysessionpool -g MyResourceGroup \\
+              --container-type CustomContainer --environment MyEnvironment --image MyImage \\
+              --cpu 0.5 --memory 1Gi --target-port 80 --registry-server myregistry.azurecr.io \\
+              --registry-identity  MyUserIdentityResourceId \\
+              --location eastasia
+    - name: Create or update a Session Pool with container type CustomContainer with system assigned and user assigned identity.
+      text: |
+          az containerapp sessionpool create -n mysessionpool -g MyResourceGroup \\
+              --container-type CustomContainer --environment MyEnvironment --image MyImage \\
+              --cpu 0.5 --memory 1Gi --target-port 80 \\
+              --mi-system-assigned --mi-user-assigned MyUserIdentityResourceId \\
+              --location eastasia
     - name: Create or update a Session Pool with container type CustomContainer with cooldown period 360s
       text: |
           az containerapp sessionpool create -n mysessionpool -g MyResourceGroup \\

src/containerapp/azext_containerapp/_params.py

@@ -398,6 +398,8 @@ def load_arguments(self, _):
     with self.argument_context('containerapp sessionpool') as c:
         c.argument('name', options_list=['--name', '-n'], help="The Session Pool name.")
         c.argument('resource_group_name', arg_type=resource_group_name_type, id_part=None)
+        c.argument('mi_system_assigned', help='Boolean indicating whether to assign system-assigned identity.', action='store_true')
+        c.argument('mi_user_assigned', nargs='+', help='Space-separated user identities to be assigned.')
 
     with self.argument_context('containerapp sessionpool', arg_group='Configuration') as c:
         c.argument('container_type', arg_type=get_enum_type(["CustomContainer", "PythonLTS", "NodeLTS"]), help="The pool type of the Session Pool, default='PythonLTS'")
@@ -424,6 +426,7 @@ def load_arguments(self, _):
         c.argument('registry_server', validator=validate_registry_server, help="The container registry server hostname, e.g. myregistry.azurecr.io.")
         c.argument('registry_pass', validator=validate_registry_pass, options_list=['--registry-password'], help="The password to log in to container registry. If stored as a secret, value must start with \'secretref:\' followed by the secret name.")
         c.argument('registry_user', validator=validate_registry_user, options_list=['--registry-username'], help="The username to log in to container registry.")
+        c.argument('registry_identity', validator=validate_registry_user, help="The managed identity with which to authenticate to the Azure Container Registry (instead of username/password). Use 'system' for a system-assigned identity, use a resource ID for a user-assigned identity. The managed identity should have been assigned acrpull permissions on the ACR before deployment (use 'az role assignment create --role acrpull ...').")
 
     # sessions code interpreter commands
     with self.argument_context('containerapp session code-interpreter') as c:

src/containerapp/azext_containerapp/containerapp_sessionpool_decorator.py

@@ -25,12 +25,13 @@
                                                            store_as_secret_and_return_secret_ref,
                                                            _ensure_location_allowed, CONTAINER_APPS_RP,
                                                            validate_container_app_name,
-                                                           safe_set, safe_get)
+                                                           safe_set, safe_get, _ensure_identity_resource_id)
 from azure.cli.command_modules.containerapp._clients import ManagedEnvironmentClient
 from azure.cli.command_modules.containerapp._client_factory import handle_non_404_status_code_exception
+from azure.cli.command_modules.containerapp._utils import is_registry_msi_system
 from azure.cli.core.commands.client_factory import get_subscription_id
 
-from ._models import SessionPool as SessionPoolModel
+from ._models import ManagedServiceIdentity, SessionPool as SessionPoolModel
 from ._client_factory import handle_raw_exception
 from ._utils import AppType
 
@@ -129,6 +130,15 @@ def get_argument_registry_pass(self):
     def get_argument_registry_user(self):
         return self.get_param("registry_user")
 
+    def get_argument_registry_identity(self):
+        return self.get_param("registry_identity")
+
+    def get_argument_system_assigned(self):
+        return self.get_param("mi_system_assigned")
+
+    def get_argument_user_assigned(self):
+        return self.get_param("mi_user_assigned")
+
     # pylint: disable=no-self-use
     def get_environment_client(self):
         return ManagedEnvironmentClient
@@ -153,6 +163,7 @@ def validate_arguments(self):
 
     def construct_payload(self):
         self.session_pool_def["location"] = self.get_argument_location()
+        self.set_up_managed_identity()
 
         # We only support 'Dynamic' type in CLI
         self.session_pool_def["properties"]["poolManagementType"] = "Dynamic"
@@ -186,6 +197,59 @@ def construct_payload(self):
         safe_set(self.session_pool_def, "properties", "dynamicPoolConfiguration", value=dynamic_pool_def)
         safe_set(self.session_pool_def, "properties", "sessionNetworkConfiguration", value=session_network_def)
         safe_set(self.session_pool_def, "properties", "scaleConfiguration", value=session_scale_def)
+        self.set_up_managed_identity_settings()
+
+    def set_up_managed_identity_settings(self):
+        managed_identity_settings = []
+        if self.get_argument_system_assigned():
+            managed_identity_setting = {
+                "identity": "system",
+                "lifecycle": "Main"
+            }
+            managed_identity_settings.append(managed_identity_setting)
+
+        if self.get_argument_user_assigned():
+            for x in self.get_argument_user_assigned():
+                managed_identity_setting = {
+                    "identity": x.lower(),
+                    "lifecycle": "Main"
+                }
+                managed_identity_settings.append(managed_identity_setting)
+        if managed_identity_settings:
+            safe_set(self.session_pool_def, "properties", "managedIdentitySettings", value=managed_identity_settings)
+
+    def set_up_managed_identity(self):
+        identity_def = deepcopy(ManagedServiceIdentity)
+        identity_def["type"] = "None"
+
+        assign_system_identity = self.get_argument_system_assigned()
+        if self.get_argument_user_assigned():
+            assign_user_identities = [x.lower() for x in self.get_argument_user_assigned()]
+        else:
+            assign_user_identities = []
+
+        identity = self.get_argument_registry_identity()
+        if identity:
+            if is_registry_msi_system(identity):
+                assign_system_identity = True
+            else:
+                assign_user_identities.append(self.get_argument_registry_identity())
+
+        if assign_system_identity and assign_user_identities:
+            identity_def["type"] = "SystemAssigned, UserAssigned"
+        elif assign_system_identity:
+            identity_def["type"] = "SystemAssigned"
+        elif assign_user_identities:
+            identity_def["type"] = "UserAssigned"
+
+        if assign_user_identities:
+            identity_def["userAssignedIdentities"] = {}
+            subscription_id = get_subscription_id(self.cmd.cli_ctx)
+
+            for r in assign_user_identities:
+                r = _ensure_identity_resource_id(subscription_id, self.get_argument_resource_group_name(), r)
+                identity_def["userAssignedIdentities"][r] = {}  # pylint: disable=unsupported-assignment-operation
+        self.session_pool_def["identity"] = identity_def
 
     def set_up_dynamic_configuration(self):
         dynamic_pool_def = {}
@@ -244,14 +308,20 @@ def set_up_registry_auth_configuration(self, secrets_def):
         if self.get_argument_registry_server() is not None:
             registry_def = {}
             registry_def["server"] = self.get_argument_registry_server()
-            registry_def["username"] = self.get_argument_registry_user()
-
-            if secrets_def is None:
-                secrets_def = []
-            registry_def["passwordSecretRef"] = store_as_secret_and_return_secret_ref(secrets_def,
-                                                                                      self.get_argument_registry_user(),
-                                                                                      self.get_argument_registry_server(),
-                                                                                      self.get_argument_registry_pass())
+
+            if self.get_argument_registry_user():
+                registry_def["username"] = self.get_argument_registry_user()
+
+                if secrets_def is None:
+                    secrets_def = []
+                registry_def["passwordSecretRef"] = store_as_secret_and_return_secret_ref(secrets_def,
+                                                                                          self.get_argument_registry_user(),
+                                                                                          self.get_argument_registry_server(),
+                                                                                          self.get_argument_registry_pass())
+
+            if self.get_argument_registry_identity():
+                registry_def["identity"] = self.get_argument_registry_identity()
+
         return registry_def, secrets_def
 
     def set_up_ingress(self):

src/containerapp/azext_containerapp/custom.py

@@ -2864,7 +2864,10 @@ def create_session_pool(cmd,
                         target_port=None,
                         registry_server=None,
                         registry_pass=None,
-                        registry_user=None):
+                        registry_user=None,
+                        mi_user_assigned=None,
+                        registry_identity=None,
+                        mi_system_assigned=False):
     raw_parameters = locals()
     session_pool_decorator = SessionPoolCreateDecorator(
         cmd=cmd,