[Cosmos] Resiliency and Documentation Improvements

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,14 +3,18 @@
 ### 4.7.1 (Unreleased)
 
 #### Features Added
+* SDK will now retry all failing outgoing requests three times before failing. See [PR 36514](https://github.com/Azure/azure-sdk-for-python/pull/36514).
 
 #### Breaking Changes
 
 #### 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)
+* 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).
+* Added retry handling logic for DatabaseAccountNotFound exceptions. See [PR 36514](https://github.com/Azure/azure-sdk-for-python/pull/36514).
 
 #### Other Changes
 
+
 ### 4.7.0 (2024-05-15)
 
 #### Features Added

sdk/cosmos/azure-cosmos/azure/cosmos/_global_endpoint_manager.py

@@ -89,29 +89,29 @@ def force_refresh(self, database_account):
         self.refresh_endpoint_list(database_account)
 
     def refresh_endpoint_list(self, database_account, **kwargs):
-        with self.refresh_lock:
-            # if refresh is not needed or refresh is already taking place, return
-            if not self.refresh_needed:
-                return
-            try:
-                self._refresh_endpoint_list_private(database_account, **kwargs)
-            except Exception as e:
-                raise e
+        if self.refresh_needed:
+            with self.refresh_lock:
+                # if refresh is not needed or refresh is already taking place, return
+                if not self.refresh_needed:
+                    return
+                try:
+                    self._refresh_endpoint_list_private(database_account, **kwargs)
+                except Exception as e:
+                    raise e
 
     def _refresh_endpoint_list_private(self, database_account=None, **kwargs):
+        self.refresh_needed = False
         if database_account:
             self.location_cache.perform_on_database_account_read(database_account)
-            self.refresh_needed = False
-
-        if (
-            self.location_cache.should_refresh_endpoints()
-            and self.location_cache.current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms
-        ):
-            if not database_account:
+        else:
+            if (
+                self.location_cache.should_refresh_endpoints()
+                and
+                self.location_cache.current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms
+            ):
                 database_account = self._GetDatabaseAccount(**kwargs)
                 self.location_cache.perform_on_database_account_read(database_account)
                 self.last_refresh_time = self.location_cache.current_time_millis()
-                self.refresh_needed = False
 
     def _GetDatabaseAccount(self, **kwargs):
         """Gets the database account.

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

@@ -25,7 +25,7 @@
 import time
 from typing import Optional
 
-from azure.core.exceptions import AzureError, ClientAuthenticationError
+from azure.core.exceptions import AzureError, ClientAuthenticationError, ServiceRequestError
 from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import RetryPolicy
 
@@ -100,7 +100,8 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs):
         except exceptions.CosmosHttpResponseError as e:
             retry_policy = defaultRetry_policy
             # Re-assign retry policy based on error code
-            if e.status_code == StatusCodes.FORBIDDEN and e.sub_status == SubStatusCodes.WRITE_FORBIDDEN:
+            if e.status_code == StatusCodes.FORBIDDEN and e.sub_status in\
+                    [SubStatusCodes.DATABASE_ACCOUNT_NOT_FOUND, SubStatusCodes.WRITE_FORBIDDEN]:
                 retry_policy = endpointDiscovery_retry_policy
             elif e.status_code == StatusCodes.TOO_MANY_REQUESTS:
                 retry_policy = resourceThrottle_retry_policy
@@ -210,6 +211,15 @@ def send(self, request):
                 timeout_error.response = response
                 timeout_error.history = retry_settings['history']
                 raise
+            except ServiceRequestError as err:
+                # the request ran into a socket timeout or failed to establish a new connection
+                # since request wasn't sent, we retry up to however many connection retries are configured (default 3)
+                if retry_settings['connect'] > 0:
+                    retry_active = self.increment(retry_settings, response=request, error=err)
+                    if retry_active:
+                        self.sleep(retry_settings, request.context.transport)
+                        continue
+                raise err
             except AzureError as err:
                 retry_error = err
                 if self._is_method_retryable(retry_settings, request.http_request):

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

@@ -59,7 +59,6 @@ def _request_body_from_data(data):
     if data is None or isinstance(data, str) or _is_readable_stream(data):
         return data
     if isinstance(data, (dict, list, tuple)):
-
         json_dumped = json.dumps(data, separators=(",", ":"))
 
         return json_dumped
@@ -97,11 +96,15 @@ def _Request(global_endpoint_manager, request_params, connection_policy, pipelin
             raise exceptions.CosmosClientTimeoutError()
 
     if request_params.endpoint_override:
-        base_url = request_params.endpoint_override
+        base_url = request_params.location_endpoint_to_route
     else:
         base_url = global_endpoint_manager.resolve_service_endpoint(request_params)
-    if base_url != pipeline_client._base_url:
-        request.url = request.url.replace(pipeline_client._base_url, base_url)
+    if base_url != request.url:
+        # TODO: remove this comment for the PR
+        # this was checking the wrong thing, pipeline_client.base_url will always be the base non-locational url
+        # due to this, the request url will never get updated since you won't find the base_url in the locational one
+        # trying to find and replace contoso.documents.azure.com in contoso-west.documents.azure.com will never happen
+        request.url = _replace_url_prefix(request.url, base_url)
 
     parse_result = urlparse(request.url)
 
@@ -167,20 +170,31 @@ def _Request(global_endpoint_manager, request_params, connection_policy, pipelin
     return result, headers
 
 
+def _replace_url_prefix(original_url, new_prefix):
+    parts = original_url.split('/', 3)
+
+    if not new_prefix.endswith('/'):
+        new_prefix += '/'
+
+    new_url = new_prefix + parts[3] if len(parts) > 3 else new_prefix
+
+    return new_url
+
+
 def _PipelineRunFunction(pipeline_client, request, **kwargs):
     # pylint: disable=protected-access
 
     return pipeline_client._pipeline.run(request, **kwargs)
 
 def SynchronizedRequest(
-    client,
-    request_params,
-    global_endpoint_manager,
-    connection_policy,
-    pipeline_client,
-    request,
-    request_data,
-    **kwargs
+        client,
+        request_params,
+        global_endpoint_manager,
+        connection_policy,
+        pipeline_client,
+        request,
+        request_data,
+        **kwargs
 ):
     """Performs one synchronized http request according to the parameters.
 

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

@@ -31,7 +31,7 @@
 from .. import exceptions
 from .. import http_constants
 from . import _retry_utility_async
-from .._synchronized_request import _request_body_from_data
+from .._synchronized_request import _request_body_from_data, _replace_url_prefix
 
 
 async def _Request(global_endpoint_manager, request_params, connection_policy, pipeline_client, request, **kwargs):
@@ -65,11 +65,15 @@ async def _Request(global_endpoint_manager, request_params, connection_policy, p
             raise exceptions.CosmosClientTimeoutError()
 
     if request_params.endpoint_override:
-        base_url = request_params.endpoint_override
+        base_url = request_params.location_endpoint_to_route
     else:
         base_url = global_endpoint_manager.resolve_service_endpoint(request_params)
-    if base_url != pipeline_client._base_url:
-        request.url = request.url.replace(pipeline_client._base_url, base_url)
+    if base_url != request.url:
+        # TODO: remove this comment for the PR
+        # this was checking the wrong thing, pipeline_client.base_url will always be the base non-locational url
+        # due to this, the request url will never get updated since you won't find the base_url in the locational one
+        # trying to find and replace contoso.documents.azure.com in contoso-west.documents.azure.com will never happen
+        request.url = _replace_url_prefix(request.url, base_url)
 
     parse_result = urlparse(request.url)
 

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

@@ -25,7 +25,7 @@
 import time
 import asyncio
 
-from azure.core.exceptions import AzureError, ClientAuthenticationError
+from azure.core.exceptions import AzureError, ClientAuthenticationError, ServiceRequestError
 from azure.core.pipeline.policies import AsyncRetryPolicy
 
 from .. import exceptions
@@ -97,7 +97,8 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg
             return result
         except exceptions.CosmosHttpResponseError as e:
             retry_policy = None
-            if e.status_code == StatusCodes.FORBIDDEN and e.sub_status == SubStatusCodes.WRITE_FORBIDDEN:
+            if e.status_code == StatusCodes.FORBIDDEN and e.sub_status in \
+                    [SubStatusCodes.DATABASE_ACCOUNT_NOT_FOUND, SubStatusCodes.WRITE_FORBIDDEN]:
                 retry_policy = endpointDiscovery_retry_policy
             elif e.status_code == StatusCodes.TOO_MANY_REQUESTS:
                 retry_policy = resourceThrottle_retry_policy
@@ -194,6 +195,15 @@ async def send(self, request):
                 timeout_error.response = response
                 timeout_error.history = retry_settings['history']
                 raise
+            except ServiceRequestError as err:
+                # the request ran into a socket timeout or failed to establish a new connection
+                # since request wasn't sent, we retry up to however many connection retries are configured (default 3)
+                if retry_settings['connect'] > 0:
+                    retry_active = self.increment(retry_settings, response=request, error=err)
+                    if retry_active:
+                        await self.sleep(retry_settings, request.context.transport)
+                        continue
+                raise err
             except AzureError as err:
                 retry_error = err
                 if self._is_method_retryable(retry_settings, request.http_request):

sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py

@@ -124,7 +124,6 @@ def _build_connection_policy(kwargs: Dict[str, Any]) -> ConnectionPolicy:
             "'connection_retry_policy' has been deprecated and will be removed from the SDK in a future release.",
             DeprecationWarning
         )
-    connection_retry = policy.ConnectionRetryConfiguration
     if not connection_retry:
         connection_retry = ConnectionRetryPolicy(
             retry_total=total_retries,

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

@@ -193,6 +193,7 @@ class HttpHeaders:
     PartitionKey = "x-ms-documentdb-partitionkey"
     EnableCrossPartitionQuery = "x-ms-documentdb-query-enablecrosspartition"
     PartitionKeyRangeID = "x-ms-documentdb-partitionkeyrangeid"
+    PhysicalPartitionId = "x-ms-cosmos-physical-partition-id"
     PartitionKeyDeletePending = "x-ms-cosmos-is-partition-key-delete-pending"
     StartEpkString = "x-ms-start-epk"
     EndEpkString = "x-ms-end-epk"
@@ -239,6 +240,8 @@ class HttpHeaders:
     GlobalCommittedLsn = "x-ms-global-committed-lsn"
     NumberOfReadRegions = "x-ms-number-of-read-regions"
     TransportRequestId = "x-ms-transport-request-id"
+    ItemLsn = "x-ms-item-lsn"
+    CosmosItemLsn = "x-ms-cosmos-item-llsn"  # cspell:disable-line
     CosmosLsn = "x-ms-cosmos-llsn"  # cspell:disable-line
     CosmosQuorumAckedLsn = "x-ms-cosmos-quorum-acked-llsn"  # cspell:disable-line
     RequestDurationMs = "x-ms-request-duration-ms"