diff --git a/src/horizondb/azext_horizondb/_help.py b/src/horizondb/azext_horizondb/_help.py index 8b75c209e64..64608e783db 100644 --- a/src/horizondb/azext_horizondb/_help.py +++ b/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 +""" diff --git a/src/horizondb/azext_horizondb/_params.py b/src/horizondb/azext_horizondb/_params.py index fedd6097eb6..00cc995ed69 100644 --- a/src/horizondb/azext_horizondb/_params.py +++ b/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() diff --git a/src/horizondb/azext_horizondb/cluster_commands.py b/src/horizondb/azext_horizondb/cluster_commands.py index 9df2c7c1322..ed7fa4385cb 100644 --- a/src/horizondb/azext_horizondb/cluster_commands.py +++ b/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') diff --git a/src/horizondb/azext_horizondb/commands/identity_commands.py b/src/horizondb/azext_horizondb/commands/identity_commands.py new file mode 100644 index 00000000000..c479a352bd4 --- /dev/null +++ b/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]}