[Cosmos] Add AutoScale Header Support 1/2 (Azure#24976)

* Update offer.py

* Update http_constants.py

* Update database.py

* Update _base.py

* Update container_management.py

* Update container_management_async.py

* Create test_auto_scale.py

* Create test_auto_scale_async.py

* Update database.py

* Update _database.py

* Update _base.py

* update readme and changelog

* Update _base.py

* Update offer.py

* Update _base.py

* Update container_management.py

* Update container_management_async.py

* Update test_auto_scale.py

* Update test_auto_scale_async.py

* Update offer.py

* Update test_auto_scale.py

* Update sdk/cosmos/azure-cosmos/samples/container_management_async.py

Co-authored-by: Anna Tisch <[email protected]>

* Update sdk/cosmos/azure-cosmos/samples/container_management.py

Co-authored-by: Anna Tisch <[email protected]>

* Update _base.py

* Update database.py

* Update offer.py

* Update test_auto_scale.py

* Update __init__.py

* Update test_auto_scale.py

* Update CHANGELOG.md

* Update __init__.py

* Update database.py

* Update test_auto_scale.py

* Update _database.py

* Update database.py

* Update test_auto_scale.py

* Update test_auto_scale_async.py

* Update database.py

* Update container.py

* Update _database.py

* Update _container.py

* Update container.py

* Update database.py

* Update _container.py

* Update _database.py

* Update database.py

* Update offer.py

* Update cosmos_client.py

* Update _cosmos_client.py

* Update offer.py

* Update test_auto_scale.py

* Update test_auto_scale_async.py

* code refactor

* Update _base.py

adding deserialize method

* Update _base.py

* Update _cosmos_client.py

* Update _database.py

* Update cosmos_client.py

* Update database.py

* renaming method

* database samples

* merge conflicts

* Update database.py

* Update _base.py

* Update _base.py

* Update _base.py

* Update __init__.py

* Update _base.py

* docstring changes

* Update CHANGELOG.md

* Update test_auto_scale_async.py

* Update test_auto_scale.py

* Update _base.py

* Update _base.py

* Update _base.py

* tests updates

* Update _base.py

* Update CHANGELOG.md

* Update _base.py

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

Co-authored-by: Travis Prescott <[email protected]>

* Update CHANGELOG.md

* Update CHANGELOG.md

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

Co-authored-by: Anna Tisch <[email protected]>

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

Co-authored-by: Anna Tisch <[email protected]>

* Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py

Co-authored-by: Anna Tisch <[email protected]>

* Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py

Co-authored-by: Anna Tisch <[email protected]>

* Update _base.py

* Update _database.py

* Update _cosmos_client.py

* Update cosmos_client.py

* Update database.py

* Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py

Co-authored-by: Simon Moreno <[email protected]>

* Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py

Co-authored-by: Anna Tisch <[email protected]>

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

Co-authored-by: Anna Tisch <[email protected]>

* Update _database.py

* Update _base.py

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

Co-authored-by: Anna Tisch <[email protected]>

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

Co-authored-by: Anna Tisch <[email protected]>

* Update offer.py

* Update database.py

* Update cosmos_client.py

* Update _database.py

* Update _base.py

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

Co-authored-by: Simon Moreno <[email protected]>

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

Co-authored-by: Simon Moreno <[email protected]>

Co-authored-by: Anna Tisch <[email protected]>
Co-authored-by: Simon Moreno <[email protected]>
Co-authored-by: Travis Prescott <[email protected]>

Author: 4 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,10 +1,11 @@
 ## Release History
 
-### 4.3.1 (2022-08-12)
+### 4.3.1 (Unreleased)
 
 #### Features Added
 - GA release of integrated cache functionality. For more information on integrated cache please see [Azure Cosmos DB integrated cache](https://docs.microsoft.com/azure/cosmos-db/integrated-cache).
 - Added ability to replace analytical ttl on containers. For more information on analytical ttl please see [Azure Cosmos DB analytical store](https://docs.microsoft.com/azure/cosmos-db/analytical-store-introduction).
+- Added the ability to create containers and databases with autoscale properties for the sync and async clients.
 
 #### Bugs Fixed
 - Fixed parsing of args for overloaded `container.read()` method.

sdk/cosmos/azure-cosmos/README.md

@@ -162,7 +162,6 @@ Currently the features below are **not supported**. For alternatives options, ch
 
 * Get CollectionSizeUsage, DatabaseUsage, and DocumentUsage metrics
 * Create Geospatial Index
-* Provision Autoscale DBs or containers
 * Update Autoscale throughput
 * Get the connection string
 * Get the minimum RU/s of a container

sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py

@@ -27,6 +27,7 @@
 from .user import UserProxy
 from .scripts import ScriptsProxy
 from .offer import Offer
+from .offer import ThroughputProperties
 from .documents import (
     ConsistencyLevel,
     DataType,
@@ -62,5 +63,6 @@
     "TriggerOperation",
     "TriggerType",
     "ConnectionRetryPolicy",
+    "ThroughputProperties",
 )
 __version__ = VERSION

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

@@ -27,18 +27,17 @@
 import json
 import uuid
 import binascii
-from typing import Dict, Any
+from typing import Dict, Any, Union
 
 from urllib.parse import quote as urllib_quote
 from urllib.parse import urlsplit
-
 from azure.core import MatchConditions
-
 from . import auth
 from . import documents
 from . import partition_key
 from . import http_constants
 from . import _runtime_constants
+from .offer import ThroughputProperties
 
 # pylint: disable=protected-access
 
@@ -297,6 +296,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("maxIntegratedCacheStaleness"):
         headers[http_constants.HttpHeaders.DedicatedGatewayCacheStaleness] = options["maxIntegratedCacheStaleness"]
 
+    if options.get("autoUpgradePolicy"):
+        headers[http_constants.HttpHeaders.AutoscaleSettings] = options["autoUpgradePolicy"]
+
     return headers
 
 
@@ -674,3 +676,60 @@ def validate_cache_staleness_value(max_integrated_cache_staleness):
     if max_integrated_cache_staleness < 0:
         raise ValueError("Parameter 'max_integrated_cache_staleness_in_ms' can only be an "
                          "integer greater than or equal to zero")
+
+
+def _stringify_auto_scale(offer: Dict[str, Any]) -> Any:
+    auto_scale_params = None
+    max_throughput = offer.auto_scale_max_throughput
+    increment_percent = offer.auto_scale_increment_percent
+    auto_scale_params = {"maxThroughput": max_throughput}
+    if increment_percent is not None:
+        auto_scale_params["autoUpgradePolicy"] = {"throughputPolicy": {"incrementPercent": increment_percent}}
+    auto_scale_settings = json.dumps(auto_scale_params)
+
+    return auto_scale_settings
+
+
+def _set_throughput_options(offer: Union[int, ThroughputProperties], request_options: Dict[str, Any]) -> Any:
+    if offer is not None:
+        try:
+            max_throughput = offer.auto_scale_max_throughput
+            increment_percent = offer.auto_scale_increment_percent
+
+            if max_throughput is not None:
+                request_options['autoUpgradePolicy'] = _stringify_auto_scale(offer=offer)
+            elif increment_percent:
+                raise ValueError("auto_scale_max_throughput must be supplied in "
+                                 "conjunction with auto_scale_increment_percent")
+            if offer.offer_throughput:
+                request_options["offerThroughput"] = offer.offer_throughput
+
+        except AttributeError:
+            if isinstance(offer, int):
+                request_options["offerThroughput"] = offer
+            else:
+                raise TypeError("offer_throughput must be int or an instance of ThroughputProperties")
+
+
+def _deserialize_throughput(throughput: list) -> Any:
+    throughput_properties = throughput
+    try:
+        max_throughput = throughput_properties[0]['content']['offerAutopilotSettings']['maxThroughput']
+    except (KeyError, TypeError):  # Adding TypeError just in case one of these dicts is None
+        max_throughput = None
+    try:
+        increment_percent = \
+            throughput_properties[0]['content']['offerAutopilotSettings']['autoUpgradePolicy']['throughputPolicy'][
+                'incrementPercent']
+    except (KeyError, TypeError):
+        increment_percent = None
+    try:
+        throughput = throughput_properties[0]["content"]["offerThroughput"]
+    except (KeyError, TypeError):
+        throughput = None
+    return ThroughputProperties(
+        auto_scale_max_throughput=max_throughput,
+        auto_scale_increment_percent=increment_percent,
+        offer_throughput=throughput,
+        properties=throughput_properties[0]
+    )

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -29,7 +29,7 @@
 from azure.core.tracing.decorator_async import distributed_trace_async  # type: ignore
 
 from ._cosmos_client_connection_async import CosmosClientConnection
-from .._base import build_options as _build_options, validate_cache_staleness_value
+from .._base import build_options as _build_options, validate_cache_staleness_value, _deserialize_throughput
 from ..exceptions import CosmosResourceNotFoundError
 from ..http_constants import StatusCodes
 from ..offer import ThroughputProperties
@@ -594,8 +594,7 @@ async def get_throughput(self, **kwargs: Any) -> ThroughputProperties:
         if response_hook:
             response_hook(self.client_connection.last_response_headers, throughput_properties)
 
-        return ThroughputProperties(offer_throughput=throughput_properties[0]["content"]["offerThroughput"],
-                                    properties=throughput_properties[0])
+        return _deserialize_throughput(throughput=throughput_properties)
 
     @distributed_trace_async
     async def replace_throughput(self, throughput: int, **kwargs: Any) -> ThroughputProperties:

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py

@@ -31,7 +31,8 @@
 
 from ..cosmos_client import _parse_connection_str, _build_auth
 from ._cosmos_client_connection_async import CosmosClientConnection
-from .._base import build_options as _build_options
+from .._base import build_options as _build_options, _set_throughput_options
+from ..offer import ThroughputProperties
 from ._retry_utility_async import _ConnectionRetryPolicy
 from ._database import DatabaseProxy
 from ..documents import ConnectionPolicy, DatabaseAccount
@@ -196,7 +197,8 @@ async def create_database(  # pylint: disable=redefined-builtin
         Create a new database with the given ID (name).
 
         :param str id: ID (name) of the database to create.
-        :keyword int offer_throughput: The provisioned throughput for this offer.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :keyword str session_token: Token for use with Session consistency.
         :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
         :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource
@@ -223,8 +225,7 @@ async def create_database(  # pylint: disable=redefined-builtin
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
         offer_throughput = kwargs.pop('offer_throughput', None)
-        if offer_throughput is not None:
-            request_options["offerThroughput"] = offer_throughput
+        _set_throughput_options(offer=offer_throughput, request_options=request_options)
 
         result = await self.client_connection.CreateDatabase(database=dict(id=id), options=request_options, **kwargs)
         if response_hook:

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py

@@ -30,7 +30,7 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from ._cosmos_client_connection_async import CosmosClientConnection
-from .._base import build_options as _build_options
+from .._base import build_options as _build_options, _set_throughput_options, _deserialize_throughput
 from ._container import ContainerProxy
 from ..offer import ThroughputProperties
 from ..http_constants import StatusCodes
@@ -164,7 +164,8 @@ async def create_container(
         :keyword dict[str, str] indexing_policy: The indexing policy to apply to the container.
         :keyword int default_ttl: Default time to live (TTL) for items in the container.
             If unspecified, items do not expire.
-        :keyword int offer_throughput: The provisioned throughput for this offer.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :keyword dict[str, str] unique_key_policy: The unique key policy to apply to the container.
         :keyword dict[str, str] conflict_resolution_policy: The conflict resolution policy to apply to the container.
         :keyword str session_token: Token for use with Session consistency.
@@ -227,8 +228,7 @@ async def create_container(
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
         offer_throughput = kwargs.pop('offer_throughput', None)
-        if offer_throughput is not None:
-            request_options["offerThroughput"] = offer_throughput
+        _set_throughput_options(offer=offer_throughput, request_options=request_options)
 
         data = await self.client_connection.CreateContainer(
             database_link=self.database_link, collection=definition, options=request_options, **kwargs
@@ -258,7 +258,8 @@ async def create_container_if_not_exists(
         :keyword dict[str, str] indexing_policy: The indexing policy to apply to the container.
         :keyword int default_ttl: Default time to live (TTL) for items in the container.
             If unspecified, items do not expire.
-        :keyword int offer_throughput: The provisioned throughput for this offer.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :keyword dict[str, str] unique_key_policy: The unique key policy to apply to the container.
         :keyword dict[str, str] conflict_resolution_policy: The conflict resolution policy to apply to the container.
         :keyword str session_token: Token for use with Session consistency.
@@ -749,9 +750,7 @@ async def get_throughput(self, **kwargs: Any) -> ThroughputProperties:
         if response_hook:
             response_hook(self.client_connection.last_response_headers, throughput_properties)
 
-        return ThroughputProperties(offer_throughput=throughput_properties[0]["content"]["offerThroughput"],
-                                    properties=throughput_properties[0])
-
+        return _deserialize_throughput(throughput=throughput_properties)
 
     @distributed_trace_async
     async def replace_throughput(self, throughput: int, **kwargs: Any) -> ThroughputProperties:

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

@@ -29,7 +29,7 @@
 from azure.core.tracing.decorator import distributed_trace  # type: ignore
 
 from ._cosmos_client_connection import CosmosClientConnection
-from ._base import build_options, validate_cache_staleness_value
+from ._base import build_options, validate_cache_staleness_value, _deserialize_throughput
 from .exceptions import CosmosResourceNotFoundError
 from .http_constants import StatusCodes
 from .offer import ThroughputProperties
@@ -665,8 +665,7 @@ def get_throughput(self, **kwargs):
         if response_hook:
             response_hook(self.client_connection.last_response_headers, throughput_properties)
 
-        return ThroughputProperties(offer_throughput=throughput_properties[0]["content"]["offerThroughput"],
-                                    properties=throughput_properties[0])
+        return _deserialize_throughput(throughput=throughput_properties)
 
     @distributed_trace
     def replace_throughput(self, throughput, **kwargs):

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

@@ -27,7 +27,8 @@
 from azure.core.tracing.decorator import distributed_trace  # type: ignore
 
 from ._cosmos_client_connection import CosmosClientConnection
-from ._base import build_options
+from ._base import build_options, _set_throughput_options
+from .offer import ThroughputProperties
 from ._retry_utility import ConnectionRetryPolicy
 from .database import DatabaseProxy
 from .documents import ConnectionPolicy, DatabaseAccount
@@ -230,15 +231,16 @@ def create_database(  # pylint: disable=redefined-builtin
         self,
         id,  # type: str
         populate_query_metrics=None,  # type: Optional[bool]
-        offer_throughput=None,  # type: Optional[int]
+        offer_throughput=None,  # type: Optional[Union[int, ThroughputProperties]]
         **kwargs  # type: Any
     ):
         # type: (...) -> DatabaseProxy
         """
         Create a new database with the given ID (name).
 
         :param id: ID (name) of the database to create.
-        :param int offer_throughput: The provisioned throughput for this offer.
+        :param offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :keyword str session_token: Token for use with Session consistency.
         :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request.
         :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource
@@ -268,8 +270,7 @@ def create_database(  # pylint: disable=redefined-builtin
                 UserWarning,
             )
             request_options["populateQueryMetrics"] = populate_query_metrics
-        if offer_throughput is not None:
-            request_options["offerThroughput"] = offer_throughput
+        _set_throughput_options(offer=offer_throughput, request_options=request_options)
 
         result = self.client_connection.CreateDatabase(database=dict(id=id), options=request_options, **kwargs)
         if response_hook:
@@ -281,7 +282,7 @@ def create_database_if_not_exists(  # pylint: disable=redefined-builtin
         self,
         id,  # type: str
         populate_query_metrics=None,  # type: Optional[bool]
-        offer_throughput=None,  # type: Optional[int]
+        offer_throughput=None,  # type: Optional[Union[int, ThroughputProperties]]
         **kwargs  # type: Any
     ):
         # type: (...) -> DatabaseProxy
@@ -296,7 +297,8 @@ def create_database_if_not_exists(  # pylint: disable=redefined-builtin
 
         :param id: ID (name) of the database to read or create.
         :param bool populate_query_metrics: Enable returning query metrics in response headers.
-        :param int offer_throughput: The provisioned throughput for this offer.
+        :param offer_throughput: The provisioned throughput for this offer.
+        :type offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :keyword str session_token: Token for use with Session consistency.
         :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request.
         :keyword str etag: An ETag value, or the wildcard character (*). Used to check if the resource

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

@@ -28,7 +28,7 @@
 from azure.core.tracing.decorator import distributed_trace  # type: ignore
 
 from ._cosmos_client_connection import CosmosClientConnection
-from ._base import build_options
+from ._base import build_options, _set_throughput_options, _deserialize_throughput
 from .container import ContainerProxy
 from .offer import ThroughputProperties
 from .http_constants import StatusCodes
@@ -155,7 +155,7 @@ def create_container(
         indexing_policy=None,  # type: Optional[Dict[str, Any]]
         default_ttl=None,  # type: Optional[int]
         populate_query_metrics=None,  # type: Optional[bool]
-        offer_throughput=None,  # type: Optional[int]
+        offer_throughput=None,  # type: Optional[Union[int, ThroughputProperties]]
         unique_key_policy=None,  # type: Optional[Dict[str, Any]]
         conflict_resolution_policy=None,  # type: Optional[Dict[str, Any]]
         **kwargs  # type: Any
@@ -170,6 +170,7 @@ def create_container(
         :param indexing_policy: The indexing policy to apply to the container.
         :param default_ttl: Default time to live (TTL) for items in the container. If unspecified, items do not expire.
         :param offer_throughput: The provisioned throughput for this offer.
+        :type offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :param unique_key_policy: The unique key policy to apply to the container.
         :param conflict_resolution_policy: The conflict resolution policy to apply to the container.
         :keyword str session_token: Token for use with Session consistency.
@@ -232,9 +233,7 @@ def create_container(
                 UserWarning,
             )
             request_options["populateQueryMetrics"] = populate_query_metrics
-        if offer_throughput is not None:
-            request_options["offerThroughput"] = offer_throughput
-
+        _set_throughput_options(offer=offer_throughput, request_options=request_options)
         data = self.client_connection.CreateContainer(
             database_link=self.database_link, collection=definition, options=request_options, **kwargs
         )
@@ -252,7 +251,7 @@ def create_container_if_not_exists(
         indexing_policy=None,  # type: Optional[Dict[str, Any]]
         default_ttl=None,  # type: Optional[int]
         populate_query_metrics=None,  # type: Optional[bool]
-        offer_throughput=None,  # type: Optional[int]
+        offer_throughput=None,  # type: Optional[Union[int, ThroughputProperties]]
         unique_key_policy=None,  # type: Optional[Dict[str, Any]]
         conflict_resolution_policy=None,  # type: Optional[Dict[str, Any]]
         **kwargs  # type: Any
@@ -270,6 +269,7 @@ def create_container_if_not_exists(
         :param default_ttl: Default time to live (TTL) for items in the container. If unspecified, items do not expire.
         :param populate_query_metrics: Enable returning query metrics in response headers.
         :param offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: int or ~azure.cosmos.ThroughputProperties.
         :param unique_key_policy: The unique key policy to apply to the container.
         :param conflict_resolution_policy: The conflict resolution policy to apply to the container.
         :keyword str session_token: Token for use with Session consistency.
@@ -774,8 +774,7 @@ def get_throughput(self, **kwargs):
         if response_hook:
             response_hook(self.client_connection.last_response_headers, throughput_properties)
 
-        return ThroughputProperties(offer_throughput=throughput_properties[0]["content"]["offerThroughput"],
-                                    properties=throughput_properties[0])
+        return _deserialize_throughput(throughput=throughput_properties)
 
     @distributed_trace
     def replace_throughput(self, throughput, **kwargs):