[CosmosDB]Support az cosmosdb private-link-resource and private-endpoint-connection

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

@@ -78,6 +78,14 @@ def cf_cosmosdb(cli_ctx, **_):
     return get_mgmt_service_client(cli_ctx, CosmosDBManagementClient)
 
 
+def cf_db_private_endpoint_connections(cli_ctx, _):
+    return cf_cosmosdb(cli_ctx).private_endpoint_connections
+
+
+def cf_db_private_link_resources(cli_ctx, _):
+    return cf_cosmosdb(cli_ctx).private_link_resources
+
+
 def cf_db_accounts(cli_ctx, _):
     return cf_cosmosdb(cli_ctx).database_accounts
 

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

@@ -486,6 +486,61 @@
 short-summary: Manage Azure Comsos DB network rules.
 """
 
+
+helps['cosmosdb private-endpoint-connection'] = """
+type: group
+short-summary: Manage Azure Comsos DB private endpoint connections.
+"""
+
+helps['cosmosdb private-endpoint-connection approve'] = """
+type: command
+short-summary: Approve the specified private endpoint connection associated with Azure Comsos DB.
+examples:
+  - name: Approve the specified private endpoint connection associated with Azure Comsos DB.
+    text: az cosmosdb private-endpoint-connection approve --account-name MyAccount --name MyPrivateEndpoint --resource-group MyResourceGroup --description "Approved"
+"""
+
+
+helps['cosmosdb private-endpoint-connection delete'] = """
+type: command
+short-summary: Delete the specified private endpoint connection associated with Azure Comsos DB.
+examples:
+  - name: Delete the specified private endpoint connection associated with Azure Comsos DB.
+    text: az cosmosdb private-endpoint-connection delete --account-name MyAccount --name MyPrivateEndpoint --resource-group MyResourceGroup
+
+"""
+
+helps['cosmosdb private-endpoint-connection reject'] = """
+type: command
+short-summary: Reject the specified private endpoint connection associated with Azure Comsos DB.
+examples:
+  - name: Reject the specified private endpoint connection associated with Azure Comsos DB.
+    text: az cosmosdb private-endpoint-connection reject --account-name MyAccount --name MyPrivateEndpoint --resource-group MyResourceGroup --description "Rejected"
+"""
+
+
+helps['cosmosdb private-endpoint-connection show'] = """
+type: command
+short-summary: Show details of a private endpoint connection associated with Azure Comsos DB.
+examples:
+  - name: Show details of a private endpoint connection associated with Azure Comsos DB.
+    text: az cosmosdb private-endpoint-connection show --account-name MyAccount --name MyPrivateEndpoint --resource-group MyResourceGroup
+"""
+
+helps['cosmosdb private-link-resource'] = """
+type: group
+short-summary: Manage Azure Comsos DB private link resources.
+"""
+
+helps['cosmosdb private-link-resource list'] = """
+type: command
+short-summary: List the private link resources supported for Azure Comsos DB.
+example:
+  - name: List the private link resources supported for Azure Comsos DB.
+    text: cosmosdb private-link-resource list --account-name MyAccount --resource-group MyResourceGroup
+"""
+
+
 helps['cosmosdb regenerate-key'] = """
 type: command
 short-summary: Regenerate an access key for a Azure Cosmos DB database account.

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

@@ -102,6 +102,25 @@ def load_arguments(self, _):
     database_name_type = CLIArgumentType(options_list=['--database-name', '-d'], help='Database name.')
     container_name_type = CLIArgumentType(options_list=['--container-name', '-c'], help='Container name.')
 
+    with self.argument_context('cosmosdb private-endpoint-connection') as c:
+        c.argument('private_endpoint_connection_name', options_list=['--name', '-n'], required=False,
+                   help='The name of the private endpoint connection associated with Azure Cosmos DB. '
+                        'Required if --connection-id is not specified')
+        c.argument('account_name', account_name_type, required=False,
+                   help='Name of the Cosmos DB database account. Required if --connection-id is not specified')
+        c.argument('resource_group_name', required=False,
+                   help='The resource group name of specified Cosmos DB account. Required if --connection-id is not specified')
+
+    for item in ['approve', 'reject', 'delete', 'show']:
+        with self.argument_context('cosmosdb private-endpoint-connection {}'.format(item)) as c:
+            c.extra('connection_id', options_list=['--id'], required=False,
+                    help='The ID of the private endpoint connection associated with Azure Cosmos DB. '
+                         'If specified --account-name --resource-group/-g and --name/-n, this should be omitted.')
+            c.argument('description', options_list=['--description'], required=False, help='Comments for the {} operation.'.format(item))
+
+    with self.argument_context('cosmosdb private-link-resource') as c:
+        c.argument('account_name', account_name_type, required=True, help="Cosmosdb account name", id_part=None)
+
 # SQL database
     with self.argument_context('cosmosdb sql database') as c:
         c.argument('account_name', account_name_type, id_part=None)

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

@@ -3,6 +3,8 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from knack.util import CLIError
+
 
 def validate_failover_policies(ns):
     """ Extracts multiple space-separated failoverPolicies in regionName=failoverPriority format """
@@ -19,6 +21,20 @@ def validate_ip_range_filter(ns):
         ns.ip_range_filter = ",".join(ns.ip_range_filter)
 
 
+def validate_private_endpoint_connection_id(ns):
+    if ns.connection_id:
+        from azure.cli.core.util import parse_proxy_resource_id
+        result = parse_proxy_resource_id(ns.connection_id)
+        ns.resource_group_name = result['resource_group']
+        ns.account_name = result['name']
+        ns.private_endpoint_connection_name = result['child_name_1']
+
+    if not all([ns.account_name, ns.resource_group_name, ns.private_endpoint_connection_name]):
+        raise CLIError(None, 'incorrect usage: [--id ID | --name NAME --account-name NAME --resource-group NAME]')
+
+    del ns.connection_id
+
+
 def validate_capabilities(ns):
     """ Extracts multiple space-separated capabilities """
     from azure.mgmt.cosmosdb.models import Capability

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

@@ -8,7 +8,16 @@
 
 from azure.cli.core.commands import CliCommandType
 
-from azure.cli.command_modules.cosmosdb._client_factory import cf_db_accounts, cf_sql_resources, cf_mongo_db_resources, cf_cassandra_resources, cf_gremlin_resources, cf_table_resources
+from azure.cli.command_modules.cosmosdb._client_factory import (
+    cf_db_accounts,
+    cf_db_private_endpoint_connections,
+    cf_db_private_link_resources,
+    cf_sql_resources,
+    cf_mongo_db_resources,
+    cf_cassandra_resources,
+    cf_gremlin_resources,
+    cf_table_resources
+)
 
 from azure.cli.command_modules.cosmosdb._format import (
     database_output,
@@ -18,6 +27,10 @@
     list_connection_strings_output
 )
 
+from ._validators import (
+    validate_private_endpoint_connection_id
+)
+
 DATABASE_DEPRECATION_INFO = 'cosmosdb sql database, cosmosdb mongodb database, cosmosdb cassandra keyspace or cosmosdb gremlin database'
 
 COLLECTION_DEPRECATON_INFO = 'cosmosdb sql container, cosmosdb mongodb collection, cosmosdb cassandra table, cosmosdb gremlin graph or cosmosdb table'
@@ -29,6 +42,14 @@ def load_command_table(self, _):
         operations_tmpl='azure.mgmt.cosmosdb.operations#DatabaseAccountsOperations.{}',
         client_factory=cf_db_accounts)
 
+    cosmosdb_private_endpoint_connections_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.cosmosdb.operations#PrivateEndpointConnectionsOperations.{}',
+        client_factory=cf_db_private_endpoint_connections)
+
+    cosmosdb_private_link_resources_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.cosmosdb.operations#PrivateLinkResourcesOperations.{}',
+        client_factory=cf_db_private_link_resources)
+
     cosmosdb_sql_sdk = CliCommandType(
         operations_tmpl='azure.mgmt.cosmosdb.operations#SqlResourcesOperations.{}',
         client_factory=cf_sql_resources)
@@ -62,6 +83,22 @@ def load_command_table(self, _):
         g.custom_command('update', 'cli_cosmosdb_update')
         g.custom_command('list', 'cli_cosmosdb_list')
 
+    with self.command_group('cosmosdb private-endpoint-connection',
+                            cosmosdb_private_endpoint_connections_sdk,
+                            client_factory=cf_db_private_endpoint_connections) as g:
+        g.custom_command('approve', 'approve_private_endpoint_connection',
+                         validator=validate_private_endpoint_connection_id)
+        g.custom_command('reject', 'reject_private_endpoint_connection',
+                         validator=validate_private_endpoint_connection_id)
+        g.command('delete', 'delete', validator=validate_private_endpoint_connection_id)
+        g.show_command('show', 'get', validator=validate_private_endpoint_connection_id)
+
+    with self.command_group('cosmosdb private-link-resource',
+                            cosmosdb_private_link_resources_sdk,
+                            client_factory=cf_db_private_link_resources) as g:
+        from azure.cli.core.commands.transform import gen_dict_to_list_transform
+        g.show_command('list', 'list_by_database_account', transform=gen_dict_to_list_transform(key='values'))
+
     # SQL api
     with self.command_group('cosmosdb sql', is_preview=True):
         pass

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

@@ -991,6 +991,41 @@ def cli_cosmosdb_network_rule_remove(cmd,
     return docdb_account
 
 
+def _update_private_endpoint_connection_status(client, resource_group_name, account_name,
+                                               private_endpoint_connection_name, is_approved=True, description=None):
+    private_endpoint_connection = client.get(resource_group_name=resource_group_name, account_name=account_name,
+                                             private_endpoint_connection_name=private_endpoint_connection_name)
+
+    new_status = "Approved" if is_approved else "Rejected"
+    private_endpoint_connection.private_link_service_connection_state.status = new_status
+    private_endpoint_connection.private_link_service_connection_state.description = description
+
+    return client.create_or_update(resource_group_name=resource_group_name,
+                                   account_name=account_name,
+                                   private_endpoint_connection_name=private_endpoint_connection_name,
+                                   private_link_service_connection_state=private_endpoint_connection.private_link_service_connection_state)
+
+
+def approve_private_endpoint_connection(client, resource_group_name, account_name, private_endpoint_connection_name,
+                                        description=None):
+    """Approve a private endpoint connection request for Azure Cosmos DB."""
+
+    return _update_private_endpoint_connection_status(
+        client, resource_group_name, account_name, private_endpoint_connection_name, is_approved=True,
+        description=description
+    )
+
+
+def reject_private_endpoint_connection(client, resource_group_name, account_name, private_endpoint_connection_name,
+                                       description=None):
+    """Reject a private endpoint connection request for Azure Cosmos DB."""
+
+    return _update_private_endpoint_connection_status(
+        client, resource_group_name, account_name, private_endpoint_connection_name, is_approved=False,
+        description=description
+    )
+
+
 ######################
 # data plane APIs
 ######################