[Cosmos] Per-Partition Automatic Failover (#41588)

* sync PPAF

* async changes

* Update test_per_partition_automatic_failover_async.py

* CI fixes

* changelog

* broken link

* Update test_location_cache.py

* change PPAF detection logic

* Update _global_partition_endpoint_manager_circuit_breaker_core.py

* Update _global_partition_endpoint_manager_circuit_breaker_core.py

* fix tests and remove environment variable

* fix tests

* revert excluded locations change

* fix analyze

* test excluded locations

* Add different error handling for 503 and 408s, update README

* mypy, cspell, pylint

* remove tag from tests since config is service based

* add threshold-based retries for 408, 5xx errors

* update constant use, rollback session token PR change

* threshold based retries

* Update _base.py

* cspell, test fixes

* Update _service_unavailable_retry_policy.py

* mypy, pylint

* 503 behavior change, use regional contexts

* mypy, pylint, tests

* special-casing 503s

* small fix

* exclude region tests

* session retry tests

* pylint, cspell

* change errors since 503 is now retried directly

* Update sdk/cosmos/azure-cosmos/README.md

Co-authored-by: Abhijeet Mohanty <mabhijeet1995@gmail.com>

* address comments

update changelog, update docs, add typehints and documentation

* Update _service_unavailable_retry_policy.py

* small test updates for 503 behavior

* further comments

* Update test_per_partition_circuit_breaker_sm_mrr.py

* test fixes

* Update test_excluded_locations.py

* small improvement to region-finding

* pylint

* Update _global_partition_endpoint_manager_per_partition_automatic_failover.py

* address comments, add threshold lock

* add more comments

* edge cases

* changes from testing

* pylint

* fixes pylint/mypy

* mypy complaining about assigning str to none

* testing changes - will roll back later

* Update _endpoint_discovery_retry_policy.py

* Update _asynchronous_request.py

* add user agent feature flags

* Update test_per_partition_automatic_failover_async.py

* move user agent logic

* sync and async match, remove print statements

* leftover timer

* Update _retry_utility.py

* use constants

* pylint

* Update CHANGELOG.md

* react to comments

* Update _retry_utility.py

* mypy pylint

* test fixes

* add lock to failure additions

---------

Co-authored-by: tvaron3 <tomas.varon1802@gmail.com>
Co-authored-by: Abhijeet Mohanty <mabhijeet1995@gmail.com>

Author: 3 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,12 +3,14 @@
 ### 4.14.3 (Unreleased)
 
 #### Features Added
+* Added support for Per Partition Automatic Failover. To enable this feature, you must follow the guide [here](https://learn.microsoft.com/azure/cosmos-db/how-to-configure-per-partition-automatic-failover). See [PR 41588](https://github.com/Azure/azure-sdk-for-python/pull/41588).
 
 #### Breaking Changes
 
 #### Bugs Fixed
 
 #### Other Changes
+* Added cross-regional retries for 503 (Service Unavailable) errors. See [PR 41588](https://github.com/Azure/azure-sdk-for-python/pull/41588).
 
 ### 4.14.2 (2025-11-14)
 

sdk/cosmos/azure-cosmos/README.md

@@ -940,6 +940,11 @@ requests to another region:
 - `AZURE_COSMOS_FAILURE_PERCENTAGE_TOLERATED`: Default is a `90` percent failure rate.
   - After a partition reaches a 90 percent failure rate for all requests, the SDK will send requests routed to that partition to another region.
 
+### Per Partition Automatic Failover (Public Preview)
+Per partition automatic failover enables the SDK to automatically redirect write requests at the partition level to another region based on service-side signals. This feature is available 
+only for single write region accounts that have at least one read-only region. When per partition automatic failover is enabled, per partition circuit breaker and cross-region hedging is enabled by default, meaning 
+all its configurable options also apply to per partition automatic failover. To enable this feature, follow the guide [here](https://learn.microsoft.com/azure/cosmos-db/how-to-configure-per-partition-automatic-failover).
+
 ## Troubleshooting
 
 ### General

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

@@ -46,9 +46,13 @@
 if TYPE_CHECKING:
     from ._cosmos_client_connection import CosmosClientConnection
     from .aio._cosmos_client_connection_async import CosmosClientConnection as AsyncClientConnection
+    from ._global_partition_endpoint_manager_per_partition_automatic_failover import (
+        _GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover)
     from ._request_object import RequestObject
+    from ._routing.routing_range import PartitionKeyRangeWrapper
 
 # pylint: disable=protected-access
+#cspell:ignore PPAF, ppaf
 
 _COMMON_OPTIONS = {
     'initial_headers': 'initialHeaders',

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

@@ -22,8 +22,9 @@
 """Class for defining internal constants in the Azure Cosmos database service.
 """
 
-
+from enum import IntEnum
 from typing_extensions import Literal
+# cspell:ignore PPAF
 
 # cspell:ignore reranker
 
@@ -40,6 +41,7 @@ class _Constants:
     Name: Literal["name"] = "name"
     DatabaseAccountEndpoint: Literal["databaseAccountEndpoint"] = "databaseAccountEndpoint"
     DefaultEndpointsRefreshTime: int = 5 * 60 * 1000 # milliseconds
+    EnablePerPartitionFailoverBehavior: Literal["enablePerPartitionFailoverBehavior"] = "enablePerPartitionFailoverBehavior" #pylint: disable=line-too-long
 
     # ServiceDocument Resource
     EnableMultipleWritableLocations: Literal["enableMultipleWriteLocations"] = "enableMultipleWriteLocations"
@@ -74,6 +76,10 @@ class _Constants:
     FAILURE_PERCENTAGE_TOLERATED = "AZURE_COSMOS_FAILURE_PERCENTAGE_TOLERATED"
     FAILURE_PERCENTAGE_TOLERATED_DEFAULT: int = 90
     # -------------------------------------------------------------------------
+    # Only applicable when per partition automatic failover is enabled --------
+    TIMEOUT_ERROR_THRESHOLD_PPAF = "AZURE_COSMOS_TIMEOUT_ERROR_THRESHOLD_FOR_PPAF"
+    TIMEOUT_ERROR_THRESHOLD_PPAF_DEFAULT: int = 10
+    # -------------------------------------------------------------------------
 
     # Error code translations
     ERROR_TRANSLATIONS: dict[int, str] = {
@@ -99,3 +105,22 @@ class Kwargs:
         """Whether to retry write operations if they fail. Used either at client level or request level."""
 
         EXCLUDED_LOCATIONS: Literal["excludedLocations"] = "excludedLocations"
+
+    class UserAgentFeatureFlags(IntEnum):
+        """
+        User agent feature flags.
+        Each flag represents a bit in a number to encode what features are enabled. Therefore, the first feature flag
+        will be 1, the second 2, the third 4, etc. When constructing the user agent suffix, the feature flags will be
+        used to encode a unique number representing the features enabled. This number will be converted into a hex
+        string following the prefix "F" to save space in the user agent as it is limited and appended to the user agent
+        suffix. This number will then be used to determine what features are enabled by decoding the hex string back
+        to a number and checking what bits are set.
+
+        Features being developed should align with the .NET SDK as a source of truth for feature flag assignments:
+        https://github.com/Azure/azure-cosmos-dotnet-v3/blob/master/Microsoft.Azure.Cosmos/src/Diagnostics/UserAgentFeatureFlags.cs
+
+        Example:
+            If the user agent suffix has "F3", this means that flags 1 and 2.
+        """
+        PER_PARTITION_AUTOMATIC_FAILOVER = 1
+        PER_PARTITION_CIRCUIT_BREAKER = 2

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

@@ -49,7 +49,7 @@
     HttpResponse  # pylint: disable=no-legacy-azure-core-http-response-import
 
 from . import _base as base
-from ._global_partition_endpoint_manager_circuit_breaker import _GlobalPartitionEndpointManagerForCircuitBreaker
+from ._global_partition_endpoint_manager_per_partition_automatic_failover import _GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover # pylint: disable=line-too-long
 from . import _query_iterable as query_iterable
 from . import _runtime_constants as runtime_constants
 from . import _session
@@ -176,7 +176,7 @@ def __init__( # pylint: disable=too-many-statements
         self.last_response_headers: CaseInsensitiveDict = CaseInsensitiveDict()
 
         self.UseMultipleWriteLocations = False
-        self._global_endpoint_manager = _GlobalPartitionEndpointManagerForCircuitBreaker(self)
+        self._global_endpoint_manager = _GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover(self)
 
         retry_policy = None
         if isinstance(self.connection_policy.ConnectionRetryConfiguration, HTTPPolicy):
@@ -2688,12 +2688,15 @@ def GetDatabaseAccount(
             database_account._ReadableLocations = result[Constants.ReadableLocations]
         if Constants.EnableMultipleWritableLocations in result:
             database_account._EnableMultipleWritableLocations = result[
-                Constants.EnableMultipleWritableLocations
-            ]
+                Constants.EnableMultipleWritableLocations]
 
         self.UseMultipleWriteLocations = (
                 self.connection_policy.UseMultipleWriteLocations and database_account._EnableMultipleWritableLocations
         )
+
+        if Constants.EnablePerPartitionFailoverBehavior in result:
+            database_account._EnablePerPartitionFailoverBehavior = result[Constants.EnablePerPartitionFailoverBehavior]
+
         if response_hook:
             response_hook(last_response_headers, result)
         return database_account

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_http_logging_policy.py

@@ -180,7 +180,7 @@ def _get_client_settings(global_endpoint_manager: Optional[_GlobalEndpointManage
             gem_client = global_endpoint_manager.client
             if gem_client and gem_client.connection_policy:
                 connection_policy: ConnectionPolicy = gem_client.connection_policy
-                client_preferred_regions = connection_policy.PreferredLocations
+                client_preferred_regions = global_endpoint_manager.location_cache.effective_preferred_locations
                 client_excluded_regions = connection_policy.ExcludedLocations
 
         if global_endpoint_manager.location_cache:

sdk/cosmos/azure-cosmos/azure/cosmos/_endpoint_discovery_retry_policy.py

@@ -23,16 +23,9 @@
 Azure Cosmos database service.
 """
 
-import logging
-from azure.cosmos.documents import _OperationType
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-log_formatter = logging.Formatter("%(levelname)s:%(message)s")
-log_handler = logging.StreamHandler()
-log_handler.setFormatter(log_formatter)
-logger.addHandler(log_handler)
+# cspell:ignore PPAF
 
+from azure.cosmos.documents import _OperationType
 
 class EndpointDiscoveryRetryPolicy(object):
     """The endpoint discovery retry policy class used for geo-replicated database accounts
@@ -43,8 +36,9 @@ class EndpointDiscoveryRetryPolicy(object):
     Max_retry_attempt_count = 120
     Retry_after_in_milliseconds = 1000
 
-    def __init__(self, connection_policy, global_endpoint_manager, *args):
+    def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper, *args):
         self.global_endpoint_manager = global_endpoint_manager
+        self.pk_range_wrapper = pk_range_wrapper
         self._max_retry_attempt_count = EndpointDiscoveryRetryPolicy.Max_retry_attempt_count
         self.failover_retry_count = 0
         self.retry_after_in_milliseconds = EndpointDiscoveryRetryPolicy.Retry_after_in_milliseconds
@@ -70,6 +64,22 @@ def ShouldRetry(self, exception):  # pylint: disable=unused-argument
 
         self.failover_retry_count += 1
 
+        # set the refresh_needed flag to ensure that endpoint list is
+        # refreshed with new writable and readable locations
+        self.global_endpoint_manager.refresh_needed = True
+
+        # If per partition automatic failover is applicable, we mark the current endpoint as unavailable
+        # and resolve the service endpoint for the partition range - otherwise, continue the default retry logic
+        if self.global_endpoint_manager.is_per_partition_automatic_failover_applicable(self.request):
+            partition_level_info = self.global_endpoint_manager.partition_range_to_failover_info[self.pk_range_wrapper]
+            location = self.global_endpoint_manager.location_cache.get_location_from_endpoint(
+                str(self.request.location_endpoint_to_route))
+            regional_endpoint = (self.global_endpoint_manager.location_cache.
+                                account_read_regional_routing_contexts_by_location.get(location))
+            partition_level_info.unavailable_regional_endpoints[location] = regional_endpoint
+            self.global_endpoint_manager.resolve_service_endpoint_for_partition(self.request, self.pk_range_wrapper)
+            return True
+
         if self.request.location_endpoint_to_route:
             context = self.__class__.__name__
             if _OperationType.IsReadOnlyOperation(self.request.operation_type):
@@ -82,16 +92,11 @@ def ShouldRetry(self, exception):  # pylint: disable=unused-argument
                     self.request.location_endpoint_to_route,
                     True, context)
 
-        # set the refresh_needed flag to ensure that endpoint list is
-        # refreshed with new writable and readable locations
-        self.global_endpoint_manager.refresh_needed = True
-
         # clear previous location-based routing directive
         self.request.clear_route_to_location()
 
         # set location-based routing directive based on retry count
-        # simulating single master writes by ensuring usePreferredLocations
-        # is set to false
+        # simulating single master writes by ensuring usePreferredLocations is set to false
         # reasoning being that 403.3 is only expected for write region failover in single writer account
         # and we must rely on account locations as they are the source of truth
         self.request.route_to_location_with_preferred_location_flag(self.failover_retry_count, False)

sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker.py

@@ -36,6 +36,8 @@
 if TYPE_CHECKING:
     from azure.cosmos._cosmos_client_connection import CosmosClientConnection
 
+#cspell:ignore ppcb
+
 class _GlobalPartitionEndpointManagerForCircuitBreaker(_GlobalEndpointManager):
     """
     This internal class implements the logic for partition endpoint management for
@@ -93,16 +95,17 @@ def create_pk_range_wrapper(self, request: RequestObject) -> Optional[PartitionK
 
         return PartitionKeyRangeWrapper(partition_range, container_rid)
 
-    def record_failure(
+    def record_ppcb_failure(
             self,
-            request: RequestObject
-    ) -> None:
+            request: RequestObject,
+            pk_range_wrapper: Optional[PartitionKeyRangeWrapper] = None)-> None:
         if self.is_circuit_breaker_applicable(request):
-            pk_range_wrapper = self.create_pk_range_wrapper(request)
+            if pk_range_wrapper is None:
+                pk_range_wrapper = self.create_pk_range_wrapper(request)
             if pk_range_wrapper:
                 self.global_partition_endpoint_manager_core.record_failure(request, pk_range_wrapper)
 
-    def resolve_service_endpoint_for_partition(
+    def _resolve_service_endpoint_for_partition_circuit_breaker(
             self,
             request: RequestObject,
             pk_range_wrapper: Optional[PartitionKeyRangeWrapper]
@@ -113,11 +116,12 @@ def resolve_service_endpoint_for_partition(
                                                                                                     pk_range_wrapper)
         return self._resolve_service_endpoint(request)
 
-    def record_success(
+    def record_ppcb_success(
             self,
-            request: RequestObject
-    ) -> None:
-        if self.global_partition_endpoint_manager_core.is_circuit_breaker_applicable(request):
-            pk_range_wrapper = self.create_pk_range_wrapper(request)
+            request: RequestObject,
+            pk_range_wrapper: Optional[PartitionKeyRangeWrapper] = None) -> None:
+        if self.is_circuit_breaker_applicable(request):
+            if pk_range_wrapper is None:
+                pk_range_wrapper = self.create_pk_range_wrapper(request)
             if pk_range_wrapper:
                 self.global_partition_endpoint_manager_core.record_success(request, pk_range_wrapper)

sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker_core.py

@@ -19,6 +19,8 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+# pylint: disable=protected-access
+
 """Internal class for global endpoint manager for circuit breaker.
 """
 import logging
@@ -60,7 +62,10 @@ def is_circuit_breaker_applicable(self, request: RequestObject) -> bool:
             return False
 
         circuit_breaker_enabled = os.environ.get(Constants.CIRCUIT_BREAKER_ENABLED_CONFIG,
-                                                 Constants.CIRCUIT_BREAKER_ENABLED_CONFIG_DEFAULT) == "True"
+                                                 Constants.CIRCUIT_BREAKER_ENABLED_CONFIG_DEFAULT).lower() == "true"
+        if not circuit_breaker_enabled and self.client._global_endpoint_manager is not None:
+            if self.client._global_endpoint_manager._database_account_cache is not None:
+                circuit_breaker_enabled = self.client._global_endpoint_manager._database_account_cache._EnablePerPartitionFailoverBehavior is True # pylint: disable=line-too-long
         if not circuit_breaker_enabled:
             return False