Added Response Headers to Control Plane Operations

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,14 +86,16 @@ 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
         self.database_link = database_link
         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

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

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 "<DatabaseProxy [{}]>".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

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

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

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 "<DatabaseProxy [{}]>".format(self.database_link)[:1024]
@@ -120,14 +123,22 @@ 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,
         populate_query_metrics: Optional[bool] = None,
         *,
         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(

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

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()

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()