Adding Throughput Bucket Header

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -316,6 +316,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("correlatedActivityId"):
         headers[http_constants.HttpHeaders.CorrelatedActivityId] = options["correlatedActivityId"]
 
+    if options.get("throughputBucket"):
+        headers[http_constants.HttpHeaders.ThroughputBucket] = options["throughputBucket"]
+
     if resource_type == "docs" and verb != "get":
         if "responsePayloadOnWriteDisabled" in options:
             responsePayloadOnWriteDisabled = options["responsePayloadOnWriteDisabled"]

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -849,6 +849,7 @@ def create_item(  # pylint:disable=docstring-missing-param
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
         no_response: Optional[bool] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> CosmosDict:
         """Create an item in the container.
@@ -876,6 +877,7 @@ def create_item(  # pylint:disable=docstring-missing-param
         :keyword bool no_response: Indicates whether service should be instructed to skip sending
             response payloads. When not specified explicitly here, the default value will be determined from kwargs or
             when also not specified there from client-level kwargs.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists.
         :returns: A CosmosDict representing the new item. The dict will be empty if `no_response` is specified.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
@@ -896,6 +898,8 @@ def create_item(  # pylint:disable=docstring-missing-param
             kwargs['match_condition'] = match_condition
         if no_response is not None:
             kwargs['no_response'] = no_response
+        if throughput_bucket is not None:
+            kwargs["throughput_bucket"] = throughput_bucket
         request_options = build_options(kwargs)
         request_options["disableAutomaticIdGeneration"] = not enable_automatic_id_generation
         if populate_query_metrics:

sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py

@@ -188,6 +188,7 @@ class CosmosClient:  # pylint: disable=client-accepts-api-version-keyword
         level (to log all requests) or at a single request level. Requests will be logged at INFO level.
     :keyword bool no_response_on_write: Indicates whether service should be instructed to skip sending 
         response payloads on rite operations for items.
+    :keyword int throughput_bucket: The desired throughput bucket for the client
 
     .. admonition:: Example:
 
@@ -265,6 +266,7 @@ def create_database(  # pylint:disable=docstring-missing-param
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> DatabaseProxy:
         """
@@ -279,6 +281,7 @@ def create_database(  # pylint:disable=docstring-missing-param
             has changed, and act according to the condition specified by the `match_condition` parameter.
         :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag.
         :keyword response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :paramtype response_hook: Callable[[Mapping[str, str]], None]
         :returns: A DatabaseProxy instance representing the new database.
         :rtype: ~azure.cosmos.DatabaseProxy
@@ -308,6 +311,8 @@ def create_database(  # pylint:disable=docstring-missing-param
                 UserWarning,
             )
             request_options["populateQueryMetrics"] = populate_query_metrics
+        if throughput_bucket is not None:
+            kwargs["throughput_bucket"] = throughput_bucket
 
         _set_throughput_options(offer=offer_throughput, request_options=request_options)
         result = self.client_connection.CreateDatabase(database={"id": id}, options=request_options, **kwargs)
@@ -326,6 +331,7 @@ def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
         initial_headers: Optional[Dict[str, str]] = None,
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> DatabaseProxy:
         """
@@ -346,6 +352,7 @@ def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
             has changed, and act according to the condition specified by the `match_condition` parameter.
         :keyword ~azure.core.MatchConditions match_condition: The match condition to use upon the etag.
         :keyword Callable response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :returns: A DatabaseProxy instance representing the database.
         :rtype: ~azure.cosmos.DatabaseProxy
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The database read or creation failed.
@@ -358,6 +365,8 @@ def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
             kwargs["etag"] = etag
         if match_condition is not None:
             kwargs["match_condition"] = match_condition
+        if throughput_bucket is not None:
+            kwargs["throughput_bucket"] = throughput_bucket
         try:
             database_proxy = self.get_database_client(id)
             database_proxy.read(

sdk/cosmos/azure-cosmos/azure/cosmos/database.py

@@ -25,6 +25,8 @@
 from typing import Any, Dict, List, Union, Optional, Mapping, Callable
 
 import warnings
+
+from adodbapi.examples.db_print import kw_args
 from azure.core import MatchConditions
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.paging import ItemPaged
@@ -178,6 +180,7 @@ def create_container(  # pylint:disable=docstring-missing-param
         vector_embedding_policy: Optional[Dict[str, Any]] = None,
         change_feed_policy: Optional[Dict[str, Any]] = None,
         full_text_policy: Optional[Dict[str, Any]] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> ContainerProxy:
         """Create a new container with the given ID (name).
@@ -212,6 +215,7 @@ def create_container(  # pylint:disable=docstring-missing-param
         :keyword Dict[str, Any] full_text_policy: **provisional** The full text policy for the container.
             Used to denote the default language to be used for all full text indexes, or to individually
             assign a language to each full text index path.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :returns: A `ContainerProxy` instance representing the new container.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The container creation failed.
         :rtype: ~azure.cosmos.ContainerProxy
@@ -267,6 +271,8 @@ def create_container(  # pylint:disable=docstring-missing-param
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if throughput_bucket is not None:
+            kwargs['throughput_bucket'] = throughput_bucket
         request_options = build_options(kwargs)
         if populate_query_metrics is not None:
             warnings.warn(
@@ -302,6 +308,7 @@ def create_container_if_not_exists(  # pylint:disable=docstring-missing-param
         vector_embedding_policy: Optional[Dict[str, Any]] = None,
         change_feed_policy: Optional[Dict[str, Any]] = None,
         full_text_policy: Optional[Dict[str, Any]] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> ContainerProxy:
         """Create a container if it does not exist already.
@@ -338,6 +345,7 @@ def create_container_if_not_exists(  # pylint:disable=docstring-missing-param
         :keyword Dict[str, Any] full_text_policy: **provisional** The full text policy for the container.
             Used to denote the default language to be used for all full text indexes, or to individually
             assign a language to each full text index path.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :returns: A `ContainerProxy` instance representing the container.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The container read or creation failed.
         :rtype: ~azure.cosmos.ContainerProxy
@@ -370,6 +378,7 @@ def create_container_if_not_exists(  # pylint:disable=docstring-missing-param
                 vector_embedding_policy=vector_embedding_policy,
                 change_feed_policy=change_feed_policy,
                 full_text_policy=full_text_policy,
+                throughput_bucket=throughput_bucket,
                 **kwargs
             )
 
@@ -563,6 +572,7 @@ def replace_container(  # pylint:disable=docstring-missing-param
         analytical_storage_ttl: Optional[int] = None,
         computed_properties: Optional[List[Dict[str, str]]] = None,
         full_text_policy: Optional[Dict[str, Any]] = None,
+        throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> ContainerProxy:
         """Reset the properties of the container.
@@ -593,6 +603,7 @@ def replace_container(  # pylint:disable=docstring-missing-param
         :keyword Dict[str, Any] full_text_policy: **provisional** The full text policy for the container.
             Used to denote the default language to be used for all full text indexes, or to individually
             assign a language to each full text index path.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
         :returns: A `ContainerProxy` instance representing the container after replace completed.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Raised if the container couldn't be replaced.
             This includes if the container with given id does not exist.
@@ -615,6 +626,8 @@ def replace_container(  # pylint:disable=docstring-missing-param
             kwargs['etag'] = etag
         if match_condition is not None:
             kwargs['match_condition'] = match_condition
+        if throughput_bucket is not None:
+            kwargs['throughput_bucket'] = throughput_bucket
         request_options = build_options(kwargs)
         if populate_query_metrics is not None:
             warnings.warn(

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -189,6 +189,7 @@ class HttpHeaders:
     DisableRUPerMinuteUsage = "x-ms-documentdb-disable-ru-per-minute-usage"
     IsRUPerMinuteUsed = "x-ms-documentdb-is-ru-per-minute-used"
     OfferIsRUPerMinuteThroughputEnabled = "x-ms-offer-is-ru-per-minute-throughput-enabled"
+    ThroughputBucket = "x-ms-cosmos-throughput-bucket" #CHECK
 
     # Partitioned collection headers
     PartitionKey = "x-ms-documentdb-partitionkey"