diff --git a/azure-cli-extensions.pyproj b/azure-cli-extensions.pyproj
index c877de0a89c..2dbd7a29457 100644
--- a/azure-cli-extensions.pyproj
+++ b/azure-cli-extensions.pyproj
@@ -743,6 +743,103 @@
+
+
+
+
+ Code
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
@@ -4668,6 +4765,15 @@
+
+
+
+
+
+
+
+
+
@@ -4791,6 +4897,17 @@
+
+
+
+
+
+
+
+
+
+
+
diff --git a/src/cosmosdb-preview/HISTORY.rst b/src/cosmosdb-preview/HISTORY.rst
index c97fa6774f4..b96d6379c3a 100644
--- a/src/cosmosdb-preview/HISTORY.rst
+++ b/src/cosmosdb-preview/HISTORY.rst
@@ -2,6 +2,10 @@
Release History
===============
+0.11.0
+++++++
+* Adding support for commands to manage identity.
+
0.10.0
++++++
* Adding support for Services APIs and Graph Resources.
diff --git a/src/cosmosdb-preview/README.md b/src/cosmosdb-preview/README.md
index f81aa322012..5aff6b61a1f 100644
--- a/src/cosmosdb-preview/README.md
+++ b/src/cosmosdb-preview/README.md
@@ -7,6 +7,7 @@ This package provides commands to
- List the different versions of databases and collections that were modified
- Trigger a point in time restore on the Azure CosmosDB continuous mode backup accounts
- Update the backup interval and backup retention of periodic mode backup accounts
+- Update the identities associated with a Cosmos DB account
## How to use ##
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py
index b656d4d79a6..dd291b33671 100644
--- a/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py
+++ b/src/cosmosdb-preview/azext_cosmosdb_preview/_help.py
@@ -213,3 +213,37 @@
type: command
short-summary: Return if the given cosmosdb graph resource exist.
"""
+
+helps['cosmosdb identity'] = """
+type: group
+short-summary: Manage Azure Cosmos DB managed service identities.
+"""
+
+helps['cosmosdb identity show'] = """
+type: command
+short-summary: Show the identities for a Azure Cosmos DB database account.
+examples:
+ - name: Show the identities for a Azure Cosmos DB database account.
+ text: az cosmosdb identity show --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup
+"""
+
+helps['cosmosdb identity assign'] = """
+type: command
+short-summary: Assign SystemAssigned identity for a Azure Cosmos DB database account.
+examples:
+ - name: Assign SystemAssigned identity for a Azure Cosmos DB database account.
+ text: az cosmosdb identity assign --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup
+ - name: Assign one UserAssigned identity for a Azure Cosmos DB database account.
+ text: az cosmosdb identity assign --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --identities /subscriptions/00000000-0000-0000-0000-00000000/resourcegroups/MyRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyID
+"""
+
+helps['cosmosdb identity remove'] = """
+type: command
+short-summary: Remove SystemAssigned identity for a Azure Cosmos DB database account.
+examples:
+ - name: Remove SystemAssigned identity for a Azure Cosmos DB database account.
+ text: az cosmosdb identity remove --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup
+ - name: Remove a UserAssigned identity for a Azure Cosmos DB database account.
+ text: az cosmosdb identity remove --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --identities /subscriptions/00000000-0000-0000-0000-00000000/resourcegroups/MyRG/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyID
+
+"""
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py
index fd031571481..ee9c0c08995 100644
--- a/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py
+++ b/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py
@@ -93,3 +93,10 @@ def load_arguments(self, _):
c.argument('service_name', options_list=['--name', '-n'], help="Service Name.")
c.argument('instance_count', options_list=['--count', '-c'], help="Instance Count.")
c.argument('instance_size', options_list=['--size'], help="Instance Size. Possible values are: Cosmos.D4s, Cosmos.D8s, Cosmos.D16s etc")
+
+ # Identity
+ with self.argument_context('cosmosdb identity assign') as c:
+ c.argument('identities', options_list=['--identities'], nargs='*', help="Space-separated identities to assign. Use '[system]' to refer to the system assigned identity. Default: '[system]'")
+
+ with self.argument_context('cosmosdb identity remove') as c:
+ c.argument('identities', options_list=['--identities'], nargs='*', help="Space-separated identities to remove. Use '[system]' to refer to the system assigned identity. Default: '[system]'")
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py
index 5aef8206d16..f540421b3e4 100644
--- a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py
+++ b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py
@@ -62,3 +62,8 @@ def load_command_table(self, _):
g.command('list', 'list')
g.show_command('show', 'get')
g.command('delete', 'begin_delete', confirmation=True, supports_no_wait=True)
+
+ with self.command_group('cosmosdb identity', client_factory=cf_service, is_preview=True) as g:
+ g.custom_show_command('show', 'cli_cosmosdb_identity_show')
+ g.custom_command('assign', 'cli_cosmosdb_identity_assign')
+ g.custom_command('remove', 'cli_cosmosdb_identity_remove')
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py b/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py
index d5f673934f5..30aaeb114b8 100644
--- a/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py
+++ b/src/cosmosdb-preview/azext_cosmosdb_preview/custom.py
@@ -216,6 +216,127 @@ def cli_cosmosdb_managed_cassandra_datacenter_update(client, resource_group_name
)
return client.begin_create_update(resource_group_name, cluster_name, data_center_name, data_center_resource)
+
+
+def cli_cosmosdb_identity_show(client, resource_group_name, account_name):
+ """ Show the identity associated with a Cosmos DB account """
+
+ cosmos_db_account = client.get(resource_group_name, account_name)
+ return cosmos_db_account.identity
+
+
+def cli_cosmosdb_identity_assign(client,
+ resource_group_name,
+ account_name,
+ identities=None):
+ """ Update the identities associated with a Cosmos DB account """
+
+ existing = client.get(resource_group_name, account_name)
+
+ SYSTEM_ID = '[system]'
+ enable_system = identities is None or SYSTEM_ID in identities
+ new_user_identities = []
+ if identities is not None:
+ new_user_identities = [x for x in identities if x != SYSTEM_ID]
+
+ only_enabling_system = enable_system and len(new_user_identities) == 0
+ system_already_added = existing.identity.type == ResourceIdentityType.system_assigned or existing.identity.type == ResourceIdentityType.system_assigned_user_assigned
+ all_new_users_already_added = new_user_identities and existing.identity and existing.identity.user_assigned_identities and all(x in existing.identity.user_assigned_identities for x in new_user_identities)
+ if only_enabling_system and system_already_added:
+ return existing.identity
+ if (not enable_system) and all_new_users_already_added:
+ return existing.identity
+ if enable_system and system_already_added and all_new_users_already_added:
+ return existing.identity
+
+ if existing.identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned:
+ identity_type = ResourceIdentityType.system_assigned_user_assigned
+ elif existing.identity and existing.identity.type == ResourceIdentityType.system_assigned and new_user_identities:
+ identity_type = ResourceIdentityType.system_assigned_user_assigned
+ elif existing.identity and existing.identity.type == ResourceIdentityType.user_assigned and enable_system:
+ identity_type = ResourceIdentityType.system_assigned_user_assigned
+ elif new_user_identities and enable_system:
+ identity_type = ResourceIdentityType.system_assigned_user_assigned
+ elif new_user_identities:
+ identity_type = ResourceIdentityType.user_assigned
+ else:
+ identity_type = ResourceIdentityType.system_assigned
+
+ if identity_type in [ResourceIdentityType.system_assigned, ResourceIdentityType.none]:
+ new_identity = ManagedServiceIdentity(type=identity_type.value)
+ else:
+ new_assigned_identities = existing.identity.user_assigned_identities or {}
+ for identity in new_user_identities:
+ new_assigned_identities[identity] = Components1Jq1T4ISchemasManagedserviceidentityPropertiesUserassignedidentitiesAdditionalproperties()
+
+ new_identity = ManagedServiceIdentity(type=identity_type.value, user_assigned_identities=new_assigned_identities)
+
+ params = DatabaseAccountUpdateParameters(identity=new_identity)
+ async_cosmos_db_update = client.begin_update(resource_group_name, account_name, params)
+ cosmos_db_account = async_cosmos_db_update.result()
+ return cosmos_db_account.identity
+
+
+def cli_cosmosdb_identity_remove(client,
+ resource_group_name,
+ account_name,
+ identities=None):
+ """ Remove the identities associated with a Cosmos DB account """
+
+ existing = client.get(resource_group_name, account_name)
+
+ SYSTEM_ID = '[system]'
+ remove_system_assigned_identity = False
+ if not identities:
+ remove_system_assigned_identity = True
+ elif SYSTEM_ID in identities:
+ remove_system_assigned_identity = True
+ identities.remove(SYSTEM_ID)
+
+ if existing.identity is None:
+ return ManagedServiceIdentity(type=ResourceIdentityType.none.value)
+ if existing.identity.user_assigned_identities:
+ existing_identities = existing.identity.user_assigned_identities.keys()
+ else:
+ existing_identities = []
+ if identities:
+ identities_to_remove = identities
+ else:
+ identities_to_remove = []
+ non_existing = [x for x in identities_to_remove if x not in set(existing_identities)]
+
+ if non_existing:
+ raise CLIError("'{}' are not associated with '{}'".format(','.join(non_existing), account_name))
+ identities_remaining = [x for x in existing_identities if x not in set(identities_to_remove)]
+ if remove_system_assigned_identity and ((not existing.identity) or (existing.identity and existing.identity.type in [ResourceIdentityType.none, ResourceIdentityType.user_assigned])):
+ raise CLIError("System-assigned identity is not associated with '{}'".format(account_name))
+
+ if identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned:
+ set_type = ResourceIdentityType.system_assigned_user_assigned
+ elif identities_remaining and remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned:
+ set_type = ResourceIdentityType.user_assigned
+ elif identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.user_assigned:
+ set_type = ResourceIdentityType.user_assigned
+ elif not identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned_user_assigned:
+ set_type = ResourceIdentityType.system_assigned
+ elif not identities_remaining and not remove_system_assigned_identity and existing.identity.type == ResourceIdentityType.system_assigned:
+ set_type = ResourceIdentityType.system_assigned
+ else:
+ set_type = ResourceIdentityType.none
+
+ new_user_identities = {}
+ for identity in identities_remaining:
+ new_user_identities[identity] = Components1Jq1T4ISchemasManagedserviceidentityPropertiesUserassignedidentitiesAdditionalproperties()
+ if set_type in [ResourceIdentityType.system_assigned_user_assigned, ResourceIdentityType.user_assigned]:
+ for removed_identity in identities_to_remove:
+ new_user_identities[removed_identity] = None
+ if not new_user_identities:
+ new_user_identities = None
+
+ params = DatabaseAccountUpdateParameters(identity=ManagedServiceIdentity(type=set_type, user_assigned_identities=new_user_identities))
+ async_cosmos_db_update = client.begin_update(resource_group_name, account_name, params)
+ cosmos_db_account = async_cosmos_db_update.result()
+ return cosmos_db_account.identity
def _handle_exists_exception(http_response_error):
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py b/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py
new file mode 100644
index 00000000000..dd7f01cf662
--- /dev/null
+++ b/src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/test_cosmosdb-identity_scenario.py
@@ -0,0 +1,123 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+from unittest import mock
+
+
+TEST_DIR = os.path.abspath(os.path.join(os.path.abspath(__file__), '..'))
+
+class Cosmosdb_previewIdentityTest(ScenarioTest):
+
+ @ResourceGroupPreparer(name_prefix='cli_test_cosmosdb_managed_service_identity')
+ @KeyVaultPreparer(name_prefix='cli', name_len=15, location='eastus2', additional_params='--enable-purge-protection')
+ def test_cosmosdb_preview_identity(self, resource_group, key_vault):
+ key_name = self.create_random_name(prefix='cli', length=15)
+ key_uri = "https://{}.vault.azure.net/keys/{}".format(key_vault, key_name)
+
+ self.kwargs.update({
+ 'acc': self.create_random_name(prefix='cli', length=15),
+ 'acc2': self.create_random_name(prefix='cli', length=15),
+ 'kv_name': key_vault,
+ 'key_name': key_name,
+ 'key_uri': key_uri,
+ 'location': "eastus2",
+ 'id1': self.create_random_name(prefix='cli', length=15),
+ 'id2': self.create_random_name(prefix='cli', length=15)
+ })
+
+ self.cmd('az keyvault set-policy -n {kv_name} -g {rg} --spn a232010e-820c-4083-83bb-3ace5fc29d0b --key-permissions get unwrapKey wrapKey')
+ self.cmd('az keyvault key create -n {key_name} --kty RSA --size 3072 --vault-name {kv_name}')
+
+ cmk_account = self.cmd('az cosmosdb create -n {acc} -g {rg} --locations regionName={location} failoverPriority=0 --key-uri {key_uri} --assign-identity [system] --default-identity FirstPartyIdentity').get_output_in_json()
+
+ assert cmk_account["keyVaultKeyUri"] == key_uri
+ assert cmk_account["defaultIdentity"] == 'FirstPartyIdentity'
+ assert cmk_account["identity"]['type'] == 'SystemAssigned'
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json()
+ assert identity_output["type"] == "None"
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned"
+
+ identity_principal_id = identity_output["principalId"]
+ self.kwargs.update({
+ 'identity_principal_id': identity_principal_id
+ })
+ self.cmd('az keyvault set-policy -n {kv_name} -g {rg} --object-id {identity_principal_id} --key-permissions get unwrapKey wrapKey')
+
+ # System assigned identity tests
+ cmk_account = self.cmd('az cosmosdb update -n {acc} -g {rg} --default-identity SystemAssignedIdentity').get_output_in_json()
+ assert cmk_account["defaultIdentity"] == 'SystemAssignedIdentity'
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json()
+ assert identity_output["type"] == "None"
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned"
+
+ # User assigned identity tests
+ user_identity1 = self.cmd('az identity create -n {id1} -g {rg}').get_output_in_json()
+ user_identity2 = self.cmd('az identity create -n {id2} -g {rg}').get_output_in_json()
+ id1 = user_identity1["id"]
+ id1principal = user_identity1["principalId"]
+ id2 = user_identity2["id"]
+ self.kwargs.update({
+ 'id1': id1,
+ 'id2': id2,
+ 'id1principal': id1principal
+ })
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned,UserAssigned"
+ assert list(identity_output["userAssignedIdentities"])[0] == id1
+ assert len(identity_output["userAssignedIdentities"]) == 1
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id2}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned,UserAssigned"
+ assert (list(identity_output["userAssignedIdentities"])[0] == id2 or list(identity_output["userAssignedIdentities"])[1] == id2)
+ assert len(identity_output["userAssignedIdentities"]) == 2
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg}').get_output_in_json()
+ assert identity_output["type"] == "UserAssigned"
+ assert len(identity_output["userAssignedIdentities"]) == 2
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id2}').get_output_in_json()
+ assert identity_output["type"] == "UserAssigned"
+ assert list(identity_output["userAssignedIdentities"])[0] == id1
+ assert len(identity_output["userAssignedIdentities"]) == 1
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id2}').get_output_in_json()
+ assert identity_output["type"] == "UserAssigned"
+ assert (list(identity_output["userAssignedIdentities"])[0] == id2 or list(identity_output["userAssignedIdentities"])[1] == id2)
+ assert len(identity_output["userAssignedIdentities"]) == 2
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1} {id2}').get_output_in_json()
+ assert identity_output["type"] == "None"
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned,UserAssigned"
+ assert len(identity_output["userAssignedIdentities"]) == 2
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id2}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned,UserAssigned"
+ assert len(identity_output["userAssignedIdentities"]) == 1
+
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1}').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned"
+
+ identity_output = self.cmd('az cosmosdb identity assign -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json()
+ assert identity_output["type"] == "SystemAssigned,UserAssigned"
+ assert len(identity_output["userAssignedIdentities"]) == 2
+ identity_output = self.cmd('az cosmosdb identity remove -n {acc} -g {rg} --identities {id1} {id2} [system]').get_output_in_json()
+ assert identity_output["type"] == "None"
+
+ # Default identity tests
+ self.cmd('az keyvault set-policy --name {kv_name} --object-id {id1principal} --key-permissions get unwrapKey wrapKey')
+ default_id_acct = self.cmd('az cosmosdb create -n {acc2} -g {rg} --locations regionName={location} failoverPriority=0 --key-uri {key_uri} --assign-identity {id1} --default-identity "UserAssignedIdentity={id1}"').get_output_in_json()
+ assert default_id_acct["identity"]["type"] == "UserAssigned"
+ assert list(default_id_acct["identity"]["userAssignedIdentities"])[0] == id1
+ assert default_id_acct["defaultIdentity"] == "UserAssignedIdentity=" + id1
\ No newline at end of file
diff --git a/src/cosmosdb-preview/setup.py b/src/cosmosdb-preview/setup.py
index 9ffdda9ef03..86aa239c392 100644
--- a/src/cosmosdb-preview/setup.py
+++ b/src/cosmosdb-preview/setup.py
@@ -16,7 +16,7 @@
# TODO: Confirm this is the right version number you want and it matches your
# HISTORY.rst entry.
-VERSION = '0.10.0'
+VERSION = '0.11.0'
# The full list of classifiers is available at
# https://pypi.python.org/pypi?%3Aaction=list_classifiers