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

src/cosmosdb-preview/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+0.18.0
+++++++
+* Retrieve and redistribute 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: "Retrieves the partition throughput of a sql container."
+    examples:
+      - name: Retrieves 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: Retrieves 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 '[{\"id\":8,\"throughput\":1200},{\"id\":6,\"throughput\":1200}]' --source-partition-info '[{\"id\":9}]'
+"""
+
+helps['cosmosdb mongodb collection retrieve-partition-throughput'] = """
+    type: command
+    short-summary: "Retrieves the partition throughput of a mongodb collection."
+    examples:
+      - name: Retrieves 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: Retrieves 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 '[{\"id\":8,\"throughput\":1200},{\"id\":6,\"throughput\":1200}]' --source-partition-info '[{\"id\":9}]'
+"""

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -331,3 +331,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'], required=False, help='comma 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', options_list=['--target-partition-info'], required=False, help="information about desired target physical partition throughput eg: '[{\"id\":0,\"throughput\":1200},{\"id\":1,\"throughput\":1200}]'")
+        c.argument('source_partition_info', options_list=['--source-partition-info'], required=False, help="information about source physical partition ids eg: '[{\"id\":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'], required=False, help='comma 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', options_list=['--target-partition-info'], required=False, help="information about desired target physical partition throughput eg: '[{\"id\":0,\"throughput\":1200},{\"id\":1,\"throughput\":1200}]'")
+        c.argument('source_partition_info', options_list=['--source-partition-info'], required=False, help="information about source physical partition ids eg: '[{\"id\":2}]'")

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)