From 9368d1f620156be41f1f55d32d0b59a8f6d96bf2 Mon Sep 17 00:00:00 2001 From: Andrew Mathew Date: Tue, 24 Jun 2025 15:59:31 -0400 Subject: [PATCH 1/5] added response headers to control plane operations --- .../azure/cosmos/aio/_container.py | 9 ++++++ .../azure/cosmos/aio/_database.py | 9 ++++++ .../azure-cosmos/azure/cosmos/container.py | 11 ++++++++ .../azure-cosmos/azure/cosmos/database.py | 9 ++++++ .../tests/test_cosmos_responses.py | 24 ++++++++++++++++ .../tests/test_cosmos_responses_async.py | 28 +++++++++++++++++++ 6 files changed, 90 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 3cac81623f00..743d56a5cdcd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -93,6 +93,7 @@ def __init__( self.container_link = "{}/colls/{}".format(database_link, self.id) self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None + self._response_headers = properties if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -112,6 +113,14 @@ async def _get_properties(self, **kwargs: Any) -> Dict[str, Any]: await self.read(**kwargs) return self.client_connection._container_properties_cache[self.container_link] + def get_response_headers(self) -> dict[str, Any]: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers.copy() + @property async def is_system_key(self) -> bool: if self._is_system_key is None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py index db4e1151c0f9..b2a023535a0d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py @@ -100,6 +100,7 @@ def __init__( self.id = id self.database_link = "dbs/{}".format(self.id) self._properties = properties + self._response_headers = properties def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -126,6 +127,14 @@ async def _get_properties(self) -> Dict[str, Any]: self._properties = await self.read() return self._properties + def get_response_headers(self) -> dict[str, Any]: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers.copy() + @distributed_trace_async async def read( self, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index e498501d5e51..76f821b54d8c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -24,6 +24,8 @@ import warnings from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload, Iterable, Callable + +from requests.structures import CaseInsensitiveDict from typing_extensions import Literal from azure.core import MatchConditions @@ -104,6 +106,7 @@ def __init__( self.client_connection = client_connection self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None + self._response_headers = properties if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -123,6 +126,14 @@ def _get_properties(self, **kwargs: Any) -> Dict[str, Any]: self.read(**kwargs) return self.__get_client_container_caches()[self.container_link] + def get_response_headers(self) -> dict[str, Any]: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers.copy() + @property def is_system_key(self) -> bool: if self._is_system_key is None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py index dba8c02905a0..0f9fe700e385 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py @@ -94,6 +94,7 @@ def __init__( self.id = id self.database_link: str = "dbs/{}".format(self.id) self._properties: Optional[Dict[str, Any]] = properties + self._response_headers = properties def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -120,6 +121,14 @@ def _get_properties(self) -> Dict[str, Any]: self._properties = self.read() return self._properties + def get_response_headers(self) -> dict[str, Any]: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers.copy() + @distributed_trace def read( # pylint:disable=docstring-missing-param self, diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py index b0cf2c920830..323992a54493 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py @@ -69,6 +69,30 @@ def test_point_operation_headers(self): assert len(batch_response.get_response_headers()) > 0 assert int(lsn) + 1 < int(batch_response.get_response_headers()['lsn']) + def test_create_database_headers(self): + first_response = self.client.create_database(id="responses_test" + str(uuid.uuid4())) + assert len(first_response.get_response_headers()) > 0 + + def test_create_database_if_not_exists_headers(self): + first_response = self.client.create_database_if_not_exists(id="responses_test" + str(uuid.uuid4())) + assert len(first_response.get_response_headers()) > 0 + + def test_create_container_headers(self): + first_response = self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + assert len(first_response.get_response_headers()) > 0 + + def test_create_container_if_not_exists_headers(self): + first_response = self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + assert len(first_response.get_response_headers()) > 0 + + def test_replace_container_headers(self): + first_response = self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + second_response = self.test_database.replace_container(first_response.id, + partition_key=PartitionKey(path="/company")) + assert len(second_response.get_response_headers()) > 0 if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py index 10e077945a18..86d7706601af 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py @@ -73,6 +73,34 @@ async def test_point_operation_headers_async(self): assert len(batch_response.get_response_headers()) > 0 assert int(lsn) + 1 < int(batch_response.get_response_headers()['lsn']) + async def test_create_database_headers(self): + first_response = await self.client.create_database(id="responses_test" + str(uuid.uuid4())) + print(first_response.get_response_headers()) + assert len(first_response.get_response_headers()) > 0 + + async def test_create_database_if_not_exists_headers(self): + first_response = await self.client.create_database_if_not_exists(id="responses_test" + str(uuid.uuid4())) + print(first_response.get_response_headers()) + assert len(first_response.get_response_headers()) > 0 + + async def test_create_container_headers(self): + first_response = await self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + print(first_response.get_response_headers()) + assert len(first_response.get_response_headers()) > 0 + + async def test_create_container_if_not_exists_headers(self): + first_response = await self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + print(first_response.get_response_headers()) + assert len(first_response.get_response_headers()) > 0 + + async def test_replace_container_headers(self): + first_response = await self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + second_response = await self.test_database.replace_container(first_response.id, + partition_key=PartitionKey(path="/company")) + assert len(second_response.get_response_headers()) > 0 if __name__ == '__main__': unittest.main() From bd5ab7bba93c8e86343eb7f838e60597bdbbbe5e Mon Sep 17 00:00:00 2001 From: Andrew Mathew Date: Tue, 24 Jun 2025 16:07:43 -0400 Subject: [PATCH 2/5] test cleanup --- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 2 -- sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py | 4 ---- 2 files changed, 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 76f821b54d8c..b24b25202001 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -24,8 +24,6 @@ import warnings from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload, Iterable, Callable - -from requests.structures import CaseInsensitiveDict from typing_extensions import Literal from azure.core import MatchConditions diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py index 86d7706601af..3ec7ba90985f 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py @@ -75,24 +75,20 @@ async def test_point_operation_headers_async(self): async def test_create_database_headers(self): first_response = await self.client.create_database(id="responses_test" + str(uuid.uuid4())) - print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 async def test_create_database_if_not_exists_headers(self): first_response = await self.client.create_database_if_not_exists(id="responses_test" + str(uuid.uuid4())) - print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 async def test_create_container_headers(self): first_response = await self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), partition_key=PartitionKey(path="/company")) - print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 async def test_create_container_if_not_exists_headers(self): first_response = await self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), partition_key=PartitionKey(path="/company")) - print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 async def test_replace_container_headers(self): From 432ebaa995adf4597a3175291be3b4d2c7ffce8a Mon Sep 17 00:00:00 2001 From: Andrew Mathew Date: Thu, 26 Jun 2025 17:28:49 -0400 Subject: [PATCH 3/5] corrected response headers and added return type of cosmosDict --- .../azure-cosmos/azure/cosmos/aio/_container.py | 7 ++++--- .../azure-cosmos/azure/cosmos/aio/_database.py | 7 ++++--- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 7 ++++--- sdk/cosmos/azure-cosmos/azure/cosmos/database.py | 7 ++++--- .../azure-cosmos/tests/test_cosmos_responses.py | 13 +++++++++++++ .../tests/test_cosmos_responses_async.py | 12 ++++++++++++ 6 files changed, 41 insertions(+), 12 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 743d56a5cdcd..dc831f4ac4bd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -47,6 +47,7 @@ from .._routing.routing_range import Range from .._session_token_helpers import get_latest_session_token from ..offer import ThroughputProperties +from .._cosmos_responses import CosmosDict, CosmosList from ..partition_key import ( NonePartitionKeyValue, _return_undefined_or_empty_partition_key, @@ -93,7 +94,7 @@ def __init__( self.container_link = "{}/colls/{}".format(database_link, self.id) self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None - self._response_headers = properties + self._response_headers = self.client_connection.last_response_headers.copy() if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -119,7 +120,7 @@ def get_response_headers(self) -> dict[str, Any]: :return: Dict of response headers :rtype: dict[str, Any] """ - return self._response_headers.copy() + return self._response_headers @property async def is_system_key(self) -> bool: @@ -176,7 +177,7 @@ async def read( priority: Optional[Literal["High", "Low"]] = None, initial_headers: Optional[Dict[str, str]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Read the container properties. :keyword bool populate_partition_key_range_statistics: Enable returning partition key diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py index b2a023535a0d..eb40207ee04c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py @@ -39,6 +39,7 @@ from ._user import UserProxy from ..documents import IndexingMode from ..partition_key import PartitionKey +from .._cosmos_responses import CosmosDict __all__ = ("DatabaseProxy",) @@ -100,7 +101,7 @@ def __init__( self.id = id self.database_link = "dbs/{}".format(self.id) self._properties = properties - self._response_headers = properties + self._response_headers = self.client_connection.last_response_headers.copy() def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -133,7 +134,7 @@ def get_response_headers(self) -> dict[str, Any]: :return: Dict of response headers :rtype: dict[str, Any] """ - return self._response_headers.copy() + return self._response_headers @distributed_trace_async async def read( @@ -141,7 +142,7 @@ async def read( *, initial_headers: Optional[Dict[str, str]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Read the database properties. :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index b24b25202001..23286d82dbf1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -45,6 +45,7 @@ from ._routing.routing_range import Range from ._session_token_helpers import get_latest_session_token from .offer import Offer, ThroughputProperties +from ._cosmos_responses import CosmosDict from .partition_key import ( NonePartitionKeyValue, PartitionKey, @@ -104,7 +105,7 @@ def __init__( self.client_connection = client_connection self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None - self._response_headers = properties + self._response_headers = self.client_connection.last_response_headers.copy() if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -130,7 +131,7 @@ def get_response_headers(self) -> dict[str, Any]: :return: Dict of response headers :rtype: dict[str, Any] """ - return self._response_headers.copy() + return self._response_headers @property def is_system_key(self) -> bool: @@ -179,7 +180,7 @@ def read( # pylint:disable=docstring-missing-param initial_headers: Optional[Dict[str, str]] = None, response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Read the container properties. :param bool populate_partition_key_range_statistics: Enable returning partition key diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py index 0f9fe700e385..92c831ce5a30 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py @@ -37,6 +37,7 @@ from .exceptions import CosmosResourceNotFoundError from .user import UserProxy from .documents import IndexingMode +from ._cosmos_responses import CosmosDict __all__ = ("DatabaseProxy",) @@ -94,7 +95,7 @@ def __init__( self.id = id self.database_link: str = "dbs/{}".format(self.id) self._properties: Optional[Dict[str, Any]] = properties - self._response_headers = properties + self._response_headers = self.client_connection.last_response_headers.copy() def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -127,7 +128,7 @@ def get_response_headers(self) -> dict[str, Any]: :return: Dict of response headers :rtype: dict[str, Any] """ - return self._response_headers.copy() + return self._response_headers @distributed_trace def read( # pylint:disable=docstring-missing-param @@ -136,7 +137,7 @@ def read( # pylint:disable=docstring-missing-param *, initial_headers: Optional[Dict[str, str]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Read the database properties. :keyword dict[str,str] initial_headers: Initial headers to be sent as part of the request. diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py index 323992a54493..1f3f21778af0 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py @@ -85,6 +85,7 @@ def test_create_container_headers(self): def test_create_container_if_not_exists_headers(self): first_response = self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), partition_key=PartitionKey(path="/company")) + print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 def test_replace_container_headers(self): @@ -94,5 +95,17 @@ def test_replace_container_headers(self): partition_key=PartitionKey(path="/company")) assert len(second_response.get_response_headers()) > 0 + def test_database_read_headers(self): + db = self.client.create_database(id="responses_test" + str(uuid.uuid4())) + first_response = db.read() + assert len(first_response.get_response_headers()) > 0 + + def test_container_read_headers(self): + container = self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + first_response = container.read() + assert len(first_response.get_response_headers()) > 0 + + if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py index 3ec7ba90985f..df8a00233414 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses_async.py @@ -98,5 +98,17 @@ async def test_replace_container_headers(self): partition_key=PartitionKey(path="/company")) assert len(second_response.get_response_headers()) > 0 + async def test_database_read_headers(self): + db = await self.client.create_database(id="responses_test" + str(uuid.uuid4())) + first_response = await db.read() + assert len(first_response.get_response_headers()) > 0 + + async def test_container_read_headers(self): + container = await self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + first_response = await container.read() + assert len(first_response.get_response_headers()) > 0 + + if __name__ == '__main__': unittest.main() From a1d3a08a6b5c6562b02381049dc7e2eace04dfeb Mon Sep 17 00:00:00 2001 From: Andrew Mathew Date: Mon, 30 Jun 2025 16:31:18 -0400 Subject: [PATCH 4/5] added response headers to throughputProperties class --- sdk/cosmos/azure-cosmos/azure/cosmos/offer.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py index 3eeaf4de24e6..26de8454167e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py @@ -22,6 +22,7 @@ from typing import Any, Dict, Optional +from ._cosmos_client_connection import CosmosClientConnection class ThroughputProperties: @@ -39,11 +40,19 @@ class ThroughputProperties: the increment percent should be greater than or equal to zero. """ - def __init__(self, *args, **kwargs) -> None: + def __init__(self, client_connection: CosmosClientConnection, *args, **kwargs) -> None: self.offer_throughput: Optional[int] = args[0] if args else kwargs.get('offer_throughput') self.properties: Optional[Dict[str, Any]] = args[1] if len(args) > 1 else kwargs.get('properties') self.auto_scale_max_throughput: Optional[int] = kwargs.get('auto_scale_max_throughput') self.auto_scale_increment_percent: Optional[int] = kwargs.get('auto_scale_increment_percent') + self._response_headers = client_connection.last_response_headers.copy() + def get_response_headers(self) -> dict[str, Any]: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers Offer = ThroughputProperties From 64829b5d410dedab4443212e73022e3e916072e2 Mon Sep 17 00:00:00 2001 From: Andrew Mathew Date: Thu, 3 Jul 2025 13:58:56 -0400 Subject: [PATCH 5/5] edited getting response headers for thread safety --- .../azure-cosmos/azure/cosmos/aio/_container.py | 7 ++++--- .../azure/cosmos/aio/_cosmos_client.py | 2 +- .../azure-cosmos/azure/cosmos/aio/_database.py | 13 ++++++++----- .../azure-cosmos/azure/cosmos/container.py | 7 ++++--- .../azure-cosmos/azure/cosmos/cosmos_client.py | 2 +- sdk/cosmos/azure-cosmos/azure/cosmos/database.py | 13 ++++++++----- sdk/cosmos/azure-cosmos/azure/cosmos/offer.py | 12 +----------- .../azure-cosmos/tests/test_cosmos_responses.py | 16 +++++++++++++++- 8 files changed, 42 insertions(+), 30 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index dc831f4ac4bd..9af3f42377a8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -86,7 +86,8 @@ def __init__( client_connection: CosmosClientConnection, database_link: str, id: str, - properties: Optional[Dict[str, Any]] = None + properties: Optional[Dict[str, Any]] = None, + header: Optional[CosmosDict] = None ) -> None: self.client_connection = client_connection self.id = id @@ -94,7 +95,7 @@ def __init__( self.container_link = "{}/colls/{}".format(database_link, self.id) self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None - self._response_headers = self.client_connection.last_response_headers.copy() + self._response_headers = header if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -114,7 +115,7 @@ async def _get_properties(self, **kwargs: Any) -> Dict[str, Any]: await self.read(**kwargs) return self.client_connection._container_properties_cache[self.container_link] - def get_response_headers(self) -> dict[str, Any]: + def get_response_headers(self) -> CosmosDict: """Returns a copy of the response headers associated to this response :return: Dict of response headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py index 7bdeba650c72..f57ea7f30e25 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py @@ -310,7 +310,7 @@ async def create_database( _set_throughput_options(offer=offer_throughput, request_options=request_options) result = await self.client_connection.CreateDatabase(database={"id": id}, options=request_options, **kwargs) - return DatabaseProxy(self.client_connection, id=result["id"], properties=result) + return DatabaseProxy(self.client_connection, id=result["id"], properties=result, header=result.get_response_headers()) @distributed_trace_async async def create_database_if_not_exists( # pylint: disable=redefined-builtin diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py index eb40207ee04c..390ab2cde718 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py @@ -90,7 +90,8 @@ def __init__( self, client_connection: CosmosClientConnection, id: str, - properties: Optional[Dict[str, Any]] = None + properties: Optional[Dict[str, Any]] = None, + header: Optional[CosmosDict] = None ) -> None: """ :param client_connection: Client from which this database was retrieved. @@ -101,7 +102,7 @@ def __init__( self.id = id self.database_link = "dbs/{}".format(self.id) self._properties = properties - self._response_headers = self.client_connection.last_response_headers.copy() + self._response_headers = header def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -128,7 +129,7 @@ async def _get_properties(self) -> Dict[str, Any]: self._properties = await self.read() return self._properties - def get_response_headers(self) -> dict[str, Any]: + def get_response_headers(self) -> CosmosDict: """Returns a copy of the response headers associated to this response :return: Dict of response headers @@ -295,7 +296,8 @@ async def create_container( data = await self.client_connection.CreateContainer( database_link=self.database_link, collection=definition, options=request_options, **kwargs ) - return ContainerProxy(self.client_connection, self.database_link, data["id"], properties=data) + return ContainerProxy(self.client_connection, self.database_link, data["id"], properties=data, + header=data.get_response_headers()) @distributed_trace_async async def create_container_if_not_exists( @@ -615,7 +617,8 @@ async def replace_container( container_link, collection=parameters, options=request_options, **kwargs ) return ContainerProxy( - self.client_connection, self.database_link, container_properties["id"], properties=container_properties + self.client_connection, self.database_link, container_properties["id"], properties=container_properties, + header=container_properties.get_response_headers() ) @distributed_trace_async diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 23286d82dbf1..c3ac66194fed 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -98,14 +98,15 @@ def __init__( client_connection: CosmosClientConnection, database_link: str, id: str, - properties: Optional[Dict[str, Any]] = None + properties: Optional[Dict[str, Any]] = None, + header: Optional[CosmosDict] = None ) -> None: self.id = id self.container_link = "{}/colls/{}".format(database_link, self.id) self.client_connection = client_connection self._is_system_key: Optional[bool] = None self._scripts: Optional[ScriptsProxy] = None - self._response_headers = self.client_connection.last_response_headers.copy() + self._response_headers = header if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -125,7 +126,7 @@ def _get_properties(self, **kwargs: Any) -> Dict[str, Any]: self.read(**kwargs) return self.__get_client_container_caches()[self.container_link] - def get_response_headers(self) -> dict[str, Any]: + def get_response_headers(self) -> CosmosDict: """Returns a copy of the response headers associated to this response :return: Dict of response headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index 2756101a57f4..875f7d4c6152 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -328,7 +328,7 @@ def create_database( # pylint:disable=docstring-missing-param result = self.client_connection.CreateDatabase(database={"id": id}, options=request_options, **kwargs) if response_hook: response_hook(self.client_connection.last_response_headers) - return DatabaseProxy(self.client_connection, id=result["id"], properties=result) + return DatabaseProxy(self.client_connection, id=result["id"], properties=result, header=result.get_response_headers()) @distributed_trace def create_database_if_not_exists( # pylint:disable=docstring-missing-param diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py index 92c831ce5a30..e8c1b9504d1f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/database.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/database.py @@ -85,7 +85,8 @@ def __init__( self, client_connection: CosmosClientConnection, id: str, - properties: Optional[Dict[str, Any]] = None + properties: Optional[Dict[str, Any]] = None, + header: Optional[CosmosDict] = None ) -> None: """ :param ClientSession client_connection: Client from which this database was retrieved. @@ -95,7 +96,7 @@ def __init__( self.id = id self.database_link: str = "dbs/{}".format(self.id) self._properties: Optional[Dict[str, Any]] = properties - self._response_headers = self.client_connection.last_response_headers.copy() + self._response_headers = header def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -122,7 +123,7 @@ def _get_properties(self) -> Dict[str, Any]: self._properties = self.read() return self._properties - def get_response_headers(self) -> dict[str, Any]: + def get_response_headers(self) -> CosmosDict: """Returns a copy of the response headers associated to this response :return: Dict of response headers @@ -294,7 +295,8 @@ def create_container( # pylint:disable=docstring-missing-param database_link=self.database_link, collection=definition, options=request_options, **kwargs ) - return ContainerProxy(self.client_connection, self.database_link, result["id"], properties=result) + return ContainerProxy(self.client_connection, self.database_link, result["id"], properties=result, + header=result.get_response_headers()) @distributed_trace def create_container_if_not_exists( # pylint:disable=docstring-missing-param @@ -683,7 +685,8 @@ def replace_container( # pylint:disable=docstring-missing-param container_link, collection=parameters, options=request_options, **kwargs) return ContainerProxy( - self.client_connection, self.database_link, container_properties["id"], properties=container_properties) + self.client_connection, self.database_link, container_properties["id"], properties=container_properties, + header=container_properties.get_response_headers()) @distributed_trace def list_users( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py index 26de8454167e..f7b3e30cd319 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py @@ -22,7 +22,6 @@ from typing import Any, Dict, Optional -from ._cosmos_client_connection import CosmosClientConnection class ThroughputProperties: @@ -40,19 +39,10 @@ class ThroughputProperties: the increment percent should be greater than or equal to zero. """ - def __init__(self, client_connection: CosmosClientConnection, *args, **kwargs) -> None: + def __init__(self, *args, **kwargs) -> None: self.offer_throughput: Optional[int] = args[0] if args else kwargs.get('offer_throughput') self.properties: Optional[Dict[str, Any]] = args[1] if len(args) > 1 else kwargs.get('properties') self.auto_scale_max_throughput: Optional[int] = kwargs.get('auto_scale_max_throughput') self.auto_scale_increment_percent: Optional[int] = kwargs.get('auto_scale_increment_percent') - self._response_headers = client_connection.last_response_headers.copy() - - def get_response_headers(self) -> dict[str, Any]: - """Returns a copy of the response headers associated to this response - - :return: Dict of response headers - :rtype: dict[str, Any] - """ - return self._response_headers Offer = ThroughputProperties diff --git a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py index 1f3f21778af0..fc25a6b234ee 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py @@ -85,7 +85,6 @@ def test_create_container_headers(self): def test_create_container_if_not_exists_headers(self): first_response = self.test_database.create_container_if_not_exists(id="responses_test" + str(uuid.uuid4()), partition_key=PartitionKey(path="/company")) - print(first_response.get_response_headers()) assert len(first_response.get_response_headers()) > 0 def test_replace_container_headers(self): @@ -106,6 +105,21 @@ def test_container_read_headers(self): first_response = container.read() assert len(first_response.get_response_headers()) > 0 + @pytest.mark.skip(reason="haven't decided on implementation yet") + def test_container_read_offer(self): + container = self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + + first_response = container.read_offer() + assert len(first_response.get_response_headers()) > 0 + + @pytest.mark.skip(reason="haven't decided on implementation yet") + def test_container_get_throughput(self): + pass + + @pytest.mark.skip(reason="haven't decided on implementation yet") + def test_container_replace_throughput(self): + pass if __name__ == '__main__': unittest.main()