[CosmosDB] API Version 2024-12-01-preview

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -1008,6 +1008,27 @@
                az cosmosdb sql container restore --resource-group resource_group --account-name database_account_name --database-name parent_database_name --name name_of_container_needs_to_be_restored --restore-timestamp 2020-07-13T16:03:41+0000
 """
 
+# sql container throughput commands
+helps['cosmosdb sql container throughput'] = """
+type: group
+short-summary: Manage throughput of SQL container under an Azure Cosmos DB account.
+"""
+
+helps['cosmosdb sql container throughput show'] = """
+type: command
+short-summary: Get the throughput of the SQL container under an Azure Cosmos DB SQL database.
+"""
+
+helps['cosmosdb sql container throughput update'] = """
+type: command
+short-summary: Update the throughput of the SQL container under an Azure Cosmos DB SQL database.
+"""
+
+helps['cosmosdb sql container throughput migrate'] = """
+type: command
+short-summary: Migrate the throughput of the SQL container between autoscale and manually provisioned.
+"""
+
 # in-account restore of a deleted mongodb database
 helps['cosmosdb mongodb database restore'] = """
     type: command

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 # pylint: disable=line-too-long, too-many-statements
 
+from enum import Enum
 import argparse
 from argcomplete.completers import FilesCompleter
 
@@ -104,6 +105,16 @@
 }"
 """
 
+SQL_THROUGHPUT_BUCKETS_EXAMPLE = """--tb "[
+    { \\"id\\": 1, \\"maxThroughputPercentage\\" : 10 },
+    { \\"id\\": 2, \\"maxThroughputPercentage\\" : 20 }
+]"
+"""
+
+class ThroughputTypes(str, Enum):
+    autoscale = "autoscale"
+    manual = "manual"
+
 
 def load_arguments(self, _):
     from knack.arguments import CLIArgumentType
@@ -499,6 +510,7 @@ def load_arguments(self, _):
             c.argument('job_name', options_list=['--job-name', '-n'], help='Name of the container copy job.', required=True)
 
     max_throughput_type = CLIArgumentType(options_list=['--max-throughput'], help='The maximum throughput resource can scale to (RU/s). Provided when the resource is autoscale enabled. The minimum value can be 4000 (RU/s)')
+    throughput_type = CLIArgumentType(options_list=['--throughput-type', '-t'], arg_type=get_enum_type(ThroughputTypes), help='The type of throughput to migrate to.')
 
 
 # SQL container
@@ -558,6 +570,19 @@ def load_arguments(self, _):
         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")
 
+    # Sql container throughput
+    with self.argument_context('cosmosdb sql container throughput') as c:
+        c.argument('account_name', account_name_type, id_part=None)
+        c.argument('database_name', database_name_type)
+        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:
+            c.argument('throughput_type', throughput_type)
+
     # 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')

src/cosmosdb-preview/azext_cosmosdb_preview/commands.py

@@ -234,6 +234,12 @@ def load_command_table(self, _):
     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)
 
+    # Get and update offer throughput for Sql containers
+    with self.command_group('cosmosdb sql container throughput', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.show_command('show', 'get_sql_container_throughput')
+        g.custom_command('update', 'cli_cosmosdb_sql_container_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_sql_container_throughput_migrate')
+
     # Merge partitions for Sql databases
     with self.command_group('cosmosdb sql database', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
         g.custom_command('merge', 'cli_begin_sql_database_partition_merge', is_preview=True)

src/cosmosdb-preview/azext_cosmosdb_preview/custom.py

@@ -63,7 +63,9 @@
     GraphAPIComputeServiceResourceCreateUpdateProperties,
     MaterializedViewsBuilderServiceResourceCreateUpdateProperties,
     DedicatedGatewayType,
-    ServiceType
+    ServiceType,
+    ThroughputSettingsResource,
+    ThroughputSettingsUpdateParameters
 )
 
 from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_mongocluster.models import (
@@ -2354,6 +2356,49 @@ def cli_begin_retrieve_sql_container_partition_throughput(client,
 
     return async_partition_retrieve_throughput_result.result()
 
+def cli_cosmosdb_sql_container_throughput_update(client,
+                                                 resource_group_name,
+                                                 account_name,
+                                                 database_name,
+                                                 container_name,
+                                                 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=throughput,
+                                                                            max_throughput = max_throughput,
+                                                                            throughput_buckets = throughput_buckets)
+    return client.begin_update_sql_container_throughput(resource_group_name,
+                                                        account_name,
+                                                        database_name,
+                                                        container_name,
+                                                        throughput_update_resource)
+
+def cli_cosmosdb_sql_container_throughput_migrate(client,
+                                                  resource_group_name,
+                                                  account_name,
+                                                  database_name,
+                                                  container_name,
+                                                  throughput_type):
+    """Migrate an Azure Cosmos DB SQL container throughput"""
+    if throughput_type == "autoscale":
+        return client.begin_migrate_sql_container_to_autoscale(resource_group_name, account_name,
+                                                               database_name, container_name)
+    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, 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_buckets=throughput_buckets)
+    elif max_throughput:
+        throughput_resource = ThroughputSettingsResource(
+            autoscale_settings=AutoscaleSettings(max_throughput=max_throughput),
+            throughput_buckets=throughput_buckets)
+
+    return ThroughputSettingsUpdateParameters(resource=throughput_resource)
 
 # pylint: disable=dangerous-default-value
 def cli_begin_redistribute_sql_container_partition_throughput(client,