[CosmosDB] Adding support to retrieve and redistribute physical partition throughput

src/cosmosdb-preview/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+0.18.0
+++++++
+* Add support for retrieving and redistributing throughput at physical partition level.
+
 0.17.0
 ++++++
 * Add support for new Continuous 7 Days backup mode

src/cosmosdb-preview/README.md

@@ -11,6 +11,7 @@ This package provides commands to
 - Update an Azure Cosmos DB database account to enable materialized views
 - Create/Delete a cosmosdb materialized views builder service resource
 - Provision and update database account with Continuous 7 days backup mode
+- Retrieve and redistribute throughput at physical partition level.
 
 ## How to use ##
 

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -644,3 +644,51 @@
         text: |-
                az cosmosdb mongodb collection merge -g my-resource-group -a my-account -d my-db --name my-mongodb-collection
 """
+
+helps['cosmosdb sql container retrieve-partition-throughput'] = """
+    type: command
+    short-summary: "Retrieve  the partition throughput of a sql container."
+    examples:
+      - name: Retrieve container container_name's throughput for specific physical partitions
+        text: |-
+               az cosmosdb sql container retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --physical-partition-ids 8 9"
+      - name: Retrieve container container_name's throughput for all physical partitions
+        text: |-
+               az cosmosdb sql container retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --all-partitions
+"""
+
+helps['cosmosdb sql container redistribute-partition-throughput'] = """
+    type: command
+    short-summary: "Redistributes the partition throughput of a sql container."
+    examples:
+      - name: Evenly distributes the partition throughput for a sql container among all physical partitions
+        text: |-
+               az cosmosdb sql container redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --evenly-distribute
+      - name: Redistributes the partition throughput for a sql container from source partitions to target partitions
+        text: |-
+               az cosmosdb sql container redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --target-partition-info 8=1200 6=1200]' --source-partition-info 9]'
+"""
+
+helps['cosmosdb mongodb collection retrieve-partition-throughput'] = """
+    type: command
+    short-summary: "Retrieve the partition throughput of a mongodb collection."
+    examples:
+      - name: Retrieve container container_name's throughput for specific physical partitions
+        text: |-
+               az cosmosdb mongodb collection retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --physical-partition-ids 8 9
+      - name: Retrieve container container_name's throughput for all physical partitions
+        text: |-
+               az cosmosdb mongodb collection retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --all-partitions
+"""
+
+helps['cosmosdb mongodb collection redistribute-partition-throughput'] = """
+    type: command
+    short-summary: "Redistributes the partition throughput of a mongodb collection."
+    examples:
+      - name: Evenly distributes the partition throughput for a mongodb collection among all physical partitions
+        text: |-
+               az cosmosdb mongodb collection redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --evenly-distribute
+      - name: Redistributes the partition throughput for a mongodb collection from source partitions to target partitions
+        text: |-
+               az cosmosdb mongodb collection redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name  --resource-group rg_name --target-partition-info 8=1200 6=1200' --source-partition-info 9'
+"""

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -18,7 +18,13 @@
     validate_mongo_user_definition_id)
 
 from azext_cosmosdb_preview.actions import (
-    CreateGremlinDatabaseRestoreResource, CreateTableRestoreResource, AddCassandraTableAction, AddSqlContainerAction)
+    CreateGremlinDatabaseRestoreResource,
+    CreateTableRestoreResource,
+    AddCassandraTableAction,
+    AddSqlContainerAction,
+    CreateTargetPhysicalPartitionThroughputInfoAction,
+    CreateSourcePhysicalPartitionThroughputInfoAction,
+    CreatePhysicalPartitionIdListAction)
 
 from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.models import (
     ContinuousTier
@@ -331,3 +337,37 @@ def load_arguments(self, _):
         c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
         c.argument('database_name', database_name_type, required=True, help='Name of the mongoDB database')
         c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the mongoDB collection')
+
+    # Sql container partition retrieve throughput
+    with self.argument_context('cosmosdb sql container retrieve-partition-throughput') as c:
+        c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
+        c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
+        c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB container')
+        c.argument('physical_partition_ids', options_list=['--physical-partition-ids', '-p'], nargs='+', action=CreatePhysicalPartitionIdListAction, required=False, help='space separated list of physical partition ids')
+        c.argument('all_partitions', arg_type=get_three_state_flag(), help="switch to retrieve throughput for all physical partitions")
+
+    # Sql container partition redistribute throughput
+    with self.argument_context('cosmosdb sql container redistribute-partition-throughput') as c:
+        c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
+        c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
+        c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB collection')
+        c.argument('evenly_distribute', arg_type=get_three_state_flag(), help="switch to distribute throughput equally among all physical partitions")        
+        c.argument('target_partition_info', nargs='+', action=CreateTargetPhysicalPartitionThroughputInfoAction, required=False, help="information about desired target physical partition throughput eg: 0=1200 1=1200")
+        c.argument('source_partition_info', nargs='+', action=CreateSourcePhysicalPartitionThroughputInfoAction, required=False, help="space separated source physical partition ids eg: 1 2")
+
+    # Mongodb collection partition retrieve throughput
+    with self.argument_context('cosmosdb mongodb collection retrieve-partition-throughput') as c:
+        c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
+        c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
+        c.argument('collection_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB container')
+        c.argument('physical_partition_ids', options_list=['--physical-partition-ids', '-p'], nargs='+', action=CreatePhysicalPartitionIdListAction, required=False, help='space separated list of physical partition ids')
+        c.argument('all_partitions', arg_type=get_three_state_flag(), help="switch to retrieve throughput for all physical partitions")
+
+    # Mongodb collection partition redistribute throughput
+    with self.argument_context('cosmosdb mongodb collection redistribute-partition-throughput') as c:
+        c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
+        c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
+        c.argument('collection_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB collection')
+        c.argument('evenly_distribute', arg_type=get_three_state_flag(), help="switch to distribute throughput equally among all physical partitions")
+        c.argument('target_partition_info', nargs='+', action=CreateTargetPhysicalPartitionThroughputInfoAction, required=False, help="information about desired target physical partition throughput eg: '0=1200 1=1200'")
+        c.argument('source_partition_info', nargs='+', action=CreateSourcePhysicalPartitionThroughputInfoAction, required=False, help="space separated source physical partition ids eg: 1 2")

src/cosmosdb-preview/azext_cosmosdb_preview/actions.py

@@ -13,7 +13,9 @@
     DatabaseRestoreResource,
     GremlinDatabaseRestoreResource,
     CosmosCassandraDataTransferDataSourceSink,
-    CosmosSqlDataTransferDataSourceSink
+    CosmosSqlDataTransferDataSourceSink,
+    PhysicalPartitionThroughputInfoResource,
+    PhysicalPartitionId
 )
 
 logger = get_logger(__name__)
@@ -173,3 +175,42 @@ def __call__(self, parser, namespace, values, option_string=None):
             namespace.dest_sql_container = sql_container
         else:
             namespace.sql_container = sql_container
+
+# pylint: disable=protected-access, too-few-public-methods
+class CreateTargetPhysicalPartitionThroughputInfoAction(argparse._AppendAction):
+    def __call__(self, parser, namespace, values, option_string=None):
+        if namespace.target_partition_info is None:
+            namespace.target_partition_info = []
+        if not values:
+            # pylint: disable=line-too-long
+            raise CLIError('usage error: --target-partition-info [PhysicalPartitionId1=Throughput1 PhysicalPartitionId2=Throughput2 ...]')
+        for item in values:
+            kvp = item.split('=', 1)
+            if len(kvp) != 2:
+                raise CLIError('usage error: --target-partition-info [PhysicalPartitionId1=Throughput1 PhysicalPartitionId2=Throughput2 ...]')
+            namespace.target_partition_info.append(
+                PhysicalPartitionThroughputInfoResource(id=kvp[0], throughput=kvp[1]))
+
+# pylint: disable=protected-access, too-few-public-methods
+class CreateSourcePhysicalPartitionThroughputInfoAction(argparse._AppendAction):
+    def __call__(self, parser, namespace, values, option_string=None):
+        if namespace.source_partition_info is None:
+            namespace.source_partition_info = []
+        if not values:
+            # pylint: disable=line-too-long
+            raise CLIError('usage error: --source-partition-info [PhysicalPartitionId1 PhysicalPartitionId2 ...]')
+        for item in values:
+            namespace.source_partition_info.append(
+                PhysicalPartitionThroughputInfoResource(id=item, throughput=0))
+
+# pylint: disable=protected-access, too-few-public-methods
+class CreatePhysicalPartitionIdListAction(argparse._AppendAction):
+    def __call__(self, parser, namespace, values, option_string=None):
+        if namespace.physical_partition_ids is None:
+            namespace.physical_partition_ids = []
+        if not values:
+            # pylint: disable=line-too-long
+            raise CLIError('usage error: --physical-partition-ids [PhysicalPartitionId1 PhysicalPartitionId2 ...]')
+        for item in values:
+            namespace.physical_partition_ids.append(
+                PhysicalPartitionId(id=item))

src/cosmosdb-preview/azext_cosmosdb_preview/commands.py

@@ -203,3 +203,19 @@ def load_command_table(self, _):
 
     with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
         g.custom_command('merge', 'cli_begin_list_mongo_db_collection_partition_merge', is_preview=True)
+
+    # Retrieve partition throughput for Sql containers
+    with self.command_group('cosmosdb sql container', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.custom_command('retrieve-partition-throughput', 'cli_begin_retrieve_sql_container_partition_throughput', is_preview=True)
+
+    # Redistribute partition throughput for Sql containers
+    with self.command_group('cosmosdb sql container', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.custom_command('redistribute-partition-throughput', 'cli_begin_redistribute_sql_container_partition_throughput', is_preview=True)
+
+    # Retrieve partition throughput for Mongo collection
+    with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
+        g.custom_command('retrieve-partition-throughput', 'cli_begin_retrieve_mongo_container_partition_throughput', is_preview=True)
+
+    # Redistribute partition throughput for Mongo collection
+    with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
+        g.custom_command('redistribute-partition-throughput', 'cli_begin_redistribute_mongo_container_partition_throughput', is_preview=True)