resolve comments

src/cosmosdb-preview/HISTORY.rst

@@ -5,7 +5,7 @@ Release History
 
 0.18.0
 ++++++
-* Retrieve and redistribute throughput at physical partition level.
+* Add support for retrieving and redistributing throughput at physical partition level.
 
 0.17.0
 ++++++

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -647,12 +647,12 @@
 
 helps['cosmosdb sql container retrieve-partition-throughput'] = """
     type: command
-    short-summary: "Retrieves the partition throughput of a sql container."
+    short-summary: "Retrieve  the partition throughput of a sql container."
     examples:
-      - name: Retrieves container container_name's throughput for specific physical partitions
+      - 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: Retrieves container container_name's throughput for all physical partitions
+               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
 """
@@ -666,17 +666,17 @@
                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}]'
+               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: "Retrieves the partition throughput of a mongodb collection."
+    short-summary: "Retrieve the partition throughput of a mongodb collection."
     examples:
-      - name: Retrieves container container_name's throughput for specific physical partitions
+      - 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: Retrieves container container_name's throughput for all physical partitions
+               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
 """
@@ -690,5 +690,5 @@
                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}]'
+               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
@@ -337,24 +343,24 @@ 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 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('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', 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}]'")
+        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'], required=False, help='comma separated list of physical partition ids')
+        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
@@ -363,5 +369,5 @@ def load_arguments(self, _):
         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}]'")
+        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))