[Cosmos] Retryable writes (#41272)

* retriable writes implementation

* pylint/mypy

* Update test_backwards_compatibility.py

* Update test_backwards_compatibility_async.py

* Update test_headers_async.py

* Update test_headers.py

* Update test_headers.py

* add tests for retryable writes

retryable writes tests

* add tests for patch item

* adds README with snippets, removes preview from VS

* add retry write at client level tests

* fix mark

* small typos

* handle service response errors

* service response request level test

* client-level check for service response

* Update retryable writes with tests and multi region retry

Adds constants for keyword retry_policy, adds logic for multi region retry for retryable writes on timeout errors, adds test for delete item retry.

* add mwr tests, fix logic

* update tests with replace item

* pylint fixes

* updates

* address comments

* Update _service_response_retry_policy.py

* update tests

Updates the tests and service response retry policy in order for the test to pass

* Update test_retryable_writes_async.py

* Update test_retryable_writes_async.py

* Update test_retryable_writes_async.py

* Update test_retryable_writes_async.py

* Update test_retryable_writes_async.py

---------

Co-authored-by: bambriz <[email protected]>

Author: simorenoh

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.13.0b3 (Unreleased)
 
 #### Features Added
+* Added option to enable automatic retries for write operations. See [PR 41272](https://github.com/Azure/azure-sdk-for-python/pull/41272).
 
 #### Breaking Changes
 * Adds cross region retries when no preferred locations are set. This is only a breaking change for customers using bounded staleness consistency. See [PR 39714](https://github.com/Azure/azure-sdk-for-python/pull/39714)

sdk/cosmos/azure-cosmos/README.md

@@ -668,7 +668,38 @@ as well as containing the list of failed responses for the failed request.
 
 For more information on Transactional Batch, see [Azure Cosmos DB Transactional Batch][cosmos_transactional_batch].
 
-### Public Preview - Vector Embeddings and Vector Indexes
+### Native Retryable Writes
+We have added native retryable writes to the SDK, a feature that can be used by customers who don't mind the
+non-idempotency of these retries and would instead like to ensure that the given operation executed in case of timeouts
+or connectivity issues (status codes 408, 5xx).
+
+This feature can be enabled either at the client level, to retry all write operations under these conditions, 
+or at the per-request level to enable the retries for an individual request.
+
+If enabled at the client level, the one exception to the rule would be patch requests, since the operations can change
+the nature of the overall request - replace, set, and copy for example would be idempotent while add, move or remove would
+not be idempotent unless combined with patch precondition checks. So, for patch we allow opting-in into automatic
+retries only on the request options level.
+
+The snippet below shows how to enable this feature at the client and request level:
+```python
+cosmos_client = CosmosClient(
+    url=URL,
+    credential=KEY,
+    retry_write=True,  # enables native retryable writes at the client level
+)
+
+database = cosmos_client.get_database_client(DATABASE_NAME)
+container = database.get_container_client(CONTAINER_NAME)
+
+container.create_item(
+    item_body,
+    retry_write=True  # enables native retryable writes at the request level
+)
+```
+
+
+### Vector Embeddings and Vector Indexes
 We have added new capabilities to utilize vector embeddings and vector indexing for users to leverage vector
 search utilizing our Cosmos SDK. These two container-level configurations have to be turned on at the account-level
 before you can use them.
@@ -755,7 +786,7 @@ database.create_container(id=container_id, partition_key=PartitionKey(path="/id"
 ```
 ***Note: vector embeddings and vector indexes CANNOT be edited by container replace operations. They are only available directly through creation.***
 
-### Public Preview - Vector Search
+### Vector Search
 
 With the addition of the vector indexing and vector embedding capabilities, the SDK can now perform order by vector search queries.
 These queries specify the VectorDistance to use as a metric within the query text. These must always use a TOP or LIMIT clause within the query though,

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -37,6 +37,7 @@
 from . import documents
 from . import http_constants
 from . import _runtime_constants
+from ._constants import _Constants as Constants
 from .auth import _get_authorization_header
 from .offer import ThroughputProperties
 from .partition_key import _Empty, _Undefined
@@ -62,6 +63,7 @@
     'query_version': 'queryVersion',
     'priority': 'priorityLevel',
     'no_response': 'responsePayloadOnWriteDisabled',
+    'retry_write': Constants.Kwargs.RETRY_WRITE,
     'max_item_count': 'maxItemCount',
     'throughput_bucket': 'throughputBucket',
     'excluded_locations': 'excludedLocations'

sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py

@@ -78,3 +78,9 @@ class _Constants:
         429: "TOO_MANY_REQUESTS",
         449: "RETRY_WITH - Conflicting request to resource has been attempted. Retry to avoid conflicts."
     }
+
+    class Kwargs:
+        """Keyword arguments used in the azure-cosmos package"""
+
+        RETRY_WRITE: Literal["retry_write"] = "retry_write"
+        """Whether to retry write operations if they fail. Used either at client level or request level."""

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

@@ -2073,6 +2073,7 @@ def PatchItem(
                                        documents._OperationType.Patch,
                                        headers)
         request_params.set_excluded_location_from_options(options)
+        request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         request_data = {}
         if options.get("filterPredicate"):
             request_data["condition"] = options.get("filterPredicate")
@@ -2692,6 +2693,7 @@ def Create(
 
         request_params = RequestObject(typ, documents._OperationType.Create, headers)
         request_params.set_excluded_location_from_options(options)
+        request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         result, last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
         self.last_response_headers = last_response_headers
 
@@ -2739,6 +2741,7 @@ def Upsert(
         # Upsert will use WriteEndpoint since it uses POST operation
         request_params = RequestObject(typ, documents._OperationType.Upsert, headers)
         request_params.set_excluded_location_from_options(options)
+        request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         result, last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
         self.last_response_headers = last_response_headers
         # update session for write request
@@ -2783,6 +2786,7 @@ def Replace(
         # Replace will use WriteEndpoint since it uses PUT operation
         request_params = RequestObject(typ, documents._OperationType.Replace, headers)
         request_params.set_excluded_location_from_options(options)
+        request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         result, last_response_headers = self.__Put(path, request_params, resource, headers, **kwargs)
         self.last_response_headers = last_response_headers
 
@@ -2865,6 +2869,7 @@ def DeleteResource(
         # Delete will use WriteEndpoint since it uses DELETE operation
         request_params = RequestObject(typ, documents._OperationType.Delete, headers)
         request_params.set_excluded_location_from_options(options)
+        request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         result, last_response_headers = self.__Delete(path, request_params, headers, **kwargs)
         self.last_response_headers = last_response_headers
 

sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py

@@ -22,6 +22,8 @@
 """Represents a request object.
 """
 from typing import Optional, Mapping, Any, Dict, List
+from .documents import _OperationType
+from ._constants import _Constants as Constants
 
 class RequestObject(object): # pylint: disable=too-many-instance-attributes
     def __init__(
@@ -43,6 +45,7 @@ def __init__(
         self.excluded_locations: Optional[List[str]] = None
         self.excluded_locations_circuit_breaker: List[str] = []
         self.healthy_tentative_location: Optional[str] = None
+        self.retry_write: bool = False
 
     def route_to_location_with_preferred_location_flag(  # pylint: disable=name-too-long
         self,
@@ -80,5 +83,13 @@ def set_excluded_location_from_options(self, options: Mapping[str, Any]) -> None
         if self._can_set_excluded_location(options):
             self.excluded_locations = options['excludedLocations']
 
+    def set_retry_write(self, request_options: Mapping[str, Any], client_retry_write: bool) -> None:
+        if request_options and request_options.get(Constants.Kwargs.RETRY_WRITE):
+            # If request retry write is True, set the option
+            self.retry_write = request_options[Constants.Kwargs.RETRY_WRITE]
+        elif client_retry_write and self.operation_type != _OperationType.Patch:
+            # If it is not a patch operation and the client config is set, set the retry write to True
+            self.retry_write = client_retry_write
+
     def set_excluded_locations_from_circuit_breaker(self, excluded_locations: List[str]) -> None: # pylint: disable=name-too-long
         self.excluded_locations_circuit_breaker = excluded_locations

sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py

@@ -263,7 +263,7 @@ def _handle_service_request_retries(
         raise exception
 
 def _handle_service_response_retries(request, client, response_retry_policy, exception, *args):
-    if request and _has_read_retryable_headers(request.headers):
+    if request and (_has_read_retryable_headers(request.headers) or (args and is_write_retryable(args[0], client))):
         # we resolve the request endpoint to the next preferred region
         # once we are out of preferred regions we stop retrying
         retry_policy = response_retry_policy
@@ -274,6 +274,11 @@ def _handle_service_response_retries(request, client, response_retry_policy, exc
     else:
         raise exception
 
+def is_write_retryable(request_params, client):
+    return (request_params.retry_write or
+            client.connection_policy.RetryNonIdempotentWrites and
+            not request_params.operation_type == _OperationType.Patch)
+
 def _configure_timeout(request: PipelineRequest, absolute: Optional[int], per_request: int) -> None:
     if absolute is not None:
         if absolute <= 0:

sdk/cosmos/azure-cosmos/azure/cosmos/_service_response_retry_policy.py

@@ -21,6 +21,13 @@ def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper,
         self.connection_policy = connection_policy
         self.request = args[0] if args else None
         if self.request:
+            if self.request.retry_write:
+                # If the request is a write operation, we set the maximum retry count to be the number of
+                # write regional routing contexts available in the global endpoint manager.
+                # This ensures that we retry the write operation across all available regions.
+                # We also ensure that we retry at least once, hence the max is set to 2 by default.
+                self.max_write_retry_count = max(len(self.global_endpoint_manager.
+                                                 location_cache.write_regional_routing_contexts), 2)
             self.location_endpoint = (self.global_endpoint_manager
                                       .resolve_service_endpoint_for_partition(self.request, pk_range_wrapper))
         self.logger = logging.getLogger('azure.cosmos.ServiceResponseRetryPolicy')
@@ -35,15 +42,21 @@ def ShouldRetry(self):
             return False
 
         # Check if the next retry about to be done is safe
-        if (self.failover_retry_count + 1) >= self.total_retries:
+        if ((self.failover_retry_count + 1) >= self.total_retries and
+                _OperationType.IsReadOnlyOperation(self.request.operation_type)):
             return False
 
         if self.request:
-            if not _OperationType.IsReadOnlyOperation(self.request.operation_type):
-                return False
 
+            if not _OperationType.IsReadOnlyOperation(self.request.operation_type) and not self.request.retry_write:
+                return False
+            if self.request.retry_write and self.failover_retry_count + 1 >= self.max_write_retry_count:
+                # If we have already retried the write operation to the maximum allowed number of times,
+                # we do not retry further.
+                return False
             self.location_endpoint = self.resolve_next_region_service_endpoint()
             self.request.route_to_location(self.location_endpoint)
+
         return True
 
     # This function prepares the request to go to the next region

sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py

@@ -11,14 +11,23 @@ class _TimeoutFailoverRetryPolicy(object):
 
     def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper, *args):
         self.retry_after_in_milliseconds = 500
+        self.args = args
+        self.request = args[0] if args else None
+
         self.global_endpoint_manager = global_endpoint_manager
         self.pk_range_wrapper = pk_range_wrapper
         # If an account only has 1 region, then we still want to retry once on the same region
-        self._max_retry_attempt_count = (len(self.global_endpoint_manager.location_cache.read_regional_routing_contexts)
-                                         + 1)
+        # We want this to be the default retry attempts as paging through a query means there are requests without
+        # a request object
+        self._max_retry_attempt_count = len(self.global_endpoint_manager.location_cache
+                                            .read_regional_routing_contexts) + 1
+       # If the request is a write operation, we only want to retry once if retry write is enabled
+        if self.request and _OperationType.IsWriteOperation(self.request.operation_type):
+            self._max_retry_attempt_count = len(
+                self.global_endpoint_manager.location_cache.write_regional_routing_contexts
+            ) + 1
         self.retry_count = 0
         self.connection_policy = connection_policy
-        self.request = args[0] if args else None
 
     def ShouldRetry(self, _exception):
         """Returns true if the request should retry based on the passed-in exception.
@@ -27,8 +36,8 @@ def ShouldRetry(self, _exception):
         :returns: a boolean stating whether the request should be retried
         :rtype: bool
         """
-        # we don't retry on write operations for timeouts or any internal server errors
-        if self.request and (not _OperationType.IsReadOnlyOperation(self.request.operation_type)):
+        # we retry only if the request is a read operation or if it is a write operation with retry enabled
+        if self.request and not self.is_operation_retryable():
             return False
 
         if not self.connection_policy.EnableEndpointDiscovery:
@@ -39,7 +48,10 @@ def ShouldRetry(self, _exception):
         if self.retry_count >= self._max_retry_attempt_count:
             return False
 
-        if self.request:
+        # second check here ensures we only do cross-regional retries for read requests
+        # non-idempotent write retries should only be retried once, using preferred locations if available (MM)
+        if self.request and (_OperationType.IsReadOnlyOperation(self.request.operation_type)
+                             or self.global_endpoint_manager.can_use_multiple_write_locations(self.request)):
             location_endpoint = self.resolve_next_region_service_endpoint()
             self.request.route_to_location(location_endpoint)
         return True
@@ -56,3 +68,8 @@ def resolve_next_region_service_endpoint(self):
         # Resolve the endpoint for the request and pin the resolution to the resolved endpoint
         # This enables marking the endpoint unavailability on endpoint failover/unreachability
         return self.global_endpoint_manager.resolve_service_endpoint_for_partition(self.request, self.pk_range_wrapper)
+
+    def is_operation_retryable(self):
+        if _OperationType.IsReadOnlyOperation(self.request.operation_type):
+            return True
+        return self.request.retry_write