Added throughput bucket param

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -105,6 +105,12 @@
 }"
 """
 
+SQL_THROUGHPUT_BUCKETS_EXAMPLE = """--tb "[
+    { \\"id\\": 1, \\"maxThroughputPercentage\\" : 10 },
+    { \\"id\\": 2, \\"maxThroughputPercentage\\" : 20 }
+]"
+"""
+
 class ThroughputTypes(str, Enum):
     autoscale = "autoscale"
     manual = "manual"
@@ -571,6 +577,7 @@ def load_arguments(self, _):
         c.argument('container_name', options_list=['--name', '-n'], help="Container name")
         c.argument('throughput', type=int, help='The throughput of SQL container (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+        c.argument('throughput_buckets', options_list=['--throughput-buckets', '-tb'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Throughput Buckets, you can enter it as a string or as a file, e.g., --tb @throughput-buckets-file.json or ' + SQL_THROUGHPUT_BUCKETS_EXAMPLE)
 
     for scope in ['sql container throughput migrate']:
         with self.argument_context('cosmosdb {}'.format(scope)) as c:

src/cosmosdb-preview/azext_cosmosdb_preview/custom.py

@@ -2362,9 +2362,12 @@ def cli_cosmosdb_sql_container_throughput_update(client,
                                                  database_name,
                                                  container_name,
                                                  throughput=None,
-                                                 max_throughput=None):
+                                                 max_throughput=None,
+                                                 throughput_buckets=None):
     """Update an Azure Cosmos DB SQL container throughput"""
-    throughput_update_resource = _get_throughput_settings_update_parameters(throughput, max_throughput)
+    throughput_update_resource = _get_throughput_settings_update_parameters(throughput=throughput,
+                                                                            max_throughput = max_throughput,
+                                                                            throughput_buckets = throughput_buckets)
     return client.begin_update_sql_container_throughput(resource_group_name,
                                                         account_name,
                                                         database_name,
@@ -2384,14 +2387,16 @@ def cli_cosmosdb_sql_container_throughput_migrate(client,
     return client.begin_migrate_sql_container_to_manual_throughput(resource_group_name, account_name,
                                                                    database_name, container_name)
 
-def _get_throughput_settings_update_parameters(throughput=None, max_throughput=None):
+def _get_throughput_settings_update_parameters(throughput=None, max_throughput=None, throughput_buckets=None):
     throughput_resource = None
     if throughput and max_throughput:
         raise CLIError("Please provide max-throughput if your resource is autoscale enabled otherwise provide throughput.")
     if throughput:
-        throughput_resource = ThroughputSettingsResource(throughput=throughput)
+        throughput_resource = ThroughputSettingsResource(throughput=throughput, throughput_buckets=throughput_buckets)
     elif max_throughput:
-        throughput_resource = ThroughputSettingsResource(autoscale_settings=AutoscaleSettings(max_throughput=max_throughput))
+        throughput_resource = ThroughputSettingsResource(
+            autoscale_settings=AutoscaleSettings(max_throughput=max_throughput),
+            throughput_buckets=throughput_buckets)
 
     return ThroughputSettingsUpdateParameters(resource=throughput_resource)