[HORIZONDB] az horizondb identity assign/list/remove/show: Introduced identity commands for HorizonDB

src/horizondb/azext_horizondb/_help.py

@@ -53,3 +53,45 @@
   - name: List Azure HorizonDB clusters in a resource group.
     text: az horizondb list --resource-group exampleresourcegroup
 """
+
+
+helps['horizondb identity'] = """
+type: group
+short-summary: Manage user assigned managed identities for an Azure HorizonDB cluster.
+"""
+
+
+helps['horizondb identity assign'] = """
+type: command
+short-summary: Add user assigned managed identities to an Azure HorizonDB cluster.
+examples:
+  - name: Assign a user assigned managed identity to a HorizonDB cluster.
+    text: az horizondb identity assign --name examplecluster --resource-group exampleresourcegroup --identity /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/exampleresourcegroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/exampleidentity
+"""
+
+
+helps['horizondb identity list'] = """
+type: command
+short-summary: List all user assigned managed identities from an Azure HorizonDB cluster.
+examples:
+  - name: List identities for a HorizonDB cluster.
+    text: az horizondb identity list --name examplecluster --resource-group exampleresourcegroup
+"""
+
+
+helps['horizondb identity remove'] = """
+type: command
+short-summary: Remove user assigned managed identities from an Azure HorizonDB cluster.
+examples:
+  - name: Remove a user assigned managed identity from a HorizonDB cluster.
+    text: az horizondb identity remove --name examplecluster --resource-group exampleresourcegroup --identity /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/exampleresourcegroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/exampleidentity
+"""
+
+
+helps['horizondb identity show'] = """
+type: command
+short-summary: Get a user assigned managed identity from an Azure HorizonDB cluster.
+examples:
+  - name: Show a specific identity from a HorizonDB cluster.
+    text: az horizondb identity show --name examplecluster --resource-group exampleresourcegroup --identity /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/exampleresourcegroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/exampleidentity
+"""

src/horizondb/azext_horizondb/_params.py

@@ -81,4 +81,26 @@ def _horizondb_params():
         with self.argument_context('horizondb delete') as c:
             c.argument('yes', arg_type=yes_arg_type)
 
+        identity_arg_type = CLIArgumentType(
+            options_list=['--identity'],
+            nargs='+',
+            help='Space-separated resource IDs of user assigned managed identities.')
+
+        identity_single_arg_type = CLIArgumentType(
+            options_list=['--identity'],
+            help='Resource ID of the user assigned managed identity.')
+
+        with self.argument_context('horizondb identity assign') as c:
+            c.argument('identity', arg_type=identity_arg_type)
+
+        with self.argument_context('horizondb identity list') as c:
+            pass
+
+        with self.argument_context('horizondb identity remove') as c:
+            c.argument('identity', arg_type=identity_arg_type)
+            c.argument('yes', arg_type=yes_arg_type)
+
+        with self.argument_context('horizondb identity show') as c:
+            c.argument('identity', arg_type=identity_single_arg_type)
+
     _horizondb_params()

src/horizondb/azext_horizondb/cluster_commands.py

@@ -26,3 +26,13 @@ def load_command_table(self, _):
         g.custom_command('delete', 'horizondb_cluster_delete')
         g.custom_command('list', 'horizondb_cluster_list')
         g.show_command('show', 'get')
+
+    identity_commands = CliCommandType(
+        operations_tmpl='azext_horizondb.commands.identity_commands#{}')
+    with self.command_group('horizondb identity', horizondb_clusters_sdk,
+                            custom_command_type=identity_commands,
+                            client_factory=cf_horizondb_clusters) as g:
+        g.custom_command('assign', 'horizondb_identity_assign', supports_no_wait=True)
+        g.custom_command('list', 'horizondb_identity_list')
+        g.custom_command('remove', 'horizondb_identity_remove', supports_no_wait=True)
+        g.custom_command('show', 'horizondb_identity_show')

src/horizondb/azext_horizondb/commands/identity_commands.py

@@ -0,0 +1,85 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+# pylint: disable=line-too-long
+
+from azure.cli.core.util import sdk_no_wait, user_confirmation
+
+
+def _get_identity_dict(cluster):
+    """Extract the identity dict from a cluster resource."""
+    try:
+        return cluster["identity"] or {}
+    except KeyError:
+        return {}
+
+
+def _get_user_assigned_identities(cluster):
+    """Extract userAssignedIdentities from cluster identity."""
+    identity = _get_identity_dict(cluster)
+    return identity.get("userAssignedIdentities") or {}
+
+
+def horizondb_identity_assign(client, resource_group_name, cluster_name, identity, no_wait=False):
+    cluster = client.get(resource_group_name=resource_group_name, cluster_name=cluster_name)
+    current_identities = _get_user_assigned_identities(cluster)
+
+    for identity_id in identity:
+        current_identities[identity_id] = {}
+
+    patch_body = {
+        "identity": {
+            "type": "UserAssigned",
+            "userAssignedIdentities": current_identities
+        }
+    }
+
+    return sdk_no_wait(no_wait, client.begin_update,
+                       resource_group_name=resource_group_name,
+                       cluster_name=cluster_name,
+                       properties=patch_body)
+
+
+def horizondb_identity_list(client, resource_group_name, cluster_name):
+    cluster = client.get(resource_group_name=resource_group_name, cluster_name=cluster_name)
+    return _get_user_assigned_identities(cluster)
+
+
+def horizondb_identity_remove(cmd, client, resource_group_name, cluster_name, identity, no_wait=False, yes=False):
+    if not yes:
+        user_confirmation(
+            "Are you sure you want to remove the specified identities from cluster '{0}' "
+            "in resource group '{1}'".format(cluster_name, resource_group_name), yes=yes)
+
+    cluster = client.get(resource_group_name=resource_group_name, cluster_name=cluster_name)
+    current_identities = _get_user_assigned_identities(cluster)
+
+    identities_to_remove = {identity_id: None for identity_id in identity}
+
+    remaining = {k: v for k, v in current_identities.items() if k not in identities_to_remove}
+
+    patch_body = {
+        "identity": {
+            "type": "UserAssigned" if remaining else "None",
+            "userAssignedIdentities": identities_to_remove
+        }
+    }
+
+    return sdk_no_wait(no_wait, client.begin_update,
+                       resource_group_name=resource_group_name,
+                       cluster_name=cluster_name,
+                       properties=patch_body)
+
+
+def horizondb_identity_show(client, resource_group_name, cluster_name, identity):
+    cluster = client.get(resource_group_name=resource_group_name, cluster_name=cluster_name)
+    user_assigned = _get_user_assigned_identities(cluster)
+
+    if identity not in user_assigned:
+        from azure.cli.core.azclierror import ResourceNotFoundError
+        raise ResourceNotFoundError(
+            "Identity '{}' not found on cluster '{}'.".format(identity, cluster_name))
+
+    return {identity: user_assigned[identity]}