[Cosmos DB] Support for managed service identity & default identity

src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py

@@ -55,6 +55,7 @@ def load_arguments(self, _):
         c.argument('account_name', completer=None)
         c.argument('key_uri', help="The URI of the key vault", is_preview=True)
         c.argument('enable_free_tier', arg_type=get_three_state_flag(), help="If enabled the account is free-tier.", is_preview=True)
+        c.argument('assign_identity', arg_type=get_three_state_flag(), help="Flag to assign system assigned identity to the account.", is_preview=True)
 
     for scope in ['cosmosdb create', 'cosmosdb update']:
         with self.argument_context(scope) as c:
@@ -79,6 +80,7 @@ def load_arguments(self, _):
             c.argument('backup_interval', type=int, help="the frequency(in minutes) with which backups are taken (only for accounts with periodic mode backups)", arg_group='Backup Policy')
             c.argument('backup_retention', type=int, help="the time(in hours) for which each backup is retained (only for accounts with periodic mode backups)", arg_group='Backup Policy')
             c.argument('server_version', arg_type=get_enum_type(ServerVersion), help="Valid only for MongoDB accounts.", is_preview=True)
+            c.argument('default_identity', help="The primary identity to access key vault in CMK related features.", is_preview=True)
 
     for scope in ['cosmosdb regenerate-key', 'cosmosdb keys regenerate']:
         with self.argument_context(scope) as c:

src/azure-cli/azure/cli/command_modules/cosmosdb/commands.py

@@ -246,6 +246,11 @@ def load_command_table(self, _):
         g.custom_command('update', 'cli_cosmosdb_table_throughput_update')
         g.custom_command('migrate', 'cli_cosmosdb_table_throughput_migrate')
 
+    with self.command_group('cosmosdb identity', None, client_factory=cf_db_accounts) as g:
+        g.custom_command('show', 'cli_cosmosdb_identity_show')
+        g.custom_command('assign', 'cli_cosmosdb_identity_assign')
+        g.custom_command('remove', 'cli_cosmosdb_identity_remove')
+
     # virtual network rules
     with self.command_group('cosmosdb network-rule', None, client_factory=cf_db_accounts) as g:
         g.custom_command('list', 'cli_cosmosdb_network_rule_list')

src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py

@@ -30,6 +30,7 @@
     SqlUserDefinedFunctionCreateUpdateParameters,
     TableResource,
     TableCreateUpdateParameters,
+    ManagedServiceIdentity,
     MongoDBDatabaseResource,
     MongoDBDatabaseCreateUpdateParameters,
     MongoDBCollectionResource,
@@ -99,7 +100,9 @@ def cli_cosmosdb_create(cmd, client,
                         network_acl_bypass=None,
                         network_acl_bypass_resource_ids=None,
                         backup_interval=None,
-                        backup_retention=None):
+                        backup_retention=None,
+                        assign_identity=None,
+                        default_identity=None):
     """Create a new Azure Cosmos DB database account."""
     consistency_policy = None
     if default_consistency_level is not None:
@@ -122,6 +125,13 @@ def cli_cosmosdb_create(cmd, client,
     if enable_public_network is not None:
         public_network_access = 'Enabled' if enable_public_network else 'Disabled'
 
+    system_assigned_identity = None
+    if assign_identity is not None:
+        if assign_identity:
+            system_assigned_identity = ManagedServiceIdentity(type='SystemAssigned')
+        else:
+            system_assigned_identity = ManagedServiceIdentity(type='None')
+
     api_properties = {}
     if kind == DatabaseAccountKind.mongo_db.value:
         api_properties['ServerVersion'] = server_version
@@ -157,7 +167,9 @@ def cli_cosmosdb_create(cmd, client,
         enable_free_tier=enable_free_tier,
         network_acl_bypass=network_acl_bypass,
         network_acl_bypass_resource_ids=network_acl_bypass_resource_ids,
-        backup_policy=backup_policy)
+        backup_policy=backup_policy,
+        identity=system_assigned_identity,
+        default_identity=default_identity)
 
     async_docdb_create = client.create_or_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_create.result()
@@ -187,7 +199,9 @@ def cli_cosmosdb_update(client,
                         network_acl_bypass_resource_ids=None,
                         server_version=None,
                         backup_interval=None,
-                        backup_retention=None):
+                        backup_retention=None,
+                        enable_system_assigned_identity=None,
+                        default_identity=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
 
@@ -216,6 +230,13 @@ def cli_cosmosdb_update(client,
     if enable_public_network is not None:
         public_network_access = 'Enabled' if enable_public_network else 'Disabled'
 
+    system_assigned_identity = None
+    if enable_system_assigned_identity is not None:
+        if enable_system_assigned_identity:
+            system_assigned_identity = ManagedServiceIdentity(type='SystemAssigned')
+        else:
+            system_assigned_identity = ManagedServiceIdentity(type='None')
+
     api_properties = {'ServerVersion': server_version}
 
     backup_policy = None
@@ -247,7 +268,9 @@ def cli_cosmosdb_update(client,
         network_acl_bypass=network_acl_bypass,
         network_acl_bypass_resource_ids=network_acl_bypass_resource_ids,
         api_properties=api_properties,
-        backup_policy=backup_policy)
+        backup_policy=backup_policy,
+        identity=system_assigned_identity,
+        default_identity=default_identity)
 
     async_docdb_update = client.update(resource_group_name, account_name, params)
     docdb_account = async_docdb_update.result()
@@ -1242,6 +1265,45 @@ def cli_cosmosdb_network_rule_list(client, resource_group_name, account_name):
     cosmos_db_account = client.get(resource_group_name, account_name)
     return cosmos_db_account.virtual_network_rules
 
+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):
+    """ Show the identity associated with a Cosmos DB account """
+
+    existing = client.get(resource_group_name, account_name)
+
+    if ("SystemAssigned" in existing.identity.type):
+        return existing.identity
+
+    identity = ManagedServiceIdentity(type='SystemAssigned')
+    params = DatabaseAccountUpdateParameters(identity=identity)
+    async_cosmos_db_update = client.update(resource_group_name, account_name, params)
+    cosmos_db_account = async_cosmos_db_update.result()
+    cosmos_db_account = client.get(resource_group_name, account_name)  # Workaround
+    return cosmos_db_account.identity
+
+def cli_cosmosdb_identity_remove(client,
+                                 resource_group_name,
+                                 account_name):
+    """ Show the identity associated with a Cosmos DB account """
+
+    existing = client.get(resource_group_name, account_name)
+
+    if (not "SystemAssigned" in existing.identity.type):
+        return existing.identity
+
+    identity = ManagedServiceIdentity(type='None')
+    params = DatabaseAccountUpdateParameters(identity=identity)
+    async_cosmos_db_update = client.update(resource_group_name, account_name, params)
+    cosmos_db_account = async_cosmos_db_update.result()
+    cosmos_db_account = client.get(resource_group_name, account_name)  # Workaround
+    return cosmos_db_account.identity
 
 def _get_virtual_network_id(cmd, resource_group_name, subnet, virtual_network):
     from azure.cli.core.commands.client_factory import get_subscription_id

src/azure-cli/azure/cli/command_modules/cosmosdb/tests/latest/test_cosmosdb_commands.py

@@ -1352,3 +1352,39 @@ def test_cosmosdb_key_vault_key_uri(self, resource_group):
         cmk_output = self.cmd('az cosmosdb create -n {acc} -g {rg} --locations regionName={location} failoverPriority=0 --key-uri {key_uri}').get_output_in_json()
 
         assert cmk_output["keyVaultKeyUri"] == key_uri
+
+    @ResourceGroupPreparer(name_prefix='cli_test_cosmosdb_managed_service_identity')
+    def test_cosmosdb_managed_service_identity(self, resource_group):
+        kv_name = self.create_random_name(prefix='cli', length=15)
+        key_name = self.create_random_name(prefix='cli', length=15)
+        key_uri = "https://{}.vault.azure.net/keys/{}".format(kv_name, key_name)
+
+        self.kwargs.update({
+            'acc': self.create_random_name(prefix='cli', length=15),
+            'kv_name': kv_name,
+            'key_name': key_name,
+            'key_uri': key_uri,
+            'location': "eastus2"
+        })
+
+        self.cmd('az keyvault create --resource-group {rg} -n {kv_name} --enable-soft-delete true --enable-purge-protection true')
+        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').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.cmd('az keyvault set-policy -n {kv_name} -g {rg} --spn {identity_principal_id} --key-permissions get unwrapKey wrapKey')
+
+        cmk_account = self.cmd('az cosmosdb update -n {acc} -g {rg} --default-identity SystemAssignedIdentity').get_output_in_json()
+        assert cmk_account["defaultIdentity"] == 'SystemAssignedIdentity'

src/azure-cli/requirements.py3.Darwin.txt

@@ -35,7 +35,7 @@ azure-mgmt-containerinstance==1.5.0
 azure-mgmt-containerregistry==3.0.0rc17
 azure-mgmt-containerservice==11.1.0
 azure-mgmt-core==1.2.1
-azure-mgmt-cosmosdb==3.0.0
+azure-mgmt-cosmosdb==5.0.0
 azure-mgmt-databoxedge==0.2.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1

src/azure-cli/requirements.py3.Linux.txt

@@ -35,7 +35,7 @@ azure-mgmt-containerinstance==1.5.0
 azure-mgmt-containerregistry==3.0.0rc17
 azure-mgmt-containerservice==11.1.0
 azure-mgmt-core==1.2.1
-azure-mgmt-cosmosdb==3.0.0
+azure-mgmt-cosmosdb==5.0.0
 azure-mgmt-databoxedge==0.2.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1

src/azure-cli/requirements.py3.windows.txt

@@ -35,7 +35,7 @@ azure-mgmt-containerinstance==1.5.0
 azure-mgmt-containerregistry==3.0.0rc17
 azure-mgmt-containerservice==11.1.0
 azure-mgmt-core==1.2.1
-azure-mgmt-cosmosdb==3.0.0
+azure-mgmt-cosmosdb==5.0.0
 azure-mgmt-databoxedge==0.2.0
 azure-mgmt-datalake-analytics==0.2.1
 azure-mgmt-datalake-nspkg==3.0.1

src/azure-cli/setup.py

@@ -76,7 +76,7 @@
     'azure-mgmt-consumption~=2.0',
     'azure-mgmt-containerinstance~=1.4',
     'azure-mgmt-containerregistry==3.0.0rc17',
-    'azure-mgmt-cosmosdb~=3.0.0',
+    'azure-mgmt-cosmosdb~=5.0.0',
     'azure-mgmt-containerservice~=11.1.0',
     'azure-mgmt-databoxedge~=0.2.0',
     'azure-mgmt-datalake-analytics~=0.2.1',