diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 5092478e483f..125caf0f5106 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,12 +3,18 @@ ### 4.7.1 (Unreleased) #### Features Added +* Added response headers directly to SDK item point operation responses. See [PR 35791](https://github.com/Azure/azure-sdk-for-python/pull/35791). * SDK will now retry all ServiceRequestErrors (failing outgoing requests) before failing. Default number of retries is 3. See [PR 36514](https://github.com/Azure/azure-sdk-for-python/pull/36514). * Added Retry Policy for Container Recreate in the Python SDK. See [PR 36043](https://github.com/Azure/azure-sdk-for-python/pull/36043) * Added option to disable write payload on writes. See [PR 37365](https://github.com/Azure/azure-sdk-for-python/pull/37365) * Added get feed ranges API. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) * Added feed range support in `query_items_change_feed`. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) +#### Breaking Changes +* Item-level point operations will now return `CosmosDict` and `CosmosList` response types. +Responses will still be able to be used directly as previously, but will now have access to their response headers without need for a response hook. See [PR 35791](https://github.com/Azure/azure-sdk-for-python/pull/35791). +For more information on this, see our README section [here](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/cosmos/azure-cosmos/README.md#using-item-operation-response-headers). + #### Bugs Fixed * Consolidated Container Properties Cache to be in the Client to cache partition key definition and container rid to avoid unnecessary container reads. See [PR 35731](https://github.com/Azure/azure-sdk-for-python/pull/35731). * Fixed bug with client hangs when running into WriteForbidden exceptions. See [PR 36514](https://github.com/Azure/azure-sdk-for-python/pull/36514). diff --git a/sdk/cosmos/azure-cosmos/README.md b/sdk/cosmos/azure-cosmos/README.md index f2ae9340bdcb..ff99c23af61f 100644 --- a/sdk/cosmos/azure-cosmos/README.md +++ b/sdk/cosmos/azure-cosmos/README.md @@ -427,13 +427,13 @@ client = CosmosClient(URL, credential=KEY) # Database DATABASE_NAME = 'testDatabase' database = client.get_database_client(DATABASE_NAME) -db_offer = database.read_offer() +db_offer = database.get_throughput() print('Found Offer \'{0}\' for Database \'{1}\' and its throughput is \'{2}\''.format(db_offer.properties['id'], database.id, db_offer.properties['content']['offerThroughput'])) # Container with dedicated throughput only. Will return error "offer not found" for containers without dedicated throughput CONTAINER_NAME = 'testContainer' container = database.get_container_client(CONTAINER_NAME) -container_offer = container.read_offer() +container_offer = container.get_throughput() print('Found Offer \'{0}\' for Container \'{1}\' and its throughput is \'{2}\''.format(container_offer.properties['id'], container.id, container_offer.properties['content']['offerThroughput'])) ``` @@ -467,6 +467,28 @@ print(json.dumps(container_props['defaultTtl'])) For more information on TTL, see [Time to Live for Azure Cosmos DB data][cosmos_ttl]. +### Using item point operation response headers + +Response headers include metadata information from the executed operations like `etag`, which allows for optimistic concurrency scenarios, or `x-ms-request-charge` which lets you know how many RUs were consumed by the request. +This applies to all item point operations in both the sync and async clients - and can be used by referencing the `get_response_headers()` method of any response as such: +```python +from azure.cosmos import CosmosClient +import os + +URL = os.environ['ACCOUNT_URI'] +KEY = os.environ['ACCOUNT_KEY'] +DATABASE_NAME = 'testDatabase' +CONTAINER_NAME = 'products' +client = CosmosClient(URL, credential=KEY) +database = client.get_database_client(DATABASE_NAME) +container = database.get_container_client(CONTAINER_NAME) + +operation_response = container.create_item({"id": "test_item", "productName": "test_item"}) +operation_headers = operation_response.get_response_headers() +etag_value = operation_headers['etag'] +request_charge = operation_headers['x-ms-request-charge'] +``` + ### Using the asynchronous client The asynchronous cosmos client is a separate client that looks and works in a similar fashion to the existing synchronous client. However, the async client needs to be imported separately and its methods need to be used with the async/await keywords. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py b/sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py index b1e3d8bf2a30..4f737f0c9450 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py @@ -20,6 +20,7 @@ # SOFTWARE. from ._version import VERSION +from ._cosmos_responses import CosmosDict, CosmosList from ._retry_utility import ConnectionRetryPolicy from .container import ContainerProxy from .cosmos_client import CosmosClient @@ -65,6 +66,8 @@ "TriggerType", "ConnectionRetryPolicy", "ThroughputProperties", + "CosmosDict", + "CosmosList", "FeedRange" ) __version__ = VERSION diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index aa0241d7f289..2b52eb4872ad 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -42,6 +42,7 @@ DistributedTracingPolicy, ProxyPolicy ) +from azure.core.utils import CaseInsensitiveDict from azure.core.pipeline.transport import HttpRequest, \ HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import @@ -60,6 +61,7 @@ from ._change_feed.change_feed_state import ChangeFeedState from ._constants import _Constants as Constants from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy +from ._cosmos_responses import CosmosDict, CosmosList from ._range_partition_resolver import RangePartitionResolver from ._request_object import RequestObject from ._retry_utility import ConnectionRetryPolicy @@ -157,7 +159,7 @@ def __init__( } # Keeps the latest response headers from the server. - self.last_response_headers: Dict[str, Any] = {} + self.last_response_headers: CaseInsensitiveDict = CaseInsensitiveDict() self.UseMultipleWriteLocations = False self._global_endpoint_manager = global_endpoint_manager._GlobalEndpointManager(self) @@ -436,7 +438,7 @@ def QueryDatabases( if options is None: options = {} - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( "/dbs", "dbs", "", lambda r: r["Databases"], lambda _, b: b, query, options, **kwargs) @@ -494,7 +496,7 @@ def QueryContainers( path = base.GetPathFromLink(database_link, "colls") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "colls", database_id, lambda r: r["DocumentCollections"], lambda _, body: body, query, options, **kwargs) @@ -723,7 +725,7 @@ def QueryUsers( path = base.GetPathFromLink(database_link, "users") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "users", database_id, lambda r: r["Users"], lambda _, b: b, query, options, **kwargs) @@ -901,7 +903,7 @@ def QueryPermissions( path = base.GetPathFromLink(user_link, "permissions") user_id = base.GetResourceIdOrFullNameFromLink(user_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "permissions", user_id, lambda r: r["Permissions"], lambda _, b: b, query, options, **kwargs) @@ -1084,7 +1086,7 @@ def QueryItems( path = base.GetPathFromLink(database_or_container_link, "docs") collection_id = base.GetResourceIdOrFullNameFromLink(database_or_container_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "docs", @@ -1173,7 +1175,7 @@ def _QueryChangeFeed( path = base.GetPathFromLink(collection_link, resource_key) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: if collection_link in self.__container_properties_cache: new_options = dict(options) new_options["containerRID"] = self.__container_properties_cache[collection_link]["_rid"] @@ -1244,7 +1246,7 @@ def _QueryPartitionKeyRanges( path = base.GetPathFromLink(collection_link, "pkranges") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], lambda _, b: b, query, options, **kwargs) @@ -1259,7 +1261,7 @@ def CreateItem( document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Creates a document in a collection. :param str database_or_container_link: @@ -1267,7 +1269,7 @@ def CreateItem( :param dict document: The Azure Cosmos document to create. :param dict options: The request options for the request. :return: The created Document. - :rtype: dict + :rtype: CosmosDict """ # Python's default arguments are evaluated once when the function is defined, # not each time the function is called (like it is in say, Ruby). This means @@ -1295,7 +1297,7 @@ def UpsertItem( document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Upserts a document in a collection. :param str database_or_container_link: @@ -1303,7 +1305,7 @@ def UpsertItem( :param dict document: The Azure Cosmos document to upsert. :param dict options: The request options for the request. :return: The upserted Document. - :rtype: dict + :rtype: CosmosDict """ # Python's default arguments are evaluated once when the function is defined, # not each time the function is called (like it is in say, Ruby). This means @@ -1370,7 +1372,7 @@ def ReadItem( document_link: str, options: Optional[Mapping[str, Any]] = None, **kwargs - ) -> Dict[str, Any]: + ) -> CosmosDict: """Reads a document. :param str document_link: @@ -1381,7 +1383,7 @@ def ReadItem( :return: The read Document. :rtype: - dict + CosmosDict """ if options is None: @@ -1442,7 +1444,7 @@ def QueryTriggers( path = base.GetPathFromLink(collection_link, "triggers") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "triggers", collection_id, lambda r: r["Triggers"], lambda _, b: b, query, options, **kwargs) @@ -1595,7 +1597,7 @@ def QueryUserDefinedFunctions( path = base.GetPathFromLink(collection_link, "udfs") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], lambda _, b: b, query, options, **kwargs) @@ -1751,7 +1753,7 @@ def QueryStoredProcedures( path = base.GetPathFromLink(collection_link, "sprocs") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "sprocs", collection_id, lambda r: r["StoredProcedures"], lambda _, b: b, query, options, **kwargs) @@ -1905,7 +1907,7 @@ def QueryConflicts( path = base.GetPathFromLink(collection_link, "conflicts") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( path, "conflicts", collection_id, lambda r: r["Conflicts"], lambda _, b: b, query, options, **kwargs) @@ -1971,7 +1973,7 @@ def ReplaceItem( new_document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces a document and returns it. :param str document_link: @@ -1983,7 +1985,7 @@ def ReplaceItem( :return: The new Document. :rtype: - dict + CosmosDict """ base._validate_resource(new_document) @@ -2012,7 +2014,7 @@ def PatchItem( operations: List[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Patches a document and returns it. :param str document_link: The link to the document. @@ -2022,7 +2024,7 @@ def PatchItem( :return: The new Document. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -2046,7 +2048,7 @@ def PatchItem( self._UpdateSessionIfRequired(headers, result, last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) def Batch( self, @@ -2054,7 +2056,7 @@ def Batch( batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> List[Dict[str, Any]]: + ) -> CosmosList: """Executes the given operations in transactional batch. :param str collection_link: The link to the collection @@ -2064,7 +2066,7 @@ def Batch( :return: The result of the batch operation. :rtype: - list + CosmosList """ response_hook = kwargs.pop("response_hook", None) @@ -2108,7 +2110,7 @@ def Batch( ) if response_hook: response_hook(last_response_headers, final_responses) - return final_responses + return CosmosList(final_responses, response_headers=last_response_headers) def _Batch( self, @@ -2117,13 +2119,13 @@ def _Batch( collection_id: Optional[str], options: Mapping[str, Any], **kwargs: Any - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) headers = base.GetHeaders(self, initial_headers, "post", path, collection_id, "docs", options) request_params = RequestObject("docs", documents._OperationType.Batch) return cast( - Tuple[List[Dict[str, Any]], Dict[str, Any]], + Tuple[List[Dict[str, Any]], CaseInsensitiveDict], self.__Post(path, request_params, batch_operations, headers, **kwargs) ) @@ -2520,7 +2522,7 @@ def QueryOffers( if options is None: options = {} - def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return self.__QueryFeed( "/offers", "offers", "", lambda r: r["Offers"], lambda _, b: b, query, options, **kwargs) @@ -2587,7 +2589,7 @@ def Create( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Creates an Azure Cosmos resource and returns it. :param dict body: @@ -2601,7 +2603,7 @@ def Create( :return: The created Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop('response_hook', None) @@ -2620,7 +2622,7 @@ def Create( self._UpdateSessionIfRequired(headers, result, last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) def Upsert( self, @@ -2631,7 +2633,7 @@ def Upsert( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Upserts an Azure Cosmos resource and returns it. :param dict body: @@ -2645,7 +2647,7 @@ def Upsert( :return: The upserted Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop('response_hook', None) @@ -2664,7 +2666,7 @@ def Upsert( self._UpdateSessionIfRequired(headers, result, last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) def Replace( self, @@ -2675,7 +2677,7 @@ def Replace( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces an Azure Cosmos resource and returns it. :param dict resource: @@ -2689,7 +2691,7 @@ def Replace( :return: The new Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop('response_hook', None) @@ -2707,7 +2709,7 @@ def Replace( self._UpdateSessionIfRequired(headers, result, self.last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) def Read( self, @@ -2717,7 +2719,7 @@ def Read( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Reads an Azure Cosmos resource and returns it. :param str path: @@ -2728,9 +2730,9 @@ def Read( The request options for the request. :return: - The upserted Azure Cosmos resource. + The requested Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop('response_hook', None) @@ -2745,7 +2747,7 @@ def Read( self.last_response_headers = last_response_headers if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) def DeleteResource( self, @@ -2793,7 +2795,7 @@ def __Get( request_params: RequestObject, req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'GET' http request. :param str path: the url to be used for the request. @@ -2821,7 +2823,7 @@ def __Post( body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'POST' http request. :param str path: the url to be used for the request. @@ -2850,7 +2852,7 @@ def __Put( body: Dict[str, Any], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'PUT' http request. :param str path: the url to be used for the request. @@ -2879,7 +2881,7 @@ def __Patch( request_data: Dict[str, Any], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'PATCH' http request. :param str path: the url to be used for the request. @@ -2907,7 +2909,7 @@ def __Delete( request_params: RequestObject, req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[None, Dict[str, Any]]: + ) -> Tuple[None, CaseInsensitiveDict]: """Azure Cosmos 'DELETE' http request. :param str path: the url to be used for the request. @@ -2936,7 +2938,7 @@ def QueryFeed( options: Mapping[str, Any], partition_key_range_id: Optional[str] = None, **kwargs: Any - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: """Query Feed for Document Collection resource. :param str path: Path to the document collection. @@ -2971,7 +2973,7 @@ def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, is_query_plan: bool = False, **kwargs: Any - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: """Query for more than one Azure Cosmos resources. :param str path: @@ -3038,7 +3040,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: initial_headers[http_constants.HttpHeaders.IsQuery] = "true" if not is_query_plan: - initial_headers[http_constants.HttpHeaders.IsQuery] = "true" + initial_headers[http_constants.HttpHeaders.IsQuery] = "true" # TODO: check why we have this weird logic if (self._query_compatibility_mode in (CosmosClientConnection._QueryCompatibilityMode.Default, CosmosClientConnection._QueryCompatibilityMode.Query)): @@ -3064,6 +3066,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: # check if query has prefix partition key isPrefixPartitionQuery = kwargs.pop("isPrefixPartitionQuery", None) if isPrefixPartitionQuery: + last_response_headers = CaseInsensitiveDict() # here get the over lapping ranges partition_key_definition = kwargs.pop("partitionKeyDefinition", None) pk_properties = partition_key_definition diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_responses.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_responses.py new file mode 100644 index 000000000000..96130b53ab97 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_responses.py @@ -0,0 +1,38 @@ +# The MIT License (MIT) +# Copyright (c) 2024 Microsoft Corporation + +from typing import Any, Dict, Iterable, Mapping, Optional, List +from azure.core.utils import CaseInsensitiveDict + + +class CosmosDict(Dict[str, Any]): + def __init__(self, original_dict: Optional[Mapping[str, Any]], /, *, response_headers: CaseInsensitiveDict) -> None: + if original_dict is None: + original_dict = {} + super().__init__(original_dict) + self._response_headers = response_headers + + def get_response_headers(self) -> CaseInsensitiveDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: ~azure.core.CaseInsensitiveDict + """ + return self._response_headers.copy() + + +class CosmosList(List[Dict[str, Any]]): + def __init__(self, original_list: Optional[Iterable[Dict[str, Any]]], /, *, + response_headers: CaseInsensitiveDict) -> None: + if original_list is None: + original_list = [] + super().__init__(original_list) + self._response_headers = response_headers + + def get_response_headers(self) -> CaseInsensitiveDict: + """Returns a copy of the response headers associated to this response + + :return: Dict of response headers + :rtype: ~azure.core.CaseInsensitiveDict + """ + return self._response_headers.copy() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 3f879ded3187..59cf276f0381 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -32,6 +32,7 @@ from azure.core.tracing.decorator_async import distributed_trace_async # type: ignore from ._cosmos_client_connection_async import CosmosClientConnection +from .._cosmos_responses import CosmosDict, CosmosList from ._scripts import ScriptsProxy from .._base import ( build_options as _build_options, @@ -204,7 +205,7 @@ async def create_item( priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Create an item in the container. To update or replace an existing item, use the @@ -228,12 +229,12 @@ async def create_item( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. - :rtype: Dict[str, Any] + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. + :returns: A CosmosDict representing the new item. The dict will be empty if `no_response` is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -261,7 +262,7 @@ async def create_item( result = await self.client_connection.CreateItem( database_or_container_link=self.container_link, document=body, options=request_options, **kwargs ) - return result or {} + return result @distributed_trace_async async def read_item( @@ -275,7 +276,7 @@ async def read_item( max_integrated_cache_staleness_in_ms: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Get the item identified by `item`. :param item: The ID (name) or dict representing item to retrieve. @@ -294,8 +295,8 @@ async def read_item( request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item couldn't be retrieved. - :returns: Dict representing the item to be retrieved. - :rtype: Dict[str, Any] + :returns: A CosmosDict representing the retrieved item. + :rtype: ~azure.cosmos.CosmosDict[str, Any] .. admonition:: Example: @@ -726,7 +727,7 @@ async def upsert_item( priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Insert or update the specified item. If the item already exists in the container, it is replaced. If the item @@ -746,13 +747,13 @@ async def upsert_item( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the upsert operation went through. The dict will be empty if + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. + :returns: A CosmosDict representing the upserted item. The dict will be empty if `no_response` is specified. - :rtype: Dict[str, Any] + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -781,7 +782,7 @@ async def upsert_item( options=request_options, **kwargs ) - return result or {} + return result @distributed_trace_async async def replace_item( @@ -798,7 +799,7 @@ async def replace_item( priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces the specified item if it exists in the container. If the item does not already exist in the container, an exception is raised. @@ -819,14 +820,14 @@ async def replace_item( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace failed or the item with - given id does not exist. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace operation failed or the item with + given id does not exist. + :returns: A CosmosDict representing the item after replace went through. The dict will be empty if `no_response` is specified. - :rtype: Dict[str, Any] + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ item_link = self._get_document_link(item) if pre_trigger_include is not None: @@ -853,7 +854,7 @@ async def replace_item( result = await self.client_connection.ReplaceItem( document_link=item_link, new_document=body, options=request_options, **kwargs ) - return result or {} + return result @distributed_trace_async async def patch_item( @@ -871,7 +872,7 @@ async def patch_item( priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """ Patches the specified item with the provided operations if it exists in the container. @@ -895,13 +896,13 @@ async def patch_item( request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the patch operation went through. The dict will be empty if - `no_response` is specified. + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. - :rtype: dict[str, Any] + :returns: A CosmosDict representing the item after the patch operations went through. The dict will be empty if + `no_response` is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -928,7 +929,7 @@ async def patch_item( item_link = self._get_document_link(item) result = await self.client_connection.PatchItem( document_link=item_link, operations=patch_operations, options=request_options, **kwargs) - return result or {} + return result @distributed_trace_async async def delete_item( @@ -1250,7 +1251,7 @@ async def execute_item_batch( match_condition: Optional[MatchConditions] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> List[Dict[str, Any]]: + ) -> CosmosList: """ Executes the transactional batch for the specified partition key. :param batch_operations: The batch of operations to be executed. @@ -1267,10 +1268,10 @@ async def execute_item_batch( request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A list representing the items after the batch operations went through. + :returns: A CosmosList representing the items after the batch operations went through. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The batch failed to execute. :raises ~azure.cosmos.exceptions.CosmosBatchOperationError: A transactional batch operation failed in the batch. - :rtype: List[Dict[str, Any]] + :rtype: ~azure.cosmos.CosmosList[Dict[str, Any]] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index eeb67225660a..b942c37e92a6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -46,6 +46,7 @@ CustomHookPolicy, DistributedTracingPolicy, ProxyPolicy) +from azure.core.utils import CaseInsensitiveDict from .. import _base as base from .._base import _set_properties_cache @@ -55,6 +56,7 @@ from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants +from .._cosmos_responses import CosmosDict, CosmosList from .. import http_constants, exceptions from . import _query_iterable_async as query_iterable from .. import _runtime_constants as runtime_constants @@ -77,6 +79,8 @@ PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long + + class CredentialDict(TypedDict, total=False): masterKey: str resourceTokens: Mapping[str, Any] @@ -161,7 +165,7 @@ def __init__( self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level # Keeps the latest response headers from the server. - self.last_response_headers: Dict[str, Any] = {} + self.last_response_headers: CaseInsensitiveDict = CaseInsensitiveDict() self.UseMultipleWriteLocations = False self._global_endpoint_manager = global_endpoint_manager_async._GlobalEndpointManager(self) @@ -516,7 +520,7 @@ async def CreateItem( document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Creates a document in a collection. :param str database_or_container_link: @@ -528,7 +532,7 @@ async def CreateItem( :return: The created Document. :rtype: - dict + CosmosDict """ # Python's default arguments are evaluated once when the function is defined, @@ -705,7 +709,7 @@ async def Create( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Creates an Azure Cosmos resource and returns it. :param dict body: @@ -718,7 +722,7 @@ async def Create( :return: The created Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -737,7 +741,7 @@ async def Create( self._UpdateSessionIfRequired(headers, result, last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, response_headers=last_response_headers) async def UpsertUser( self, @@ -797,7 +801,7 @@ async def UpsertItem( document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Upserts a document in a collection. :param str database_or_container_link: @@ -809,7 +813,7 @@ async def UpsertItem( :return: The upserted Document. :rtype: - dict + CosmosDict """ # Python's default arguments are evaluated once when the function is defined, @@ -841,7 +845,7 @@ async def Upsert( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Upserts an Azure Cosmos resource and returns it. :param dict body: @@ -854,7 +858,7 @@ async def Upsert( :return: The upserted Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -874,7 +878,8 @@ async def Upsert( self._UpdateSessionIfRequired(headers, result, self.last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, + response_headers=last_response_headers) async def __Post( self, @@ -883,7 +888,7 @@ async def __Post( body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'POST' async http request. :param str path: the url to be used for the request. @@ -960,7 +965,7 @@ async def ReadItem( document_link: str, options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Reads a document. :param str document_link: @@ -971,7 +976,7 @@ async def ReadItem( :return: The read Document. :rtype: - dict + CosmosDict """ if options is None: @@ -1143,7 +1148,7 @@ async def Read( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Reads an Azure Cosmos resource and returns it. :param str path: @@ -1154,9 +1159,9 @@ async def Read( The request options for the request. :return: - The upserted Azure Cosmos resource. + The retrieved Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -1171,7 +1176,8 @@ async def Read( self.last_response_headers = last_response_headers if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, + response_headers=last_response_headers) async def __Get( self, @@ -1179,7 +1185,7 @@ async def __Get( request_params: _request_object.RequestObject, req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'GET' async http request. :param str path: the url to be used for the request. @@ -1359,7 +1365,7 @@ async def ReplaceItem( new_document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces a document and returns it. :param str document_link: @@ -1370,7 +1376,7 @@ async def ReplaceItem( :return: The new Document. :rtype: - dict + CosmosDict """ base._validate_resource(new_document) @@ -1399,7 +1405,7 @@ async def PatchItem( operations: List[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Patches a document and returns it. :param str document_link: The link to the document. @@ -1408,7 +1414,7 @@ async def PatchItem( :return: The new Document. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -1434,7 +1440,8 @@ async def PatchItem( self._UpdateSessionIfRequired(headers, result, last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, + response_headers=last_response_headers) async def ReplaceOffer( self, @@ -1501,7 +1508,7 @@ async def Replace( initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces an Azure Cosmos resource and returns it. :param dict resource: @@ -1514,7 +1521,7 @@ async def Replace( :return: The new Azure Cosmos resource. :rtype: - dict + CosmosDict """ response_hook = kwargs.pop("response_hook", None) @@ -1532,7 +1539,8 @@ async def Replace( self._UpdateSessionIfRequired(headers, result, self.last_response_headers) if response_hook: response_hook(last_response_headers, result) - return result + return CosmosDict(result, + response_headers=last_response_headers) async def __Put( self, @@ -1541,7 +1549,7 @@ async def __Put( body: Dict[str, Any], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'PUT' async http request. :param str path: the url to be used for the request. @@ -1570,7 +1578,7 @@ async def __Patch( request_data: Dict[str, Any], req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + ) -> Tuple[Dict[str, Any], CaseInsensitiveDict]: """Azure Cosmos 'PATCH' http request. :param str path: the url to be used for the request. @@ -1860,7 +1868,7 @@ async def __Delete( request_params: _request_object.RequestObject, req_headers: Dict[str, Any], **kwargs: Any - ) -> Tuple[None, Dict[str, Any]]: + ) -> Tuple[None, CaseInsensitiveDict]: """Azure Cosmos 'DELETE' async http request. :param str path: the url to be used for the request. @@ -1887,7 +1895,7 @@ async def Batch( batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any - ) -> List[Dict[str, Any]]: + ) -> CosmosList: """Executes the given operations in transactional batch. :param str collection_link: The link to the collection @@ -1897,7 +1905,7 @@ async def Batch( :return: The result of the batch operation. :rtype: - list + CosmosList """ response_hook = kwargs.pop("response_hook", None) @@ -1943,7 +1951,8 @@ async def Batch( ) if response_hook: response_hook(last_response_headers, final_responses) - return final_responses + return CosmosList(final_responses, + response_headers=last_response_headers) async def _Batch( self, @@ -1952,13 +1961,13 @@ async def _Batch( collection_id: Optional[str], options: Mapping[str, Any], **kwargs: Any - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) headers = base.GetHeaders(self, initial_headers, "post", path, collection_id, "docs", options) request_params = _request_object.RequestObject("docs", documents._OperationType.Batch) result = await self.__Post(path, request_params, batch_operations, headers, **kwargs) - return cast(Tuple[List[Dict[str, Any]], Dict[str, Any]], result) + return cast(Tuple[List[Dict[str, Any]], CaseInsensitiveDict], result) def _ReadPartitionKeyRanges( self, @@ -2008,7 +2017,7 @@ def _QueryPartitionKeyRanges( path = base.GetPathFromLink(collection_link, "pkranges") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], @@ -2060,7 +2069,7 @@ def QueryDatabases( if options is None: options = {} - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( "/dbs", "dbs", "", lambda r: r["Databases"], @@ -2120,7 +2129,7 @@ def QueryContainers( path = base.GetPathFromLink(database_link, "colls") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "colls", database_id, lambda r: r["DocumentCollections"], @@ -2197,7 +2206,7 @@ def QueryItems( path = base.GetPathFromLink(database_or_container_link, "docs") collection_id = base.GetResourceIdOrFullNameFromLink(database_or_container_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, @@ -2288,7 +2297,7 @@ def _QueryChangeFeed( path = base.GetPathFromLink(collection_link, resource_key) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: if collection_link in self.__container_properties_cache: new_options = dict(options) new_options["containerRID"] = self.__container_properties_cache[collection_link]["_rid"] @@ -2337,7 +2346,7 @@ def QueryOffers( if options is None: options = {} - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( "/offers", "offers", "", lambda r: r["Offers"], lambda _, b: b, query, options, **kwargs @@ -2402,7 +2411,7 @@ def QueryUsers( path = base.GetPathFromLink(database_link, "users") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "users", database_id, lambda r: r["Users"], @@ -2464,7 +2473,7 @@ def QueryPermissions( path = base.GetPathFromLink(user_link, "permissions") user_id = base.GetResourceIdOrFullNameFromLink(user_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "permissions", user_id, lambda r: r["Permissions"], lambda _, b: b, query, options, **kwargs @@ -2525,7 +2534,7 @@ def QueryStoredProcedures( path = base.GetPathFromLink(collection_link, "sprocs") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "sprocs", collection_id, lambda r: r["StoredProcedures"], @@ -2587,7 +2596,7 @@ def QueryTriggers( path = base.GetPathFromLink(collection_link, "triggers") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "triggers", collection_id, lambda r: r["Triggers"], lambda _, b: b, query, options, **kwargs @@ -2648,7 +2657,7 @@ def QueryUserDefinedFunctions( path = base.GetPathFromLink(collection_link, "udfs") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], @@ -2709,7 +2718,7 @@ def QueryConflicts( path = base.GetPathFromLink(collection_link, "conflicts") collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: return ( await self.__QueryFeed( path, "conflicts", collection_id, lambda r: r["Conflicts"], @@ -2730,7 +2739,7 @@ async def QueryFeed( options: Mapping[str, Any], partition_key_range_id: Optional[str] = None, **kwargs: Any - ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + ) -> Tuple[List[Dict[str, Any]], CaseInsensitiveDict]: """Query Feed for Document Collection resource. :param str path: Path to the document collection. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_query_iterable_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_query_iterable_async.py index a2d213d1514a..9c14a7fd1058 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_query_iterable_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_query_iterable_async.py @@ -82,7 +82,7 @@ async def _unpack(self, block): continuation = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get("x-ms-continuation") or \ - self._client.last_response_headers.get('etag') + self._client.last_response_headers.get('etag') if block: self._did_a_call_already = False return continuation, block diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index e602aca419b9..ed9a211b0091 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -39,6 +39,7 @@ _set_properties_cache ) from ._cosmos_client_connection import CosmosClientConnection +from ._cosmos_responses import CosmosDict, CosmosList from ._feed_range import FeedRange, FeedRangeEpk from ._routing.routing_range import Range from .offer import Offer, ThroughputProperties @@ -204,7 +205,7 @@ def read_item( # pylint:disable=docstring-missing-param max_integrated_cache_staleness_in_ms: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Get the item identified by `item`. :param item: The ID (name) or dict representing item to retrieve. @@ -221,9 +222,9 @@ def read_item( # pylint:disable=docstring-missing-param :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :returns: Dict representing the item to be retrieved. + :returns: A CosmosDict representing the item to be retrieved. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item couldn't be retrieved. - :rtype: dict[str, Any] + :rtype: ~azure.cosmos.CosmosDict[str, Any] .. admonition:: Example: @@ -693,7 +694,7 @@ def replace_item( # pylint:disable=docstring-missing-param priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Replaces the specified item if it exists in the container. If the item does not already exist in the container, an exception is raised. @@ -714,13 +715,13 @@ def replace_item( # pylint:disable=docstring-missing-param request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` - is specified. + sending response payloads. When not specified explicitly here, the default value will be determined from + kwargs or when also not specified there from client-level kwargs. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace operation failed or the item with given id does not exist. - :rtype: Dict[str, Any] + :returns: A CosmosDict representing the item after replace went through. The dict will be empty if `no_response` + is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ item_link = self._get_document_link(item) if pre_trigger_include is not None: @@ -750,10 +751,9 @@ def replace_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - result = self.client_connection.ReplaceItem( document_link=item_link, new_document=body, options=request_options, **kwargs) - return result or {} + return result @distributed_trace def upsert_item( # pylint:disable=docstring-missing-param @@ -770,7 +770,7 @@ def upsert_item( # pylint:disable=docstring-missing-param priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Insert or update the specified item. If the item already exists in the container, it is replaced. If the item @@ -789,12 +789,12 @@ def upsert_item( # pylint:disable=docstring-missing-param :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :keyword bool no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or - when also not specified there from client-level kwargs. - :returns: A dict representing the upserted item. The dict will be empty if `no_response` is specified. + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or + when also not specified there from client-level kwargs. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. - :rtype: Dict[str, Any] + :returns: A CosmosDict representing the upserted item. The dict will be empty if `no_response` is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -829,7 +829,7 @@ def upsert_item( # pylint:disable=docstring-missing-param options=request_options, **kwargs ) - return result or {} + return result @distributed_trace def create_item( # pylint:disable=docstring-missing-param @@ -848,7 +848,7 @@ def create_item( # pylint:disable=docstring-missing-param priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """Create an item in the container. To update or replace an existing item, use the @@ -871,12 +871,12 @@ def create_item( # pylint:disable=docstring-missing-param :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :keyword bool no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. - :rtype: Dict[str, Any] + :returns: A CosmosDict representing the new item. The dict will be empty if `no_response` is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -908,7 +908,7 @@ def create_item( # pylint:disable=docstring-missing-param request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.CreateItem( database_or_container_link=self.container_link, document=body, options=request_options, **kwargs) - return result or {} + return result @distributed_trace def patch_item( @@ -926,7 +926,7 @@ def patch_item( priority: Optional[Literal["High", "Low"]] = None, no_response: Optional[bool] = None, **kwargs: Any - ) -> Dict[str, Any]: + ) -> CosmosDict: """ Patches the specified item with the provided operations if it exists in the container. @@ -949,14 +949,14 @@ def patch_item( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :keyword bool no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after the patch operations went through. The dict will be empty - if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. - :rtype: Dict[str, Any] + :returns: A CosmosDict representing the item after the patch operations went through. The dict will be empty + if `no_response` is specified. + :rtype: ~azure.cosmos.CosmosDict[str, Any] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include @@ -983,7 +983,7 @@ def patch_item( item_link = self._get_document_link(item) result = self.client_connection.PatchItem( document_link=item_link, operations=patch_operations, options=request_options, **kwargs) - return result or {} + return result @distributed_trace def execute_item_batch( @@ -998,7 +998,7 @@ def execute_item_batch( match_condition: Optional[MatchConditions] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> List[Dict[str, Any]]: + ) -> CosmosList: """ Executes the transactional batch for the specified partition key. :param batch_operations: The batch of operations to be executed. @@ -1015,10 +1015,10 @@ def execute_item_batch( request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword Callable response_hook: A callable invoked with the response metadata. - :returns: A list representing the item after the batch operations went through. + :returns: A CosmosList representing the items after the batch operations went through. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The batch failed to execute. :raises ~azure.cosmos.exceptions.CosmosBatchOperationError: A transactional batch operation failed in the batch. - :rtype: List[Dict[str, Any]] + :rtype: ~azure.cosmos.CosmosList[Dict[str, Any]] """ if pre_trigger_include is not None: kwargs['pre_trigger_include'] = pre_trigger_include diff --git a/sdk/cosmos/azure-cosmos/samples/document_management.py b/sdk/cosmos/azure-cosmos/samples/document_management.py index 17a48978ed46..aa019ffbea8f 100644 --- a/sdk/cosmos/azure-cosmos/samples/document_management.py +++ b/sdk/cosmos/azure-cosmos/samples/document_management.py @@ -205,7 +205,7 @@ def execute_item_batch(database): # We create three items to use for the sample. container.create_item(get_sales_order("read_item")) container.create_item(get_sales_order("delete_item")) - container.create_item(get_sales_order("replace_item")) + create_response = container.create_item(get_sales_order("replace_item")) # We create our batch operations create_item_operation = ("create", (get_sales_order("create_item"),)) @@ -217,12 +217,11 @@ def execute_item_batch(database): replace_item_if_match_operation = ("replace", ("replace_item", {"id": "replace_item", 'account_number': 'Account1', "message": "item was replaced"}), - {"if_match_etag": container.client_connection.last_response_headers.get("etag")}) + {"if_match_etag": create_response.get_response_headers().get("etag")}) replace_item_if_none_match_operation = ("replace", ("replace_item", {"id": "replace_item", 'account_number': 'Account1', "message": "item was replaced"}), - {"if_none_match_etag": - container.client_connection.last_response_headers.get("etag")}) + {"if_none_match_etag": create_response.get_response_headers().get("etag")}) # Put our operations into a list batch_operations = [ diff --git a/sdk/cosmos/azure-cosmos/samples/document_management_async.py b/sdk/cosmos/azure-cosmos/samples/document_management_async.py index 00b547768c01..9f98ca7f47e2 100644 --- a/sdk/cosmos/azure-cosmos/samples/document_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/document_management_async.py @@ -224,7 +224,7 @@ async def execute_item_batch(database): # We create three items to use for the sample. await container.create_item(get_sales_order("read_item")) await container.create_item(get_sales_order("delete_item")) - await container.create_item(get_sales_order("replace_item")) + create_response = await container.create_item(get_sales_order("replace_item")) # We create our batch operations create_item_operation = ("create", (get_sales_order("create_item"),)) @@ -236,12 +236,11 @@ async def execute_item_batch(database): replace_item_if_match_operation = ("replace", ("replace_item", {"id": "replace_item", 'account_number': 'Account1', "message": "item was replaced"}), - {"if_match_etag": container.client_connection.last_response_headers.get("etag")}) + {"if_match_etag": create_response.get_response_headers().get("etag")}) replace_item_if_none_match_operation = ("replace", ("replace_item", {"id": "replace_item", 'account_number': 'Account1', "message": "item was replaced"}), - {"if_none_match_etag": - container.client_connection.last_response_headers.get("etag")}) + {"if_none_match_etag": create_response.get_response_headers().get("etag")}) # Put our operations into a list batch_operations = [ diff --git a/sdk/cosmos/azure-cosmos/test/test_cosmos_responses.py b/sdk/cosmos/azure-cosmos/test/test_cosmos_responses.py new file mode 100644 index 000000000000..b0cf2c920830 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_cosmos_responses.py @@ -0,0 +1,74 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest + +import test_config +from azure.cosmos import CosmosClient, PartitionKey, DatabaseProxy +from azure.cosmos.http_constants import HttpHeaders + + +# TODO: add more tests in this file once we have response headers in control plane operations +# TODO: add query tests once those changes are available + +@pytest.mark.cosmosEmulator +class TestCosmosResponses(unittest.TestCase): + """Python Cosmos Responses Tests. + """ + + configs = test_config.TestConfig + host = configs.host + masterKey = configs.masterKey + client: CosmosClient = None + test_database: DatabaseProxy = None + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + if (cls.masterKey == '[YOUR_KEY_HERE]' or + cls.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + cls.client = CosmosClient(cls.host, cls.masterKey) + cls.test_database = cls.client.get_database_client(cls.TEST_DATABASE_ID) + + def test_point_operation_headers(self): + container = self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + first_response = container.upsert_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + lsn = first_response.get_response_headers()['lsn'] + + create_response = container.create_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + assert len(create_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(create_response.get_response_headers()['lsn']) + lsn = create_response.get_response_headers()['lsn'] + + read_response = container.read_item(create_response['id'], create_response['company']) + assert len(read_response.get_response_headers()) > 0 + + upsert_response = container.upsert_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + assert len(upsert_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(upsert_response.get_response_headers()['lsn']) + lsn = upsert_response.get_response_headers()['lsn'] + + upsert_response['replace'] = True + replace_response = container.replace_item(upsert_response['id'], upsert_response) + assert replace_response['replace'] is not None + assert len(replace_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(replace_response.get_response_headers()['lsn']) + + batch = [] + for i in range(50): + batch.append(("create", ({"id": "item" + str(i), "company": "Microsoft"},))) + batch_response = container.execute_item_batch(batch_operations=batch, partition_key="Microsoft") + assert len(batch_response.get_response_headers()) > 0 + assert int(lsn) + 1 < int(batch_response.get_response_headers()['lsn']) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_cosmos_responses_async.py b/sdk/cosmos/azure-cosmos/test/test_cosmos_responses_async.py new file mode 100644 index 000000000000..10e077945a18 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_cosmos_responses_async.py @@ -0,0 +1,78 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest + +import test_config +from azure.cosmos import PartitionKey +from azure.cosmos.aio import CosmosClient +from azure.cosmos.http_constants import HttpHeaders + + +# TODO: add more tests in this file once we have response headers in control plane operations +# TODO: add query tests once those changes are available + +@pytest.mark.cosmosEmulator +class TestCosmosResponsesAsync(unittest.IsolatedAsyncioTestCase): + """Python Cosmos Responses Tests. + """ + + configs = test_config.TestConfig + host = configs.host + masterKey = configs.masterKey + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + if (cls.masterKey == '[YOUR_KEY_HERE]' or + cls.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + + async def asyncSetUp(self): + self.client = CosmosClient(self.host, self.masterKey) + self.test_database = self.client.get_database_client(self.TEST_DATABASE_ID) + + async def asyncTearDown(self): + await self.client.close() + + async def test_point_operation_headers_async(self): + container = await self.test_database.create_container(id="responses_test" + str(uuid.uuid4()), + partition_key=PartitionKey(path="/company")) + first_response = await container.upsert_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + lsn = first_response.get_response_headers()['lsn'] + + create_response = await container.create_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + assert len(create_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(create_response.get_response_headers()['lsn']) + lsn = create_response.get_response_headers()['lsn'] + + read_response = await container.read_item(create_response['id'], create_response['company']) + assert len(read_response.get_response_headers()) > 0 + + upsert_response = await container.upsert_item({"id": str(uuid.uuid4()), "company": "Microsoft"}) + assert len(upsert_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(upsert_response.get_response_headers()['lsn']) + lsn = upsert_response.get_response_headers()['lsn'] + + upsert_response['replace'] = True + replace_response = await container.replace_item(upsert_response['id'], upsert_response) + assert replace_response['replace'] is not None + assert len(replace_response.get_response_headers()) > 0 + assert int(lsn) + 1 == int(replace_response.get_response_headers()['lsn']) + + batch = [] + for i in range(50): + batch.append(("create", ({"id": "item" + str(i), "company": "Microsoft"},))) + batch_response = await container.execute_item_batch(batch_operations=batch, partition_key="Microsoft") + assert len(batch_response.get_response_headers()) > 0 + assert int(lsn) + 1 < int(batch_response.get_response_headers()['lsn']) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_globaldb.py b/sdk/cosmos/azure-cosmos/test/test_globaldb.py index 412ab5c3e5b2..36901d817137 100644 --- a/sdk/cosmos/azure-cosmos/test/test_globaldb.py +++ b/sdk/cosmos/azure-cosmos/test/test_globaldb.py @@ -123,15 +123,12 @@ def test_global_db_read_write_endpoints(self): # Delay to get these resources replicated to read location due to Eventual consistency time.sleep(5) - self.test_coll.read_item(item=created_document, partition_key=created_document['pk']) - content_location = str(client.client_connection.last_response_headers[HttpHeaders.ContentLocation]) + read_response = self.test_coll.read_item(item=created_document, partition_key=created_document['pk']) + content_location = str(read_response.get_response_headers()[HttpHeaders.ContentLocation]) - content_location_url = urlparse(content_location) - host_url = urlparse(TestGlobalDB.host) - - # When EnableEndpointDiscovery is False, ReadEndpoint is set to the endpoint passed while creating the client instance - self.assertEqual(str(content_location_url.hostname), str(host_url.hostname)) - self.assertEqual(client.client_connection.ReadEndpoint, TestGlobalDB.host) + # When EnableEndpointDiscovery is False, ReadEndpoint is set to the endpoint passed while creating the client + # instance + assert client.client_connection.ReadEndpoint == TestGlobalDB.host connection_policy.EnableEndpointDiscovery = True document_definition['id'] = 'doc2' @@ -150,16 +147,16 @@ def test_global_db_read_write_endpoints(self): # Delay to get these resources replicated to read location due to Eventual consistency time.sleep(5) - container.read_item(item=created_document, partition_key=created_document['pk']) - content_location = str(client.client_connection.last_response_headers[HttpHeaders.ContentLocation]) + read_response = container.read_item(item=created_document, partition_key=created_document['pk']) + content_location = str(read_response.get_response_headers()[HttpHeaders.ContentLocation]) content_location_url = urlparse(content_location) write_location_url = urlparse(TestGlobalDB.write_location_host) # If no preferred locations is set, we return the write endpoint as ReadEndpoint for better latency performance if is_not_default_host(TestGlobalDB.write_location_host): - self.assertEqual(str(content_location_url.hostname), str(write_location_url.hostname)) - self.assertEqual(client.client_connection.ReadEndpoint, TestGlobalDB.write_location_host) + assert str(content_location_url.hostname) == str(write_location_url.hostname) + assert client.client_connection.ReadEndpoint == TestGlobalDB.write_location_host def test_global_db_endpoint_discovery(self): connection_policy = documents.ConnectionPolicy() @@ -223,8 +220,8 @@ def test_global_db_preferred_locations(self): # Delay to get these resources replicated to read location due to Eventual consistency time.sleep(5) - item = container.read_item(item=created_document, partition_key=created_document['pk']) - content_location = str(client.client_connection.last_response_headers[HttpHeaders.ContentLocation]) + read_response = container.read_item(item=created_document, partition_key=created_document['pk']) + content_location = str(read_response.get_response_headers()[HttpHeaders.ContentLocation]) content_location_url = urlparse(content_location) write_location_url = urlparse(self.write_location_host) @@ -249,8 +246,8 @@ def test_global_db_preferred_locations(self): # Delay to get these resources replicated to read location due to Eventual consistency time.sleep(5) - container.read_item(item=created_document, partition_key=created_document['pk']) - content_location = str(client.client_connection.last_response_headers[HttpHeaders.ContentLocation]) + read_response = container.read_item(item=created_document, partition_key=created_document['pk']) + content_location = str(read_response.get_response_headers()[HttpHeaders.ContentLocation]) content_location_url = urlparse(content_location) read_location2_url = urlparse(self.read_location2_host) diff --git a/sdk/cosmos/azure-cosmos/test/test_streaming_failover.py b/sdk/cosmos/azure-cosmos/test/test_streaming_failover.py index 09624c938578..07e5b40fa8c1 100644 --- a/sdk/cosmos/azure-cosmos/test/test_streaming_failover.py +++ b/sdk/cosmos/azure-cosmos/test/test_streaming_failover.py @@ -67,7 +67,7 @@ def test_streaming_fail_over(self): created_document = created_container.create_item(document_definition) self.assertDictEqual(created_document, {}) - self.assertDictEqual(client.client_connection.last_response_headers, {}) + self.assertDictEqual(created_document.get_response_headers(), {}) self.assertEqual(self.counter, 10) # First request is an initial read collection. diff --git a/sdk/cosmos/azure-cosmos/test/test_transactional_batch.py b/sdk/cosmos/azure-cosmos/test/test_transactional_batch.py index 20d074ee6be8..33bdbcd7a5f6 100644 --- a/sdk/cosmos/azure-cosmos/test/test_transactional_batch.py +++ b/sdk/cosmos/azure-cosmos/test/test_transactional_batch.py @@ -340,8 +340,8 @@ def test_batch_lsn(self): container.upsert_item({"id": "patch_item", "company": "Microsoft"}) container.upsert_item({"id": "delete_item", "company": "Microsoft"}) - container.read_item(item="read_item", partition_key="Microsoft") - lsn = container.client_connection.last_response_headers.get(HttpHeaders.LSN) + read_response = container.read_item(item="read_item", partition_key="Microsoft") + lsn = read_response.get_response_headers().get(HttpHeaders.LSN) batch = [("create", ({"id": "create_item", "company": "Microsoft"},)), ("replace", ("replace_item", {"id": "replace_item", "company": "Microsoft", "value": True})), @@ -352,7 +352,7 @@ def test_batch_lsn(self): batch_response = container.execute_item_batch(batch_operations=batch, partition_key="Microsoft") assert len(batch_response) == 6 - assert int(lsn) == int(container.client_connection.last_response_headers.get(HttpHeaders.LSN)) - 1 + assert int(lsn) == int(batch_response.get_response_headers().get(HttpHeaders.LSN)) - 1 self.test_database.delete_container(container.id) diff --git a/sdk/cosmos/azure-cosmos/test/test_transactional_batch_async.py b/sdk/cosmos/azure-cosmos/test/test_transactional_batch_async.py index e0c49b18e519..ad26b766627b 100644 --- a/sdk/cosmos/azure-cosmos/test/test_transactional_batch_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_transactional_batch_async.py @@ -348,8 +348,8 @@ async def test_batch_lsn_async(self): await container.upsert_item({"id": "patch_item", "company": "Microsoft"}) await container.upsert_item({"id": "delete_item", "company": "Microsoft"}) - await container.read_item(item="read_item", partition_key="Microsoft") - lsn = container.client_connection.last_response_headers.get(HttpHeaders.LSN) + read_response = await container.read_item(item="read_item", partition_key="Microsoft") + lsn = read_response.get_response_headers().get(HttpHeaders.LSN) batch = [("create", ({"id": "create_item", "company": "Microsoft"},)), ("replace", ("replace_item", {"id": "replace_item", "company": "Microsoft", "value": True})), @@ -360,7 +360,7 @@ async def test_batch_lsn_async(self): batch_response = await container.execute_item_batch(batch_operations=batch, partition_key="Microsoft") assert len(batch_response) == 6 - assert int(lsn) == int(container.client_connection.last_response_headers.get(HttpHeaders.LSN)) - 1 + assert int(lsn) == int(batch_response.get_response_headers().get(HttpHeaders.LSN)) - 1 await self.test_database.delete_container(container.id) diff --git a/sdk/cosmos/azure-cosmos/test/test_user_configs.py b/sdk/cosmos/azure-cosmos/test/test_user_configs.py index b6bd2ff6b33e..850d39517c24 100644 --- a/sdk/cosmos/azure-cosmos/test/test_user_configs.py +++ b/sdk/cosmos/azure-cosmos/test/test_user_configs.py @@ -64,17 +64,17 @@ def test_default_account_consistency(self): # Testing the session token logic works without user passing in Session explicitly database = client.create_database(database_id) container = database.create_container(id=container_id, partition_key=PartitionKey(path="/id")) - container.create_item(body=get_test_item()) - session_token = client.client_connection.last_response_headers[http_constants.CookieHeaders.SessionToken] + create_response = container.create_item(body=get_test_item()) + session_token = create_response.get_response_headers()[http_constants.CookieHeaders.SessionToken] item2 = get_test_item() - container.create_item(body=item2) - session_token2 = client.client_connection.last_response_headers[http_constants.CookieHeaders.SessionToken] + create_response = container.create_item(body=item2) + session_token2 = create_response.get_response_headers()[http_constants.CookieHeaders.SessionToken] # Check Session token is being updated to reflect new item created self.assertNotEqual(session_token, session_token2) - container.read_item(item=item2.get("id"), partition_key=item2.get("id")) - read_session_token = client.client_connection.last_response_headers[http_constants.CookieHeaders.SessionToken] + read_response = container.read_item(item=item2.get("id"), partition_key=item2.get("id")) + read_session_token = read_response.get_response_headers()[http_constants.CookieHeaders.SessionToken] # Check Session token remains the same for read operation as with previous create item operation self.assertEqual(session_token2, read_session_token)