Change default read timeout for database account (Azure#39596)

* change default read timeout

* fix tests

* Add read timeout tests for database account calls

* fix timeout retry policy

* Fixed the timeout logic

* Fixed the timeout retry policy

* Mock tests for timeout and failover retry policy

* Fixed async tests

* Updated versions for release and changelog

* Add tests for cross region tries

---------

Co-authored-by: Kushagra Thapar <[email protected]>
Co-authored-by: Kushagra Thapar <[email protected]>

Author: tvaron3

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,14 +1,11 @@
 ## Release History
 
-### 4.9.1b4 (Unreleased)
-
-#### Features Added
-
-#### Breaking Changes
+### 4.9.1b4 (2025-02-06)
 
 #### Bugs Fixed
-
-#### Other Changes
+* Improved retry logic for read requests to failover on other regions in case of timeouts and any error codes >= 500. See [PR 39596](https://github.com/Azure/azure-sdk-for-python/pull/39596)
+* Fixed a regression where read operations were not retrying on timeouts. See [PR 39596](https://github.com/Azure/azure-sdk-for-python/pull/39596)
+* Updated default read timeout for `getDatabaseAccount` calls to 3 seconds. See [PR 39596](https://github.com/Azure/azure-sdk-for-python/pull/39596)
 
 ### 4.9.1b3 (2025-02-04)
 

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

@@ -131,7 +131,6 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs):
                     sub_status_code=SubStatusCodes.THROUGHPUT_OFFER_NOT_FOUND)
             return result
         except exceptions.CosmosHttpResponseError as e:
-            retry_policy = defaultRetry_policy
             if request and _has_database_account_header(request.headers):
                 retry_policy = database_account_retry_policy
             # Re-assign retry policy based on error code
@@ -173,8 +172,12 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs):
 
                     retry_policy.container_rid = cached_container["_rid"]
                     request.headers[retry_policy._intended_headers] = retry_policy.container_rid
-            elif e.status_code in [StatusCodes.REQUEST_TIMEOUT, StatusCodes.SERVICE_UNAVAILABLE]:
+            elif e.status_code == StatusCodes.REQUEST_TIMEOUT:
                 retry_policy = timeout_failover_retry_policy
+            elif e.status_code >= StatusCodes.INTERNAL_SERVER_ERROR:
+                retry_policy = timeout_failover_retry_policy
+            else:
+                retry_policy = defaultRetry_policy
 
             # If none of the retry policies applies or there is no retry needed, set the
             # throttle related response headers and re-throw the exception back arg[0]

sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py

@@ -93,6 +93,8 @@ def _Request(global_endpoint_manager, request_params, connection_policy, pipelin
     if request_params.resource_type != http_constants.ResourceType.DatabaseAccount:
         global_endpoint_manager.refresh_endpoint_list(None, **kwargs)
     else:
+        # always override database account call timeouts
+        read_timeout = connection_policy.DBAReadTimeout
         connection_timeout = connection_policy.DBAConnectionTimeout
     if client_timeout is not None:
         kwargs['timeout'] = client_timeout - (time.time() - start_time)

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

@@ -5,20 +5,18 @@
 Cosmos database service.
 """
 from azure.cosmos.documents import _OperationType
-from . import http_constants
 
 
 class _TimeoutFailoverRetryPolicy(object):
 
     def __init__(self, connection_policy, global_endpoint_manager, *args):
-        self._max_retry_attempt_count = 120
-        self._max_service_unavailable_retry_count = 1
-        self.retry_after_in_milliseconds = 0
+        self.retry_after_in_milliseconds = 500
         self.args = args
 
         self.global_endpoint_manager = global_endpoint_manager
+        # 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_endpoints) + 1
         self.retry_count = 0
-        self.location_index = 0
         self.connection_policy = connection_policy
         self.request = args[0] if args else None
 
@@ -29,64 +27,32 @@ def ShouldRetry(self, _exception):
         :returns: a boolean stating whether the request should be retried
         :rtype: bool
         """
-        if self.request:
-            if _OperationType.IsReadOnlyOperation(self.request.operation_type):
-                return False
+        # 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)):
+            return False
 
         if not self.connection_policy.EnableEndpointDiscovery:
             return False
 
-
-        # Check if the next retry about to be done is safe
-        if _exception.status_code == http_constants.StatusCodes.SERVICE_UNAVAILABLE and \
-                self.retry_count >= self._max_service_unavailable_retry_count:
-            return False
         self.retry_count += 1
         # Check if the next retry about to be done is safe
         if self.retry_count >= self._max_retry_attempt_count:
             return False
 
         if self.request:
-            # Update the last routed location to where this request was routed previously.
-            # So that we can check in location cache if we need to return the current or previous
-            # based on where the request was routed previously.
-            self.request.last_routed_location_endpoint_within_region = self.request.location_endpoint_to_route
-
-            if _OperationType.IsReadOnlyOperation(self.request.operation_type):
-                # We just directly got to the next location in case of read requests
-                # We don't retry again on the same region for regional endpoint
-                location_endpoint = self.resolve_next_region_service_endpoint()
-            else:
-                self.global_endpoint_manager.swap_regional_endpoint_values(self.request)
-                location_endpoint = self.resolve_current_region_service_endpoint()
-                # This is the case where both current and previous point to the same writable endpoint
-                # In this case we don't want to retry again, rather failover to the next region
-                if self.request.last_routed_location_endpoint_within_region == location_endpoint:
-                    location_endpoint = self.resolve_next_region_service_endpoint()
-
+            location_endpoint = self.resolve_next_region_service_endpoint()
             self.request.route_to_location(location_endpoint)
         return True
 
-
-    # This function prepares the request to go to the second endpoint in the same region
-    def resolve_current_region_service_endpoint(self):
-        # clear previous location-based routing directive
-        self.request.clear_route_to_location()
-        # resolve the next service endpoint in the same region
-        # since we maintain 2 endpoints per region for write operations
-        self.request.route_to_location_with_preferred_location_flag(self.location_index, True)
-        return self.global_endpoint_manager.resolve_service_endpoint(self.request)
-
     # This function prepares the request to go to the next region
     def resolve_next_region_service_endpoint(self):
-        self.location_index += 1
         # clear previous location-based routing directive
         self.request.clear_route_to_location()
         # clear the last routed endpoint within same region since we are going to a new region now
         self.request.last_routed_location_endpoint_within_region = None
         # set location-based routing directive based on retry count
         # ensuring usePreferredLocations is set to True for retry
-        self.request.route_to_location_with_preferred_location_flag(self.location_index, True)
+        self.request.route_to_location_with_preferred_location_flag(self.retry_count, True)
         # 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(self.request)

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

@@ -62,6 +62,8 @@ async def _Request(global_endpoint_manager, request_params, connection_policy, p
     if request_params.resource_type != http_constants.ResourceType.DatabaseAccount:
         await global_endpoint_manager.refresh_endpoint_list(None, **kwargs)
     else:
+        # always override database account call timeouts
+        read_timeout = connection_policy.DBAReadTimeout
         connection_timeout = connection_policy.DBAConnectionTimeout
     if client_timeout is not None:
         kwargs['timeout'] = client_timeout - (time.time() - start_time)

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

@@ -130,7 +130,6 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg
 
             return result
         except exceptions.CosmosHttpResponseError as e:
-            retry_policy = None
             if request and _has_database_account_header(request.headers):
                 retry_policy = database_account_retry_policy
             elif e.status_code == StatusCodes.FORBIDDEN and e.sub_status in \
@@ -171,7 +170,9 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg
 
                     retry_policy.container_rid = cached_container["_rid"]
                     request.headers[retry_policy._intended_headers] = retry_policy.container_rid
-            elif e.status_code in [StatusCodes.REQUEST_TIMEOUT, StatusCodes.SERVICE_UNAVAILABLE]:
+            elif e.status_code == StatusCodes.REQUEST_TIMEOUT:
+                retry_policy = timeout_failover_retry_policy
+            elif e.status_code >= StatusCodes.INTERNAL_SERVER_ERROR:
                 retry_policy = timeout_failover_retry_policy
             else:
                 retry_policy = defaultRetry_policy

sdk/cosmos/azure-cosmos/azure/cosmos/documents.py

@@ -332,12 +332,14 @@ class ConnectionPolicy:  # pylint: disable=too-many-instance-attributes
     __defaultRequestTimeout: int = 5  # seconds
     __defaultDBAConnectionTimeout: int = 3  # seconds
     __defaultReadTimeout: int = 65  # seconds
+    __defaultDBAReadTimeout: int = 3 # seconds
     __defaultMaxBackoff: int = 1 # seconds
 
     def __init__(self) -> None:
         self.RequestTimeout: int = self.__defaultRequestTimeout
         self.DBAConnectionTimeout: int = self.__defaultDBAConnectionTimeout
         self.ReadTimeout: int = self.__defaultReadTimeout
+        self.DBAReadTimeout: int = self.__defaultDBAReadTimeout
         self.MaxBackoff: int = self.__defaultMaxBackoff
         self.ConnectionMode: int = ConnectionMode.Gateway
         self.SSLConfiguration: Optional[SSLConfiguration] = None

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -400,7 +400,6 @@ class StatusCodes:
     RETRY_WITH = 449
 
     INTERNAL_SERVER_ERROR = 500
-    SERVICE_UNAVAILABLE = 503
 
     # Operation pause and cancel. These are FAKE status codes for QOS logging purpose only.
     OPERATION_PAUSED = 1200

sdk/cosmos/azure-cosmos/test/test_crud.py

@@ -1821,7 +1821,14 @@ def test_client_request_timeout(self):
                 container = databaseForTest.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
                 container.create_item(body={'id': str(uuid.uuid4()), 'name': 'sample'})
 
-
+    async def test_read_timeout_async(self):
+        connection_policy = documents.ConnectionPolicy()
+        # making timeout 0 ms to make sure it will throw
+        connection_policy.DBAReadTimeout = 0.000000000001
+        with self.assertRaises(ServiceResponseError):
+            # this will make a get database account call
+            with cosmos_client.CosmosClient(self.host, self.masterKey, connection_policy=connection_policy):
+                print('initialization')
 
     def test_client_request_timeout_when_connection_retry_configuration_specified(self):
         connection_policy = documents.ConnectionPolicy()

sdk/cosmos/azure-cosmos/test/test_crud_async.py

@@ -1689,6 +1689,15 @@ async def test_client_request_timeout_async(self):
                     await container.create_item(body={'id': str(uuid.uuid4()), 'name': 'sample'})
                     print('Async initialization')
 
+    async def test_read_timeout_async(self):
+        connection_policy = documents.ConnectionPolicy()
+        # making timeout 0 ms to make sure it will throw
+        connection_policy.DBAReadTimeout = 0.000000000001
+        with self.assertRaises(ServiceResponseError):
+            # this will make a get database account call
+            async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy):
+                print('Async initialization')
+
     async def test_client_request_timeout_when_connection_retry_configuration_specified_async(self):
         connection_policy = documents.ConnectionPolicy()
         # making timeout 0 ms to make sure it will throw