[CosmosDB] API Version 2024-12-01-preview (#8323)

* Adding custom SDK from azure-rest-api-specs PR

* Recordings

* Changes to add "sql container throughput" commands

* Added throughput type to sql container throughput migrate

* Added throughput bucket param

* Changes to add "sql container throughput" commands

* Added throughput type to sql container throughput migrate

* Added throughput bucket param

* Removed unused imports

* Remaining test recordings

* Fixed linters

* Fixed examples

* Flake8 changes

* Updated to --throughput-buckets and rerecorded test

* Added final recording

---------

Co-authored-by: Achint-Agrawal <[email protected]>

Author: pjohari-ms

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -1008,6 +1008,35 @@
                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.
+examples:
+      - name: Update the throughput of the SQL container under an Azure Cosmos DB SQL database.
+        text: |-
+               az cosmosdb sql container throughput update --resource-group resource_group --account-name database_account_name --database-name parent_database_name --name name_of_collection_needs_to_be_restored --throughput 1000
+"""
+
+helps['cosmosdb sql container throughput migrate'] = """
+type: command
+short-summary: Migrate the throughput of the SQL container between autoscale and manually provisioned.
+examples:
+      - name: Migrate the throughput of the SQL container between autoscale and manually provisioned.
+        text: |-
+               az cosmosdb sql container throughput migrate --resource-group resource_group --account-name database_account_name --database-name parent_database_name --name name_of_collection_needs_to_be_restored --throughput-type Manual
+"""
+
 # 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,17 @@
 }"
 """
 
+SQL_THROUGHPUT_BUCKETS_EXAMPLE = """--throughput-buckets "[
+    { \\"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 +511,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 +571,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'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Throughput Buckets, you can enter it as a string or as a file, e.g., --throughput-buckets @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 (
@@ -2355,6 +2357,53 @@ 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,
                                                               resource_group_name,