[Bugfix] Partition Scoping Incorrect for PPCB (#42751)

* add logs and improve diagnostics

* fix logs

* fix partitionkey issue

* add tests

* remove logs and add change log

* fix pylint

* fix tests

Author: tvaron3

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -12,6 +12,7 @@
 * Fixed bug where during health checks read regions were marked as unavailable for write operations. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
 * Fixed bug where containers named with spaces or special characters using session consistency would fall back to eventual consistency. See [PR 42608](https://github.com/Azure/azure-sdk-for-python/pull/42608)
 * Fixed bug where `excluded_locations` was not being honored for some metadata calls. See [PR 42266](https://github.com/Azure/azure-sdk-for-python/pull/42266).
+* Fixed partition scoping for per partition circuit breaker. See [PR 42751](https://github.com/Azure/azure-sdk-for-python/pull/42751)
 
 #### Other Changes
 * Added session token false progress merge logic. See [42393](https://github.com/Azure/azure-sdk-for-python/pull/42393)

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

@@ -2122,7 +2122,8 @@ def PatchItem(
         # Patch will use WriteEndpoint since it uses PUT operation
         request_params = RequestObject(resource_type,
                                        documents._OperationType.Patch,
-                                       headers)
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
@@ -2216,7 +2217,8 @@ def _Batch(
                                   documents._OperationType.Batch, options)
         request_params = RequestObject(http_constants.ResourceType.Document,
                                        documents._OperationType.Batch,
-                                       headers)
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         return cast(
@@ -2280,7 +2282,8 @@ def DeleteAllItemsByPartitionKey(
                                   http_constants.ResourceType.PartitionKey, documents._OperationType.Delete, options)
         request_params = RequestObject(http_constants.ResourceType.PartitionKey,
                                        documents._OperationType.Delete,
-                                       headers)
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         _, last_response_headers = self.__Post(
             path=path,
@@ -2654,7 +2657,7 @@ def GetDatabaseAccount(
         request_params = RequestObject(http_constants.ResourceType.DatabaseAccount,
                                        documents._OperationType.Read,
                                        headers,
-                                       url_connection)
+                                       endpoint_override=url_connection)
         result, last_response_headers = self.__Get("", request_params, headers, **kwargs)
         self.last_response_headers = last_response_headers
         database_account = DatabaseAccount()
@@ -2706,7 +2709,7 @@ def _GetDatabaseAccountCheck(
         request_params = RequestObject(http_constants.ResourceType.DatabaseAccount,
                                        documents._OperationType.Read,
                                        headers,
-                                       url_connection)
+                                       endpoint_override=url_connection)
         self.__Get("", request_params, headers, **kwargs)
 
 
@@ -2744,7 +2747,10 @@ def Create(
         headers = base.GetHeaders(self, initial_headers, "post", path, id, resource_type,
                                     documents._OperationType.Create, options)
         # Create will use WriteEndpoint since it uses POST operation
-        request_params = RequestObject(resource_type, documents._OperationType.Create, headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.Create,
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
@@ -2792,7 +2798,10 @@ def Upsert(
                                     documents._OperationType.Upsert, options)
         headers[http_constants.HttpHeaders.IsUpsert] = True
         # Upsert will use WriteEndpoint since it uses POST operation
-        request_params = RequestObject(resource_type, documents._OperationType.Upsert, headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.Upsert,
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
@@ -2838,7 +2847,10 @@ def Replace(
         headers = base.GetHeaders(self, initial_headers, "put", path, id, resource_type,
                                     documents._OperationType.Replace, options)
         # Replace will use WriteEndpoint since it uses PUT operation
-        request_params = RequestObject(resource_type, documents._OperationType.Replace, headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.Replace,
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
@@ -2883,7 +2895,10 @@ def Read(
         headers = base.GetHeaders(self, initial_headers, "get", path, id, resource_type,
                                     documents._OperationType.Read, options)
         # Read will use ReadEndpoint since it uses GET operation
-        request_params = RequestObject(resource_type, documents._OperationType.Read, headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.Read,
+                                       headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         base.set_session_token_header(self, headers, path, request_params, options)
         result, last_response_headers = self.__Get(path, request_params, headers, **kwargs)
@@ -2927,7 +2942,10 @@ def DeleteResource(
         headers = base.GetHeaders(self, initial_headers, "delete", path, id, resource_type,
                                     documents._OperationType.Delete, options)
         # Delete will use WriteEndpoint since it uses DELETE operation
-        request_params = RequestObject(resource_type, documents._OperationType.Delete, headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.Delete,
+                                       headers,
+                                       options.get("partitionKey", None))
         base.set_session_token_header(self, headers, path, request_params, options)
         request_params.set_retry_write(options, self.connection_policy.RetryNonIdempotentWrites)
         request_params.set_excluded_location_from_options(options)
@@ -3178,7 +3196,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]:
             request_params = RequestObject(
                 resource_type,
                 op_type,
-                headers
+                headers,
+                options.get("partitionKey", None)
             )
             request_params.set_excluded_location_from_options(options)
             base.set_session_token_header(self, headers, path, request_params, options, partition_key_range_id)
@@ -3219,7 +3238,10 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]:
             options,
             partition_key_range_id
         )
-        request_params = RequestObject(resource_type, documents._OperationType.SqlQuery, req_headers)
+        request_params = RequestObject(resource_type,
+                                       documents._OperationType.SqlQuery,
+                                       req_headers,
+                                       options.get("partitionKey", None))
         request_params.set_excluded_location_from_options(options)
         if not is_query_plan:
             req_headers[http_constants.HttpHeaders.IsQuery] = "true"

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

@@ -68,8 +68,8 @@ def create_pk_range_wrapper(self, request: RequestObject) -> Optional[PartitionK
         options = {}
         if request.excluded_locations:
             options[_Constants.Kwargs.EXCLUDED_LOCATIONS] = request.excluded_locations
-        if HttpHeaders.PartitionKey in request.headers:
-            partition_key_value = request.headers[HttpHeaders.PartitionKey]
+        if request.pk_val:
+            partition_key_value = request.pk_val
             # get the partition key range for the given partition key
             epk_range = [partition_key._get_epk_range_for_partition_key(partition_key_value)] # pylint: disable=protected-access
             partition_ranges = (self.client._routing_map_provider # pylint: disable=protected-access

sdk/cosmos/azure-cosmos/azure/cosmos/_partition_health_tracker.py

@@ -172,6 +172,10 @@ def check_stale_partition_info(
                         # only one request should be used to recover
                         with self.stale_partition_lock:
                             if _should_mark_healthy_tentative(partition_health_info, current_time):
+                                logger.debug("Attempting recovery for %s in %s where health info is %s.",
+                                            pk_range_wrapper,
+                                            location,
+                                            partition_health_info)
                                 # this will trigger one attempt to recover
                                 partition_health_info.transition_health_status(UNHEALTHY, current_time)
                                 request.healthy_tentative_location = location
@@ -235,6 +239,11 @@ def add_failure(
 
         # Retrieve the consecutive failure threshold from the environment.
         consecutive_failure_threshold = int(os.environ.get(env_key, default_consecutive_threshold))
+        # log the current stats
+        logger.debug("Failure for partition %s in location %s has %s",
+                    pk_range_wrapper,
+                     location,
+                     self.pk_range_wrapper_to_health_info[pk_range_wrapper][location])
 
         # Call the threshold checker with the current stats.
         self._check_thresholds(

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

@@ -32,6 +32,7 @@ def __init__(
             resource_type: str,
             operation_type: str,
             headers: Dict[str, Any],
+            pk_val: Optional[Any] = None,
             endpoint_override: Optional[str] = None,
     ) -> None:
         self.resource_type = resource_type
@@ -46,6 +47,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.pk_val = pk_val
         self.retry_write: bool = False
 
     def route_to_location_with_preferred_location_flag(  # pylint: disable=name-too-long