Skip to content

Commit 7c94732

Browse files
tvaron3Copilot
tvaron3
and
Copilot
authored
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]>
1 parent 9958caf commit 7c94732

33 files changed

+193
-49
lines changed

‎sdk/cosmos/azure-cosmos/CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#### Bugs Fixed
1111
* 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).
1212
* 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).
13+
* 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).
1314

1415
#### Other Changes
1516
* 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‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@
6868
'retry_write': Constants.Kwargs.RETRY_WRITE,
6969
'max_item_count': 'maxItemCount',
7070
'throughput_bucket': 'throughputBucket',
71-
'excluded_locations': 'excludedLocations'
71+
'excluded_locations': Constants.Kwargs.EXCLUDED_LOCATIONS
7272
}
7373

7474
# Cosmos resource ID validation regex breakdown:
@@ -364,7 +364,8 @@ def set_session_token_header(
364364
options.get('partitionKey'),
365365
cosmos_client_connection._container_properties_cache,
366366
cosmos_client_connection._routing_map_provider,
367-
partition_key_range_id))
367+
partition_key_range_id,
368+
options))
368369
if session_token != "":
369370
headers[http_constants.HttpHeaders.SessionToken] = session_token
370371

@@ -391,7 +392,8 @@ async def set_session_token_header_async(
391392
options.get('partitionKey'),
392393
cosmos_client_connection._container_properties_cache,
393394
cosmos_client_connection._routing_map_provider,
394-
partition_key_range_id)
395+
partition_key_range_id,
396+
options)
395397
if session_token != "":
396398
headers[http_constants.HttpHeaders.SessionToken] = session_token
397399

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,3 +94,5 @@ class Kwargs:
9494

9595
RETRY_WRITE: Literal["retry_write"] = "retry_write"
9696
"""Whether to retry write operations if they fail. Used either at client level or request level."""
97+
98+
EXCLUDED_LOCATIONS: Literal["excludedLocations"] = "excludedLocations"

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -258,5 +258,7 @@ async def _get_target_partition_key_range(self, target_all_ranges):
258258
return [item async for item in self._client._ReadPartitionKeyRanges(collection_link=self._resource_link)]
259259
query_ranges = self._partitioned_query_ex_info.get_query_ranges()
260260
return await self._routing_provider.get_overlapping_ranges(
261-
self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
261+
self._resource_link,
262+
[routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
263+
self._options
262264
)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,9 @@ async def _get_target_partition_key_range(self):
110110

111111
query_ranges = self._partitioned_query_ex_info.get_query_ranges()
112112
return await self._routing_provider.get_overlapping_ranges(
113-
self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
113+
self._resource_link,
114+
[routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
115+
self._options
114116
)
115117

116118
async def _configure_partition_ranges(self):

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -385,5 +385,7 @@ def _get_target_partition_key_range(self, target_all_ranges):
385385
return list(self._client._ReadPartitionKeyRanges(collection_link=self._resource_link))
386386
query_ranges = self._partitioned_query_ex_info.get_query_ranges()
387387
return self._routing_provider.get_overlapping_ranges(
388-
self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
388+
self._resource_link,
389+
[routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
390+
self._options
389391
)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,5 +152,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
152152
def _get_target_partition_key_range(self):
153153
query_ranges = self._partitioned_query_ex_info.get_query_ranges()
154154
return self._routing_provider.get_overlapping_ranges(
155-
self._resource_link, [routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges]
155+
self._resource_link,
156+
[routing_range.Range.ParseFromDict(range_as_dict) for range_as_dict in query_ranges],
157+
self._options
156158
)

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
"""
2424
from typing import TYPE_CHECKING, Optional
2525

26+
from azure.cosmos._constants import _Constants
2627
from azure.cosmos.partition_key import _get_partition_key_from_partition_key_definition
2728
from azure.cosmos._global_partition_endpoint_manager_circuit_breaker_core import \
2829
_GlobalPartitionEndpointManagerForCircuitBreakerCore
@@ -64,17 +65,20 @@ def create_pk_range_wrapper(self, request: RequestObject) -> Optional[PartitionK
6465
partition_key_definition = properties["partitionKey"]
6566
partition_key = _get_partition_key_from_partition_key_definition(partition_key_definition)
6667

68+
options = {}
69+
if request.excluded_locations:
70+
options[_Constants.Kwargs.EXCLUDED_LOCATIONS] = request.excluded_locations
6771
if HttpHeaders.PartitionKey in request.headers:
6872
partition_key_value = request.headers[HttpHeaders.PartitionKey]
6973
# get the partition key range for the given partition key
7074
epk_range = [partition_key._get_epk_range_for_partition_key(partition_key_value)] # pylint: disable=protected-access
7175
partition_ranges = (self.client._routing_map_provider # pylint: disable=protected-access
72-
.get_overlapping_ranges(container_link, epk_range))
76+
.get_overlapping_ranges(container_link, epk_range, options))
7377
partition_range = Range.PartitionKeyRangeToRange(partition_ranges[0])
7478
elif HttpHeaders.PartitionKeyRangeID in request.headers:
7579
pk_range_id = request.headers[HttpHeaders.PartitionKeyRangeID]
7680
epk_range =(self.client._routing_map_provider # pylint: disable=protected-access
77-
.get_range_by_partition_key_range_id(container_link, pk_range_id))
81+
.get_range_by_partition_key_range_id(container_link, pk_range_id, options))
7882
if not epk_range:
7983
self.global_partition_endpoint_manager_core.log_warn_or_debug(
8084
"Illegal state: partition key range cache not initialized correctly. "

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def __init__(self, client):
5050
# keeps the cached collection routing map by collection id
5151
self._collection_routing_map_by_item = {}
5252

53-
async def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options = None, **kwargs):
53+
async def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options, **kwargs):
5454
"""Given a partition key range and a collection, return the list of
5555
overlapping partition key ranges.
5656
@@ -69,7 +69,7 @@ async def init_collection_routing_map_if_needed(
6969
self,
7070
collection_link: str,
7171
collection_id: str,
72-
feed_options: Optional[Dict[str, Any]] = None,
72+
feed_options: Dict[str, Any],
7373
**kwargs: Dict[str, Any]
7474
):
7575
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(
9191
self,
9292
collection_link: str,
9393
partition_key_range_id: int,
94+
feed_options: Dict[str, Any],
9495
**kwargs: Dict[str, Any]
9596
) -> Optional[Dict[str, Any]]:
9697
collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
97-
await self.init_collection_routing_map_if_needed(collection_link, collection_id, **kwargs)
98+
await self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
9899

99100
return self._collection_routing_map_by_item[collection_id].get_range_by_partition_key_range_id(
100101
partition_key_range_id)

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from . import routing_range
3030
from .routing_range import PartitionKeyRange
3131

32+
3233
# pylint: disable=protected-access
3334

3435

@@ -55,7 +56,7 @@ def init_collection_routing_map_if_needed(
5556
self,
5657
collection_link: str,
5758
collection_id: str,
58-
feed_options: Optional[Dict[str, Any]] = None,
59+
feed_options: Optional[Dict[str, Any]],
5960
**kwargs: Dict[str, Any]
6061
):
6162
collection_routing_map = self._collection_routing_map_by_item.get(collection_id)
@@ -72,7 +73,7 @@ def init_collection_routing_map_if_needed(
7273
)
7374
self._collection_routing_map_by_item[collection_id] = collection_routing_map
7475

75-
def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options = None, **kwargs):
76+
def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_options, **kwargs):
7677
"""Given a partition key range and a collection, return the list of
7778
overlapping partition key ranges.
7879
@@ -91,10 +92,11 @@ def get_range_by_partition_key_range_id(
9192
self,
9293
collection_link: str,
9394
partition_key_range_id: int,
95+
feed_options: Dict[str, Any],
9496
**kwargs: Dict[str, Any]
9597
) -> Optional[Dict[str, Any]]:
9698
collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
97-
self.init_collection_routing_map_if_needed(collection_link, collection_id, **kwargs)
99+
self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
98100

99101
return (self._collection_routing_map_by_item[collection_id]
100102
.get_range_by_partition_key_range_id(partition_key_range_id))

0 commit comments

Comments
 (0)