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

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

@@ -348,6 +348,35 @@
 short-summary: Update an Gremlin graph under an Azure Cosmos DB Gremlin database.
 """
 
+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
+"""
+
+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
+"""
+
 helps['cosmosdb keys'] = """
 type: group
 short-summary: Manage Azure Cosmos DB keys.

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', nargs='*', help="accept system or user assigned identities separated by spaces. Use '[system]' to refer system assigned identity. Currently only system assigned identity is supported.", 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. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more.", 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', client_factory=cf_db_accounts, 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')
+
     # 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

@@ -9,6 +9,7 @@
 from knack.log import get_logger
 from knack.util import CLIError
 from msrestazure.azure_exceptions import CloudError
+from azure.cli.core.azclierror import InvalidArgumentValueError
 
 from azure.mgmt.cosmosdb.models import (
     ConsistencyPolicy,
@@ -22,6 +23,7 @@
     SqlContainerResource,
     SqlContainerCreateUpdateParameters,
     ContainerPartitionKey,
+    ResourceIdentityType,
     SqlStoredProcedureResource,
     SqlStoredProcedureCreateUpdateParameters,
     SqlTriggerResource,
@@ -30,6 +32,7 @@
     SqlUserDefinedFunctionCreateUpdateParameters,
     TableResource,
     TableCreateUpdateParameters,
+    ManagedServiceIdentity,
     MongoDBDatabaseResource,
     MongoDBDatabaseCreateUpdateParameters,
     MongoDBCollectionResource,
@@ -99,7 +102,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 +127,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 == [] or (len(assign_identity) == 1 and assign_identity[0] == '[system]'):
+            system_assigned_identity = ManagedServiceIdentity(type=ResourceIdentityType.system_assigned.value)
+        else:
+            raise InvalidArgumentValueError("Only '[system]' is supported right now for command '--assign-identity'.")
+
     api_properties = {}
     if kind == DatabaseAccountKind.mongo_db.value:
         api_properties['ServerVersion'] = server_version
@@ -157,7 +169,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 +201,8 @@ def cli_cosmosdb_update(client,
                         network_acl_bypass_resource_ids=None,
                         server_version=None,
                         backup_interval=None,
-                        backup_retention=None):
+                        backup_retention=None,
+                        default_identity=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
 
@@ -247,7 +262,8 @@ 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,
+        default_identity=default_identity)
 
     async_docdb_update = client.update(resource_group_name, account_name, params)
     docdb_account = async_docdb_update.result()
@@ -1243,6 +1259,53 @@ def cli_cosmosdb_network_rule_list(client, 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 ResourceIdentityType.system_assigned.value in existing.identity.type:
+        return existing.identity
+
+    if existing.identity.type == ResourceIdentityType.user_assigned.value:
+        identity = ManagedServiceIdentity(type=ResourceIdentityType.system_assigned_user_assigned.value)
+    else:
+        identity = ManagedServiceIdentity(type=ResourceIdentityType.system_assigned.value)
+    params = DatabaseAccountUpdateParameters(identity=identity)
+    async_cosmos_db_update = client.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):
+    """ Remove the SystemAssigned identity associated with a Cosmos DB account """
+
+    existing = client.get(resource_group_name, account_name)
+
+    if ResourceIdentityType.system_assigned.value not in existing.identity.type:
+        return existing.identity
+
+    if ResourceIdentityType.user_assigned.value in existing.identity.type:
+        identity = ManagedServiceIdentity(type=ResourceIdentityType.user_assigned.value)
+    else:
+        identity = ManagedServiceIdentity(type=ResourceIdentityType.none.value)
+    params = DatabaseAccountUpdateParameters(identity=identity)
+    async_cosmos_db_update = client.update(resource_group_name, account_name, params)
+    cosmos_db_account = async_cosmos_db_update.result()
+    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
     from msrestazure.tools import is_valid_resource_id, resource_id