Added Response Headers to Control Plane Operations (#41742)

* added response headers to control plane operations

* test cleanup

* corrected response headers and added return type of cosmosDict

* added response headers to throughputProperties class

* edited getting response headers for thread safety

* added overload methods and tests for control plane operations

* changed return_type to return_header boolean to simplify code

* added headers for replace throughput

* added create container and database if not exists negative handling and tests

* changed return_headers to return_properties

* added changes from git comments

* returns tuple for previous container and db proxy operations, added samples

* matching sync and async positional parameters

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

Co-authored-by: Tomas Varon <[email protected]>

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

Co-authored-by: Tomas Varon <[email protected]>

* added changes from pr comments

* bug fix for replace_container

* fixed pylint errors

* fixed positional arg errors

* fixed pylint errors

* fixed pos args ttl test

* changed back positional parameters to avoid breaking changes

* fixed return typing errors

* bugfixes from mypy pipeline

* testing changes for mypy errors

* testing type change for mypy errors

* testing fix for mypy errors

* database fixes for mypy errors

* testing mypy/pylint fixes for overloads

* adding docstring changes for overloaded functions

* pylint fixes

* fixing pylint errors for param types

* testing changes for return value error in mypy and pylint fixes

* reverting return type back to dict[str, Any] for get_client_container_caches

* testing variable changes for mypy

* testing changes for mypy

* testing changes for mypy back to original dict

* got rid of extra library import

* fixed indent issue

* reverted back to original type of Dict for get_properties()

* minor fixes

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

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

* made suggested changes

* added suggestions for sync and async methods

* added to docstring for overloaded methods

* changed wording for return_properties in docstring

* added suggested docstring changes

* wording change for database proxy return type

* minor bugfix for docstring

* added to changelog

---------

Co-authored-by: Andrew Mathew <[email protected]>
Co-authored-by: Tomas Varon <[email protected]>
Co-authored-by: Anna Tisch <[email protected]>

Author: 4 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,7 +3,7 @@
 ### 4.14.0b5 (Unreleased)
 
 #### Features Added
-
+* Added ability to return a tuple of a DatabaseProxy/ContainerProxy with the associated database/container properties when creating or reading databases/containers through `return_properties` parameter. See [PR 41742](https://github.com/Azure/azure-sdk-for-python/pull/41742)
 #### Breaking Changes
 
 #### Bugs Fixed

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -370,7 +370,7 @@ def CreateDatabase(
         database: Dict[str, Any],
         options: Optional[Mapping[str, Any]] = None,
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Creates a database.
 
         :param dict database:
@@ -393,7 +393,7 @@ def ReadDatabase(
         database_link: str,
         options: Optional[Mapping[str, Any]] = None,
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Reads a database.
 
         :param str database_link:
@@ -527,7 +527,7 @@ def CreateContainer(
         collection: Dict[str, Any],
         options: Optional[Mapping[str, Any]] = None,
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Creates a collection in a database.
 
         :param str database_link:
@@ -556,7 +556,7 @@ def ReplaceContainer(
         collection: Dict[str, Any],
         options: Optional[Mapping[str, Any]] = None,
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Replaces a collection and return it.
 
         :param str collection_link:
@@ -586,7 +586,7 @@ def ReadContainer(
         collection_link: str,
         options: Optional[Mapping[str, Any]] = None,
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Reads a collection.
 
         :param str collection_link:
@@ -2554,7 +2554,7 @@ def ReplaceOffer(
         offer_link: str,
         offer: Dict[str, Any],
         **kwargs: Any
-    ) -> Dict[str, Any]:
+    ) -> CosmosDict:
         """Replaces an offer and returns it.
 
         :param str offer_link:

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

@@ -172,7 +172,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

@@ -22,7 +22,7 @@
 """Create, read, and delete databases in the Azure Cosmos DB SQL API service.
 """
 
-from typing import Any, Dict, List, Optional, Union, cast, Mapping, Iterable, Callable
+from typing import Any, Dict, List, Optional, Union, cast, Mapping, Iterable, Callable, overload, Literal
 import warnings
 from azure.core.async_paging import AsyncItemPaged
 from azure.core.credentials import TokenCredential
@@ -41,6 +41,7 @@
 from ._database import DatabaseProxy, _get_database_link
 from ..documents import ConnectionPolicy, DatabaseAccount
 from ..exceptions import CosmosResourceNotFoundError
+from .._cosmos_responses import CosmosDict
 
 # pylint: disable=docstring-keyword-should-match-keyword-only
 
@@ -255,14 +256,15 @@ def from_connection_string(
             **kwargs
         )
 
-    @distributed_trace_async
+    @overload
     async def create_database(
         self,
         id: str,
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
         throughput_bucket: Optional[int] = None,
+        return_properties: Literal[False] = False,
         **kwargs: Any
     ) -> DatabaseProxy:
         """
@@ -275,8 +277,10 @@ async def create_database(
         :keyword response_hook: A callable invoked with the response metadata.
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
         :raises ~azure.cosmos.exceptions.CosmosResourceExistsError: Database with the given ID already exists.
-        :returns: A DatabaseProxy instance representing the new database.
+        :returns: A `DatabaseProxy` instance representing the database.
         :rtype: ~azure.cosmos.aio.DatabaseProxy
 
         .. admonition:: Example:
@@ -289,6 +293,84 @@ async def create_database(
                 :caption: Create a database in the Cosmos DB account:
                 :name: create_database
         """
+        ...
+
+    @overload
+    async def create_database(
+        self,
+        id: str,
+        *,
+        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
+        initial_headers: Optional[Dict[str, str]] = None,
+        throughput_bucket: Optional[int] = None,
+        return_properties: Literal[True],
+        **kwargs: Any
+    ) -> tuple[DatabaseProxy, CosmosDict]:
+        """
+        Create a new database with the given ID (name).
+
+        :param str id: ID (name) of the database to create.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: Union[int, ~azure.cosmos.ThroughputProperties]
+        :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
+        :keyword response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
+        :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
+        :raises ~azure.cosmos.exceptions.CosmosResourceExistsError: Database with the given ID already exists.
+        :returns: A tuple of `DatabaseProxy` and CosmosDict with the database properties.
+        :rtype: tuple [~azure.cosmos.aio.DatabaseProxy, ~azure.cosmos.CosmosDict]
+
+        .. admonition:: Example:
+
+            .. literalinclude:: ../samples/examples_async.py
+                :start-after: [START create_database]
+                :end-before: [END create_database]
+                :language: python
+                :dedent: 0
+                :caption: Create a database in the Cosmos DB account:
+                :name: create_database
+        """
+        ...
+
+    @distributed_trace_async
+    async def create_database( # pylint:disable=docstring-should-be-keyword
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> Union[DatabaseProxy, tuple[DatabaseProxy, CosmosDict]]:
+        """
+        Create a new database with the given ID (name).
+
+        :param Any args: args
+        :param str id: ID (name) of the database to read or create.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: Union[int, ~azure.cosmos.ThroughputProperties]
+        :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
+        :keyword response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
+        :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
+        :raises ~azure.cosmos.exceptions.CosmosResourceExistsError: Database with the given ID already exists.
+        :returns: A DatabaseProxy instance representing the database or a tuple of DatabaseProxy
+            and CosmosDict with the database properties.
+        :rtype: ~azure.cosmos.aio.DatabaseProxy or tuple [~azure.cosmos.aio.DatabaseProxy, ~azure.cosmos.CosmosDict]
+
+        .. admonition:: Example:
+
+            .. literalinclude:: ../samples/examples_async.py
+                :start-after: [START create_database]
+                :end-before: [END create_database]
+                :language: python
+                :dedent: 0
+                :caption: Create a database in the Cosmos DB account:
+                :name: create_database
+        """
+        id = args[0] if args else kwargs.pop("id")
+        if len(args) > 1:
+            raise TypeError(f"Unexpected positional parameters: {args[1:]}")
         session_token = kwargs.get('session_token')
         if session_token is not None:
             warnings.warn(
@@ -307,24 +389,27 @@ async def create_database(
                 "The 'match_condition' flag does not apply to this method and is always ignored even if passed."
                 " It will now be removed in the future.",
                 DeprecationWarning)
-        if initial_headers is not None:
-            kwargs["initial_headers"] = initial_headers
-        if throughput_bucket is not None:
-            kwargs["throughput_bucket"] = throughput_bucket
+
+        offer_throughput = kwargs.pop("offer_throughput", None)
+        return_properties = kwargs.pop("return_properties", False)
+
         request_options = _build_options(kwargs)
         _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)
+        if not return_properties:
+            return DatabaseProxy(self.client_connection, id=result["id"], properties=result)
+        return  DatabaseProxy(self.client_connection, id=result["id"], properties=result), result
 
-    @distributed_trace_async
+    @overload
     async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
         self,
         id: str,
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
         throughput_bucket: Optional[int] = None,
+        return_properties: Literal[False] = False,
         **kwargs: Any
     ) -> DatabaseProxy:
         """
@@ -343,10 +428,84 @@ async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
         :keyword response_hook: A callable invoked with the response metadata.
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The database read or creation failed.
         :returns: A DatabaseProxy instance representing the database.
-        :rtype: ~azure.cosmos.DatabaseProxy
+        :rtype: ~azure.cosmos.aio.DatabaseProxy
+        """
+        ...
+
+    @overload
+    async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
+        self,
+        id: str,
+        *,
+        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
+        initial_headers: Optional[Dict[str, str]] = None,
+        throughput_bucket: Optional[int] = None,
+        return_properties: Literal[True],
+        **kwargs: Any
+    ) -> tuple[DatabaseProxy, CosmosDict]:
+        """
+        Create the database if it does not exist already.
+
+        If the database already exists, the existing settings are returned.
+
+        ..note::
+            This function does not check or update existing database settings or
+            offer throughput if they differ from what is passed in.
+
+        :param str id: ID (name) of the database to read or create.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: Union[int, ~azure.cosmos.ThroughputProperties]
+        :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
+        :keyword response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
+        :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
+        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The database read or creation failed.
+        :returns: A tuple of DatabaseProxy and CosmosDict with the database properties.
+        :rtype: tuple [~azure.cosmos.aio.DatabaseProxy, ~azure.cosmos.CosmosDict]
+        """
+        ...
+
+    @distributed_trace_async
+    async def create_database_if_not_exists( # pylint:disable=docstring-should-be-keyword
+        self,
+        *args: Any,
+        **kwargs: Any
+    ) -> Union[DatabaseProxy, tuple[DatabaseProxy, CosmosDict]]:
         """
+        Create the database if it does not exist already.
+
+        If the database already exists, the existing settings are returned.
+
+        ..note::
+            This function does not check or update existing database settings or
+            offer throughput if they differ from what is passed in.
+
+        :param Any args: args
+        :param str id: ID (name) of the database to read or create.
+        :keyword offer_throughput: The provisioned throughput for this offer.
+        :paramtype offer_throughput: Union[int, ~azure.cosmos.ThroughputProperties]
+        :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.
+        :keyword response_hook: A callable invoked with the response metadata.
+        :keyword int throughput_bucket: The desired throughput bucket for the client
+        :paramtype response_hook: Callable[[Dict[str, str], Dict[str, Any]], None]
+        :keyword bool return_properties: Specifies whether to return either a DatabaseProxy
+            or a Tuple containing a DatabaseProxy and the associated database properties.
+        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The database read or creation failed.
+        :returns: A DatabaseProxy instance representing the database or a tuple of DatabaseProxy
+            and CosmosDict with the database properties.
+        :rtype: ~azure.cosmos.aio.DatabaseProxy or tuple [~azure.cosmos.aio.DatabaseProxy, ~azure.cosmos.CosmosDict]
+        """
+
+        id = args[0] if args else kwargs.pop("id")
+        if len(args) > 1:
+            raise TypeError(f"Unexpected positional parameters: {args[1:]}")
+
         session_token = kwargs.get('session_token')
         if session_token is not None:
             warnings.warn(
@@ -365,18 +524,21 @@ async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
                 "The 'match_condition' flag does not apply to this method and is always ignored even if passed."
                 " It will now be removed in the future.",
                 DeprecationWarning)
-        if throughput_bucket is not None:
-            kwargs["throughput_bucket"] = throughput_bucket
-        if initial_headers is not None:
-            kwargs["initial_headers"] = initial_headers
+
+        offer_throughput = kwargs.pop("offer_throughput", None)
+        return_properties = kwargs.pop("return_properties", False)
+
         try:
             database_proxy = self.get_database_client(id)
-            await database_proxy.read(**kwargs)
-            return database_proxy
+            result = await database_proxy.read(**kwargs)
+            if not return_properties:
+                return database_proxy
+            return database_proxy, result
         except CosmosResourceNotFoundError:
             return await self.create_database(
                 id,
                 offer_throughput=offer_throughput,
+                return_properties=return_properties,
                 **kwargs
             )