diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 3cac81623f00..9af3f42377a8 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, @@ -85,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 @@ -93,6 +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 = header if properties: self.client_connection._set_container_properties_cache(self.container_link, _build_properties_cache(properties, @@ -112,6 +115,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) -> CosmosDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers + @property async def is_system_key(self) -> bool: if self._is_system_key is None: @@ -167,7 +178,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/_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 db4e1151c0f9..390ab2cde718 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",) @@ -89,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. @@ -100,6 +102,7 @@ def __init__( self.id = id self.database_link = "dbs/{}".format(self.id) self._properties = properties + self._response_headers = header def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -126,13 +129,21 @@ async def _get_properties(self) -> Dict[str, Any]: self._properties = await self.read() return self._properties + def get_response_headers(self) -> CosmosDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers + @distributed_trace_async async def read( self, *, 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. @@ -285,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( @@ -605,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 e498501d5e51..c3ac66194fed 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, @@ -97,13 +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 = header 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) -> CosmosDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers + @property def is_system_key(self) -> bool: if self._is_system_key is None: @@ -170,7 +181,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/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 dba8c02905a0..e8c1b9504d1f 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",) @@ -84,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. @@ -94,6 +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 = header def __repr__(self) -> str: return "".format(self.database_link)[:1024] @@ -120,6 +123,14 @@ def _get_properties(self) -> Dict[str, Any]: self._properties = self.read() return self._properties + def get_response_headers(self) -> CosmosDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: dict[str, Any] + """ + return self._response_headers + @distributed_trace def read( # pylint:disable=docstring-missing-param self, @@ -127,7 +138,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. @@ -284,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 @@ -673,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 3eeaf4de24e6..f7b3e30cd319 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py @@ -45,5 +45,4 @@ def __init__(self, *args, **kwargs) -> None: 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') - 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 b0cf2c920830..fc25a6b234ee 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py +++ b/sdk/cosmos/azure-cosmos/tests/test_cosmos_responses.py @@ -69,6 +69,57 @@ 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 + + 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 + + @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() 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..df8a00233414 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,42 @@ 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())) + 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())) + 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")) + 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")) + 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 + + 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()