Add MSI support for enterprise tier

src/spring-cloud/azext_spring_cloud/_enterprise.py

@@ -44,7 +44,7 @@ def app_list_enterprise(cmd, client, resource_group, service):
     return apps
 
 def app_create_enterprise(cmd, client, resource_group, service, name, 
-                          assign_endpoint, cpu, memory, instance_count, jvm_options, env):
+                          assign_endpoint, cpu, memory, instance_count, jvm_options, env, assign_identity):
     '''app_create_enterprise
     Create app with an active deployment, deployment should be deployed with default banner
     1. Create app
@@ -54,9 +54,9 @@ def app_create_enterprise(cmd, client, resource_group, service, name,
     _ensure_app_not_exist(client, resource_group, service, name)
     need_update_app_after_deployment = assign_endpoint
     total_steps = 3 if need_update_app_after_deployment else 2
-    
+
     logger.warning("[1/{}] Creating app {}".format(total_steps, name))
-    app_poller = _create_empty_app(client, resource_group, service, name)
+    app_poller = _create_app(client, resource_group, service, name, assign_identity)
     _wait_till_end(cmd, app_poller)
 
     logger.warning('[2/{}] Create default deployment with name {} {}'.format(total_steps, DEFAULT_DEPLOYMENT_NAME, 
@@ -431,12 +431,14 @@ def _ensure_app_not_exist(client, resource_group, service, name):
         raise CLIError('App {} already exist.'.format(app.id))
 
 
-def _create_empty_app(client, resource_group, service, name):
+def _create_app(client, resource_group, service, name, assign_identity):
     resource = models.AppResource(
                     properties=models.AppResourceProperties(
                         temporary_disk=models.TemporaryDisk(size_in_gb=5, mount_path='/tmp')
                     )
                 )
+    if assign_identity is True:
+        resource.identity = models.ManagedIdentityProperties(type="systemassigned")
     return client.apps.begin_create_or_update(resource_group, service, name, resource)
 
 

src/spring-cloud/azext_spring_cloud/commands.py

@@ -100,7 +100,7 @@ def load_command_table(self, _):
                          supports_no_wait=True)
         g.custom_command('logs', 'app_tail_log')
 
-    with self.command_group('spring-cloud app identity', client_factory=cf_spring_cloud,
+    with self.command_group('spring-cloud app identity', custom_command_type=app_routing_util,
                             exception_handler=handle_asc_exception) as g:
         g.custom_command('assign', 'app_identity_assign')
         g.custom_command('remove', 'app_identity_remove')

src/spring-cloud/azext_spring_cloud/custom.py

@@ -604,11 +604,11 @@ def app_tail_log(cmd, client, resource_group, service, name,
         raise exceptions[0]
 
 
-def app_identity_assign(cmd, client, resource_group, service, name, role=None, scope=None):
+def app_identity_assign(cmd, client, models, resource_group, service, name, role=None, scope=None):
     _check_active_deployment_exist(client, resource_group, service, name)
-    app_resource = models_20210601preview.AppResource()
-    identity = models_20210601preview.ManagedIdentityProperties(type="systemassigned")
-    properties = models_20210601preview.AppResourceProperties()
+    app_resource = models.AppResource()
+    identity = models.ManagedIdentityProperties(type="systemassigned")
+    properties = models.AppResourceProperties()
     resource = client.services.get(resource_group, service)
     location = resource.location
 
@@ -649,10 +649,10 @@ def app_identity_assign(cmd, client, resource_group, service, name, role=None, s
     return app
 
 
-def app_identity_remove(cmd, client, resource_group, service, name):
-    app_resource = models_20210601preview.AppResource()
-    identity = models_20210601preview.ManagedIdentityProperties(type="none")
-    properties = models_20210601preview.AppResourceProperties()
+def app_identity_remove(cmd, client, model, resource_group, service, name):
+    app_resource = models.AppResource()
+    identity = models.ManagedIdentityProperties(type="none")
+    properties = models.AppResourceProperties()
     resource = client.services.get(resource_group, service)
     location = resource.location
 

src/spring-cloud/azext_spring_cloud/tier_routing_app.py

@@ -12,8 +12,13 @@
                      app_create as app_create_standard, app_update as app_update_standard, 
                      app_scale as app_scale_standard, app_deploy as app_deploy_standard,
                      app_delete as app_delete_standard, app_restart as app_restart_standard,
-                     app_start as app_start_standard, app_stop as app_stop_standard)
+                     app_start as app_start_standard, app_stop as app_stop_standard,
+                     app_identity_assign as identity_assign,
+                     app_identity_remove as identity_remove,
+                     app_identity_show as identity_show)
 from knack.log import get_logger
+from .vendored_sdks.appplatform.v2022_05_01_preview import models as models_20220501preview
+from .vendored_sdks.appplatform.v2021_06_01_preview import models as models_20210601preview
 
 logger = get_logger(__name__)
 
@@ -48,7 +53,7 @@ def app_create(cmd, client, resource_group, service, name,
         # runtime_version, enable_persistent_storage assign_ideneity not support
         return app_create_enterprise(cmd, get_client(cmd), resource_group, service, name,
                                      assign_endpoint, cpu, memory, instance_count, jvm_options, 
-                                     env)
+                                     env, assign_identity)
     else:
         return app_create_standard(cmd, client, resource_group, service, name,
                                    assign_endpoint, cpu, memory, instance_count, runtime_version, 
@@ -155,4 +160,25 @@ def app_deploy(cmd, client, resource_group, service, name,
         # config_file_patterns not support
         return app_deploy_standard(cmd, client, resource_group, service, name,
                                    version, deployment, artifact_path, target_module, runtime_version, 
-                                   jvm_options, main_entry, env, no_wait)
+                                   jvm_options, main_entry, env, no_wait)
+
+
+def app_identity_assign(cmd, client, resource_group, service, name, role=None, scope=None):
+    models = models_20210601preview
+    if is_enterprise_tier(cmd, resource_group, service):
+        client = get_client(cmd)
+        models = models_20220501preview
+    return identity_assign(cmd, client, models, resource_group, service, name, role, scope)
+
+
+def app_identity_remove(cmd, client, resource_group, service, name):
+    models = models_20210601preview
+    if is_enterprise_tier(cmd, resource_group, service):
+        client = get_client(cmd)
+        models = models_20220501preview
+    return identity_remove(cmd, client, models, resource_group, service, name)
+
+
+def app_identity_show(cmd, client, resource_group, service, name):
+    client = get_client(cmd) if is_enterprise_tier(cmd, resource_group, service) else client
+    return identity_show(cmd, client, resource_group, service, name)