Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
108 commits
Select commit Hold shift + click to select a range
e1e178f
Update offer.py
JericHunter Jun 24, 2022
c32ac36
Update http_constants.py
JericHunter Jun 24, 2022
19feed2
Update database.py
JericHunter Jun 24, 2022
8e41793
Update _base.py
JericHunter Jun 24, 2022
b67b545
Update container_management.py
JericHunter Jul 8, 2022
7ef9375
Update container_management_async.py
JericHunter Jul 8, 2022
9d57b15
Create test_auto_scale.py
JericHunter Jul 8, 2022
b04a659
Create test_auto_scale_async.py
JericHunter Jul 8, 2022
6b5af70
Update database.py
JericHunter Jul 8, 2022
465097b
Update _database.py
JericHunter Jul 8, 2022
00d2756
Update _base.py
JericHunter Jul 8, 2022
4a74315
update readme and changelog
JericHunter Jul 8, 2022
b279ddd
Update _base.py
JericHunter Jul 8, 2022
30cf250
Update offer.py
JericHunter Jul 8, 2022
63122ee
Update _base.py
JericHunter Jul 13, 2022
6fe9862
Update container_management.py
JericHunter Jul 13, 2022
39521e4
Update container_management_async.py
JericHunter Jul 13, 2022
f8e4d87
Update test_auto_scale.py
JericHunter Jul 13, 2022
be72c3a
Update test_auto_scale_async.py
JericHunter Jul 13, 2022
7f983c8
Update offer.py
JericHunter Jul 13, 2022
840c179
Update test_auto_scale.py
JericHunter Jul 21, 2022
c9cc5a9
Update sdk/cosmos/azure-cosmos/samples/container_management_async.py
JericHunter Jul 26, 2022
eb58d5e
Update sdk/cosmos/azure-cosmos/samples/container_management.py
JericHunter Jul 26, 2022
a2b5d26
Update _base.py
JericHunter Jul 26, 2022
ed806dc
Update database.py
JericHunter Jul 26, 2022
8d150d2
Update offer.py
JericHunter Jul 26, 2022
869ddea
Update test_auto_scale.py
JericHunter Jul 26, 2022
0563ddd
Update __init__.py
JericHunter Aug 8, 2022
17ac051
Update test_auto_scale.py
JericHunter Aug 8, 2022
a664700
Update CHANGELOG.md
JericHunter Aug 8, 2022
706bdb7
Update __init__.py
JericHunter Aug 8, 2022
a792cf4
Update database.py
JericHunter Aug 8, 2022
4518b96
Update test_auto_scale.py
JericHunter Aug 8, 2022
c7d040a
Update _database.py
JericHunter Aug 9, 2022
3536aec
Update database.py
JericHunter Aug 9, 2022
98a2c43
Update test_auto_scale.py
JericHunter Aug 10, 2022
c5d68a6
Update test_auto_scale_async.py
JericHunter Aug 10, 2022
2214464
Update database.py
JericHunter Aug 10, 2022
3ab778f
Update container.py
JericHunter Aug 10, 2022
d73247b
Update _database.py
JericHunter Aug 10, 2022
13afb9f
Update _container.py
JericHunter Aug 10, 2022
e5cdf39
Update container.py
JericHunter Aug 10, 2022
a33bbc5
Update database.py
JericHunter Aug 10, 2022
f522365
Update _container.py
JericHunter Aug 10, 2022
fced131
Update _database.py
JericHunter Aug 10, 2022
7238446
Update database.py
JericHunter Aug 10, 2022
b408647
Update offer.py
JericHunter Aug 10, 2022
6a9203f
Update cosmos_client.py
JericHunter Aug 10, 2022
9ddbe8d
Update _cosmos_client.py
JericHunter Aug 10, 2022
7094c15
Update offer.py
JericHunter Aug 10, 2022
82b3063
Update test_auto_scale.py
JericHunter Aug 10, 2022
6eeac0f
Update test_auto_scale_async.py
JericHunter Aug 11, 2022
095c55d
code refactor
JericHunter Aug 11, 2022
5f908bc
Update _base.py
JericHunter Aug 11, 2022
9337842
Update _base.py
JericHunter Aug 11, 2022
60cebb4
Update _cosmos_client.py
JericHunter Aug 11, 2022
e7137ee
Update _database.py
JericHunter Aug 11, 2022
260d66e
Update cosmos_client.py
JericHunter Aug 11, 2022
0be7d49
Update database.py
JericHunter Aug 11, 2022
16fdf83
renaming method
JericHunter Aug 11, 2022
f643289
database samples
JericHunter Aug 12, 2022
a8eaaa5
merge conflicts
JericHunter Aug 12, 2022
178648c
Merge branch 'main' into jh-auto-scale-v2
JericHunter Aug 12, 2022
5fffb43
Update database.py
simorenoh Aug 12, 2022
fe350a5
Update _base.py
JericHunter Aug 12, 2022
6117b4c
Merge branch 'jh-auto-scale-v2' of https://github.com/Azure/azure-sdk…
JericHunter Aug 12, 2022
b5f2e43
Update _base.py
JericHunter Aug 12, 2022
6e6fbfb
Update _base.py
JericHunter Aug 12, 2022
89092c5
Update __init__.py
JericHunter Aug 12, 2022
9458b1c
Update _base.py
JericHunter Aug 24, 2022
be4e54a
docstring changes
JericHunter Aug 24, 2022
bbbd4fb
Update CHANGELOG.md
JericHunter Aug 24, 2022
831e355
Update test_auto_scale_async.py
JericHunter Aug 24, 2022
e76fcc9
Update test_auto_scale.py
JericHunter Aug 24, 2022
e26ae85
Update _base.py
JericHunter Aug 24, 2022
7ad381f
Update _base.py
JericHunter Aug 24, 2022
afc6e79
Update _base.py
JericHunter Aug 29, 2022
b68a1ce
tests updates
JericHunter Aug 30, 2022
1aa14d1
Update _base.py
JericHunter Aug 30, 2022
8aac6e6
Merge branch 'main' into jh-auto-scale-v2
JericHunter Aug 30, 2022
a201fec
Update CHANGELOG.md
JericHunter Aug 30, 2022
0bfea4c
Update _base.py
JericHunter Aug 30, 2022
6176a46
Update sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
JericHunter Aug 30, 2022
970d349
Update CHANGELOG.md
JericHunter Aug 31, 2022
814fdb0
Update CHANGELOG.md
simorenoh Sep 2, 2022
a851a14
Update sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py
JericHunter Sep 6, 2022
3cdf785
Update sdk/cosmos/azure-cosmos/azure/cosmos/database.py
JericHunter Sep 6, 2022
5032739
Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py
JericHunter Sep 6, 2022
1f317c9
Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py
JericHunter Sep 6, 2022
b37afc8
Update _base.py
JericHunter Sep 6, 2022
613958a
Update _database.py
JericHunter Sep 6, 2022
a158372
Update _cosmos_client.py
JericHunter Sep 6, 2022
b639f2a
Update cosmos_client.py
JericHunter Sep 6, 2022
c913996
Update database.py
JericHunter Sep 6, 2022
ee5777e
Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py
JericHunter Sep 6, 2022
bfe449f
Update sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py
JericHunter Sep 8, 2022
0433b44
Update sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py
JericHunter Sep 8, 2022
12bbd3a
Update _database.py
JericHunter Sep 8, 2022
4083940
Update _base.py
JericHunter Sep 8, 2022
7446dcb
Update sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
JericHunter Sep 8, 2022
74b9250
Update sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
JericHunter Sep 9, 2022
b186c67
Update offer.py
JericHunter Sep 9, 2022
60fe615
Update database.py
JericHunter Sep 9, 2022
c83dc57
Update cosmos_client.py
JericHunter Sep 9, 2022
f3fd4dd
Update _database.py
JericHunter Sep 9, 2022
4ac759b
Update _base.py
JericHunter Sep 9, 2022
8c1145f
Update sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py
JericHunter Sep 9, 2022
6d5386e
Update sdk/cosmos/azure-cosmos/azure/cosmos/database.py
JericHunter Sep 9, 2022
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion sdk/cosmos/azure-cosmos/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
3 changes: 1 addition & 2 deletions sdk/cosmos/azure-cosmos/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ client = CosmosClient(URL, credential=KEY)
### AAD Authentication

You can also authenticate a client utilizing your service principal's AAD credentials and the azure identity package.
You can directly pass in the credentials information to ClientSecretCrednetial, or use the DefaultAzureCredential:
You can directly pass in the credentials information to ClientSecretCredential, or use the DefaultAzureCredential:
```Python
from azure.cosmos import CosmosClient
from azure.identity import ClientSecretCredential, DefaultAzureCredential
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -62,5 +63,6 @@
"TriggerOperation",
"TriggerType",
"ConnectionRetryPolicy",
"ThroughputProperties",
)
__version__ = VERSION
65 changes: 62 additions & 3 deletions sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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]
)
5 changes: 2 additions & 3 deletions sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
9 changes: 5 additions & 4 deletions sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down
15 changes: 7 additions & 8 deletions sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
5 changes: 2 additions & 3 deletions sdk/cosmos/azure-cosmos/azure/cosmos/container.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
16 changes: 9 additions & 7 deletions sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down
15 changes: 7 additions & 8 deletions sdk/cosmos/azure-cosmos/azure/cosmos/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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
)
Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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):
Expand Down
1 change: 1 addition & 0 deletions sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ class HttpHeaders(object):
# Offer type.
OfferType = "x-ms-offer-type"
OfferThroughput = "x-ms-offer-throughput"
AutoscaleSettings = "x-ms-cosmos-offer-autopilot-settings"

# Custom RUs/minute headers
DisableRUPerMinuteUsage = "x-ms-documentdb-disable-ru-per-minute-usage"
Expand Down
Loading