Staging signoff fixes (#42266)

* change workloads based on feedback

* get locations from database account call

* add getting account info to conftest

* convert more tests to only use get database account regions

* fix tests

* fix test

* fix staging tests

* fix staging tests

* fix staging tests

* fix staging tests

* fix staging tests

* fix tests

* fix tests

* fix tests

* Update sdk/cosmos/azure-cosmos/tests/test_session.py

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

* add staging yml file

* add staging pipeline

* staging tests

* staging tests

* add staging pipeline

---------

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

Author: tvaron3

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,6 +10,7 @@
 #### Bugs Fixed
 * Improved the resilience of Database Account Read metadata operation against short-lived network issues by increasing number of retries. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
 * 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 `excluded_locations` was not being honored for some metadata calls. See [PR 42266](https://github.com/Azure/azure-sdk-for-python/pull/42266).
 
 #### 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/_base.py

@@ -68,7 +68,7 @@
     'retry_write': Constants.Kwargs.RETRY_WRITE,
     'max_item_count': 'maxItemCount',
     'throughput_bucket': 'throughputBucket',
-    'excluded_locations': 'excludedLocations'
+    'excluded_locations': Constants.Kwargs.EXCLUDED_LOCATIONS
 }
 
 # Cosmos resource ID validation regex breakdown:
@@ -364,7 +364,8 @@ def set_session_token_header(
                                                                 options.get('partitionKey'),
                                                                 cosmos_client_connection._container_properties_cache,
                                                                 cosmos_client_connection._routing_map_provider,
-                                                                partition_key_range_id))
+                                                                partition_key_range_id,
+                                                                options))
                 if session_token != "":
                     headers[http_constants.HttpHeaders.SessionToken] = session_token
 
@@ -391,7 +392,8 @@ async def set_session_token_header_async(
                                                                 options.get('partitionKey'),
                                                                 cosmos_client_connection._container_properties_cache,
                                                                 cosmos_client_connection._routing_map_provider,
-                                                                partition_key_range_id)
+                                                                partition_key_range_id,
+                                                                options)
                 if session_token != "":
                     headers[http_constants.HttpHeaders.SessionToken] = session_token
 

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

@@ -94,3 +94,5 @@ class Kwargs:
 
         RETRY_WRITE: Literal["retry_write"] = "retry_write"
         """Whether to retry write operations if they fail. Used either at client level or request level."""
+
+        EXCLUDED_LOCATIONS: Literal["excludedLocations"] = "excludedLocations"

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py

@@ -258,5 +258,7 @@ async def _get_target_partition_key_range(self, target_all_ranges):
             return [item async for item in self._client._ReadPartitionKeyRanges(collection_link=self._resource_link)]
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return await self._routing_provider.get_overlapping_ranges(
-            self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
+            self._resource_link,
+            [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
+            self._options
         )

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/non_streaming_order_by_aggregator.py

@@ -110,7 +110,9 @@ async def _get_target_partition_key_range(self):
 
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return await self._routing_provider.get_overlapping_ranges(
-            self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
+            self._resource_link,
+            [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
+            self._options
         )
 
     async def _configure_partition_ranges(self):

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py

@@ -385,5 +385,7 @@ def _get_target_partition_key_range(self, target_all_ranges):
             return list(self._client._ReadPartitionKeyRanges(collection_link=self._resource_link))
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return self._routing_provider.get_overlapping_ranges(
-            self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
+            self._resource_link,
+            [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
+            self._options
         )

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/non_streaming_order_by_aggregator.py

@@ -152,5 +152,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
     def _get_target_partition_key_range(self):
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return self._routing_provider.get_overlapping_ranges(
-            self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
+            self._resource_link,
+            [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
+            self._options
         )

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

@@ -23,6 +23,7 @@
 """
 from typing import TYPE_CHECKING, Optional
 
+from azure.cosmos._constants import _Constants
 from azure.cosmos.partition_key import _get_partition_key_from_partition_key_definition
 from azure.cosmos._global_partition_endpoint_manager_circuit_breaker_core import \
     _GlobalPartitionEndpointManagerForCircuitBreakerCore
@@ -64,17 +65,20 @@ def create_pk_range_wrapper(self, request: RequestObject) -> Optional[PartitionK
         partition_key_definition = properties["partitionKey"]
         partition_key = _get_partition_key_from_partition_key_definition(partition_key_definition)
 
+        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]
             # 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
-                                      .get_overlapping_ranges(container_link, epk_range))
+                                      .get_overlapping_ranges(container_link, epk_range, options))
             partition_range = Range.PartitionKeyRangeToRange(partition_ranges[0])
         elif HttpHeaders.PartitionKeyRangeID in request.headers:
             pk_range_id = request.headers[HttpHeaders.PartitionKeyRangeID]
             epk_range =(self.client._routing_map_provider # pylint: disable=protected-access
-                    .get_range_by_partition_key_range_id(container_link, pk_range_id))
+                    .get_range_by_partition_key_range_id(container_link, pk_range_id, options))
             if not epk_range:
                 self.global_partition_endpoint_manager_core.log_warn_or_debug(
                     "Illegal state: partition key range cache not initialized correctly. "

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py

@@ -50,7 +50,7 @@ def __init__(self, client):
         # keeps the cached collection routing map by collection id
         self._collection_routing_map_by_item = {}
 
-    async def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options = None, **kwargs):
+    async def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options, **kwargs):
         """Given a partition key range and a collection, return the list of
         overlapping partition key ranges.
 
@@ -69,7 +69,7 @@ async def init_collection_routing_map_if_needed(
             self,
             collection_link: str,
             collection_id: str,
-            feed_options: Optional[Dict[str, Any]] = None,
+            feed_options: Dict[str, Any],
             **kwargs: Dict[str, Any]
     ):
         collection_routing_map = self._collection_routing_map_by_item.get(collection_id)
@@ -91,10 +91,11 @@ async def get_range_by_partition_key_range_id(
             self,
             collection_link: str,
             partition_key_range_id: int,
+            feed_options: Dict[str, Any],
             **kwargs: Dict[str, Any]
     ) -> Optional[Dict[str, Any]]:
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        await self.init_collection_routing_map_if_needed(collection_link, collection_id, **kwargs)
+        await self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
 
         return self._collection_routing_map_by_item[collection_id].get_range_by_partition_key_range_id(
             partition_key_range_id)

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py

@@ -29,6 +29,7 @@
 from . import routing_range
 from .routing_range import PartitionKeyRange
 
+
 # pylint: disable=protected-access
 
 
@@ -55,7 +56,7 @@ def init_collection_routing_map_if_needed(
             self,
             collection_link: str,
             collection_id: str,
-            feed_options: Optional[Dict[str, Any]] = None,
+            feed_options: Optional[Dict[str, Any]],
             **kwargs: Dict[str, Any]
     ):
         collection_routing_map = self._collection_routing_map_by_item.get(collection_id)
@@ -72,7 +73,7 @@ def init_collection_routing_map_if_needed(
             )
             self._collection_routing_map_by_item[collection_id] = collection_routing_map
 
-    def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options = None, **kwargs):
+    def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options, **kwargs):
         """Given a partition key range and a collection, return the list of
         overlapping partition key ranges.
 
@@ -91,10 +92,11 @@ def get_range_by_partition_key_range_id(
             self,
             collection_link: str,
             partition_key_range_id: int,
+            feed_options: Dict[str, Any],
             **kwargs: Dict[str, Any]
     ) -> Optional[Dict[str, Any]]:
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        self.init_collection_routing_map_if_needed(collection_link, collection_id, **kwargs)
+        self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
 
         return (self._collection_routing_map_by_item[collection_id]
                 .get_range_by_partition_key_range_id(partition_key_range_id))