From c811dd4ae82946929ec6fa28e8a7f4549c4108db Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Wed, 11 Jun 2025 16:35:17 -0700 Subject: [PATCH 01/20] Add feed_range to query_items API --- .../_container_recreate_retry_policy.py | 8 +- .../azure/cosmos/_cosmos_client_connection.py | 34 +++-- .../azure/cosmos/aio/_container.py | 25 ++-- .../aio/_cosmos_client_connection_async.py | 3 +- .../azure-cosmos/azure/cosmos/container.py | 49 ++++--- .../azure/cosmos/partition_key.py | 58 +++++--- .../azure-cosmos/tests/test_change_feed.py | 1 - sdk/cosmos/azure-cosmos/tests/test_config.py | 45 +++++- .../tests/test_query_feed_range.py | 128 ++++++++++++++++++ .../tests/test_query_feed_range_async.py | 128 ++++++++++++++++++ 10 files changed, 407 insertions(+), 72 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py create mode 100644 sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py index 53ee57b8c3f8..d7574e9aa79e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py @@ -28,7 +28,7 @@ from azure.core.pipeline.transport._base import HttpRequest from . import http_constants -from .partition_key import _Empty, _Undefined +from .partition_key import _Empty, _Undefined, PartitionKeyKind # pylint: disable=protected-access @@ -88,7 +88,7 @@ def should_extract_partition_key(self, container_cache: Optional[Dict[str, Any]] if self._headers and http_constants.HttpHeaders.PartitionKey in self._headers: current_partition_key = self._headers[http_constants.HttpHeaders.PartitionKey] partition_key_definition = container_cache["partitionKey"] if container_cache else None - if partition_key_definition and partition_key_definition["kind"] == "MultiHash": + if partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: # A null in the multihash partition key indicates a failure in extracting partition keys # from the document definition return 'null' in current_partition_key @@ -110,7 +110,7 @@ def _extract_partition_key(self, client: Optional[Any], container_cache: Optiona elif options and isinstance(options["partitionKey"], _Empty): new_partition_key = [] # else serialize using json dumps method which apart from regular values will serialize None into null - elif partition_key_definition and partition_key_definition["kind"] == "MultiHash": + elif partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: new_partition_key = json.dumps(options["partitionKey"], separators=(',', ':')) else: new_partition_key = json.dumps([options["partitionKey"]]) @@ -131,7 +131,7 @@ async def _extract_partition_key_async(self, client: Optional[Any], elif isinstance(options["partitionKey"], _Empty): new_partition_key = [] # else serialize using json dumps method which apart from regular values will serialize None into null - elif partition_key_definition and partition_key_definition["kind"] == "MultiHash": + elif partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: new_partition_key = json.dumps(options["partitionKey"], separators=(',', ':')) else: new_partition_key = json.dumps([options["partitionKey"]]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 3ee11eeccdd8..3f1b795fa831 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -60,6 +60,7 @@ from ._base import _build_properties_cache from ._change_feed.change_feed_iterable import ChangeFeedIterable from ._change_feed.change_feed_state import ChangeFeedState +from ._change_feed.feed_range_internal import FeedRangeInternalEpk from ._constants import _Constants as Constants from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy from ._cosmos_responses import CosmosDict, CosmosList @@ -72,8 +73,9 @@ _Undefined, _Empty, PartitionKey, + PartitionKeyKind, _return_undefined_or_empty_partition_key, - NonePartitionKeyValue + NonePartitionKeyValue, ) PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long @@ -3161,19 +3163,25 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: # check if query has prefix partition key isPrefixPartitionQuery = kwargs.pop("isPrefixPartitionQuery", None) - if isPrefixPartitionQuery and "partitionKeyDefinition" in kwargs: + if (("feedRange" in options) or + (isPrefixPartitionQuery and "partitionKeyDefinition" in kwargs)): + last_response_headers = CaseInsensitiveDict() # here get the over lapping ranges - # Default to empty Dictionary, but unlikely to be empty as we first check if we have it in kwargs - pk_properties: Union[PartitionKey, Dict] = kwargs.pop("partitionKeyDefinition", {}) - partition_key_definition = PartitionKey( - path=pk_properties["paths"], - kind=pk_properties["kind"], - version=pk_properties["version"]) - partition_key_value = pk_properties["partition_key"] - feedrangeEPK = partition_key_definition._get_epk_range_for_prefix_partition_key( - partition_key_value - ) # cspell:disable-line + if "feedRange" in options: + feedrangeEPK = FeedRangeInternalEpk.from_json(options.pop("feedRange")).get_normalized_range() + else: + # Default to empty Dictionary, but unlikely to be empty as we first check if we have it in kwargs + pk_properties: Union[PartitionKey, Dict] = kwargs.pop("partitionKeyDefinition", {}) + partition_key_definition = PartitionKey( + path=pk_properties["paths"], + kind=pk_properties["kind"], + version=pk_properties["version"]) + partition_key_value = pk_properties["partition_key"] + feedrangeEPK = partition_key_definition._get_epk_range_for_prefix_partition_key( + partition_key_value + ) # cspell:disable-line + over_lapping_ranges = self._routing_map_provider.get_overlapping_ranges(resource_id, [feedrangeEPK], options) # It is possible to get more than one over lapping range. We need to get the query results for each one @@ -3337,7 +3345,7 @@ def _ExtractPartitionKey( partitionKeyDefinition: Mapping[str, Any], document: Mapping[str, Any] ) -> Union[List[Optional[Union[str, float, bool]]], str, float, bool, _Empty, _Undefined]: - if partitionKeyDefinition["kind"] == "MultiHash": + if partitionKeyDefinition["kind"] == PartitionKeyKind.MULTI_HASH: ret: List[Optional[Union[str, float, bool]]] = [] for partition_key_level in partitionKeyDefinition["paths"]: # Parses the paths into a list of token each representing a property diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index d167a3e469ed..c31d138bb261 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -430,18 +430,19 @@ def query_items( self, query: str, *, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: Optional[int] = None, + continuation_token_limit: Optional[int] = None, enable_scan_in_query: Optional[bool] = None, - populate_query_metrics: Optional[bool] = None, - populate_index_metrics: Optional[bool] = None, - session_token: Optional[str] = None, + feed_range: Optional[Dict[str, Any]] = None, initial_headers: Optional[Dict[str, str]] = None, max_integrated_cache_staleness_in_ms: Optional[int] = None, + max_item_count: Optional[int] = None, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + populate_index_metrics: Optional[bool] = None, + populate_query_metrics: Optional[bool] = None, priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, throughput_bucket: Optional[int] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: @@ -506,12 +507,14 @@ def query_items( :caption: Parameterized query to get all products that have been discontinued: :name: query_items_param """ - if session_token is not None: - kwargs['session_token'] = session_token + if feed_range is not None: + kwargs['feed_range'] = feed_range if initial_headers is not None: kwargs['initial_headers'] = initial_headers if priority is not None: kwargs['priority'] = priority + if session_token is not None: + kwargs['session_token'] = session_token if throughput_bucket is not None: kwargs["throughput_bucket"] = throughput_bucket feed_options = _build_options(kwargs) @@ -765,8 +768,8 @@ def query_items_change_feed( # pylint: disable=unused-argument cast(PartitionKeyType, partition_key_value)) change_feed_state_context["partitionKeyFeedRange"] = self._get_epk_range_for_partition_key( partition_key_value, feed_options) - if "feed_range" in kwargs: - change_feed_state_context["feedRange"] = kwargs.pop('feed_range') + if "feedRange" in kwargs: + change_feed_state_context["feedRange"] = feed_options.pop('feedRange') if "continuation" in feed_options: change_feed_state_context["continuation"] = feed_options.pop("continuation") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index cbcd3ccafba7..4670e1edb06d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -72,6 +72,7 @@ from ..partition_key import ( _Undefined, PartitionKey, + PartitionKeyKind, _return_undefined_or_empty_partition_key, NonePartitionKeyValue, _Empty ) @@ -3182,7 +3183,7 @@ async def _AddPartitionKey(self, collection_link, document, options): # Extracts the partition key from the document using the partitionKey definition def _ExtractPartitionKey(self, partitionKeyDefinition, document): - if partitionKeyDefinition["kind"] == "MultiHash": + if partitionKeyDefinition["kind"] == PartitionKeyKind.MULTI_HASH: ret = [] for partition_key_level in partitionKeyDefinition.get("paths"): # Parses the paths into a list of token each representing a property diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 5647a66a99f7..8819171cfac4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -577,6 +577,7 @@ def query_items_change_feed( change_feed_state_context["startTime"] = kwargs.pop("start_time") container_properties = self._get_properties_with_options(feed_options) + # TODO: validate partition_key and feed_range exclusive check here to avoid any extra API calls if "partition_key" in kwargs: partition_key = kwargs.pop("partition_key") change_feed_state_context["partitionKey"] = self._set_partition_key(cast(PartitionKeyType, partition_key)) @@ -599,6 +600,7 @@ def query_items_change_feed( ) return result + # TODO: add override methods to give hint to users such as exclusive options: partition_key and feed_range @distributed_trace def query_items( # pylint:disable=docstring-missing-param self, @@ -610,14 +612,16 @@ def query_items( # pylint:disable=docstring-missing-param enable_scan_in_query: Optional[bool] = None, populate_query_metrics: Optional[bool] = None, *, - populate_index_metrics: Optional[bool] = None, - session_token: Optional[str] = None, + continuation_token_limit: Optional[int] = None, + #TODO: consider adding `excluded_locations` as kwarg option here + feed_range: Optional[Dict[str, Any]] = None, initial_headers: Optional[Dict[str, str]] = None, max_integrated_cache_staleness_in_ms: Optional[int] = None, + populate_index_metrics: Optional[bool] = None, priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, - throughput_bucket: Optional[int] = None, response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, + throughput_bucket: Optional[int] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Return all results matching the given `query`. @@ -641,27 +645,28 @@ def query_items( # pylint:disable=docstring-missing-param :param bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :param bool populate_query_metrics: Enable returning query metrics in response headers. - :keyword str session_token: Token for use with Session consistency. - :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. - :keyword response_hook: A callable invoked with the response metadata. - :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). + :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations + in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + If all preferred locations were excluded, primary/hub location will be used. + This excluded_location will override existing excluded_locations in client level. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, responses are guaranteed to be no staler than this value. - :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each - request. Once the user has reached their provisioned throughput, low priority requests are throttled - before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used existing indexes and how it could use potential new indexes. Please note that this options will incur overhead, so it should be enabled only when debugging slow queries. + :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + request. Once the user has reached their provisioned throughput, low priority requests are throttled + before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword response_hook: A callable invoked with the response metadata. + :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. :keyword int throughput_bucket: The desired throughput bucket for the client - :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. - If all preferred locations were excluded, primary/hub location will be used. - This excluded_location will override existing excluded_locations in client level. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] @@ -681,12 +686,18 @@ def query_items( # pylint:disable=docstring-missing-param :dedent: 0 :caption: Parameterized query to get all products that have been discontinued: """ - if session_token is not None: - kwargs['session_token'] = session_token + if feed_range is not None: + kwargs['feed_range'] = feed_range + if 'feed_range' in kwargs and partition_key is not None: + warnings.warn("'feed_range' and 'partition_key' are mutually exclusive, " + "if both were given, the 'feed_range' will be used automatically.") + partition_key = None if initial_headers is not None: kwargs['initial_headers'] = initial_headers if priority is not None: kwargs['priority'] = priority + if session_token is not None: + kwargs['session_token'] = session_token if throughput_bucket is not None: kwargs["throughputBucket"] = throughput_bucket feed_options = build_options(kwargs) @@ -699,7 +710,9 @@ def query_items( # pylint:disable=docstring-missing-param if populate_index_metrics is not None: feed_options["populateIndexMetrics"] = populate_index_metrics properties = self._get_properties_with_options(feed_options) - if partition_key is not None: + if "feed_range" in kwargs: + feed_options["feedRange"] = kwargs.pop('feed_range') + elif partition_key is not None: partition_key_value = self._set_partition_key(partition_key) if is_prefix_partition_key(properties, partition_key): kwargs["isPrefixPartitionQuery"] = True diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index abb9b1698ecd..b17e9a89ef0a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -68,6 +68,13 @@ class _PartitionKeyComponentType: Float = 0x14 Infinity = 0xFF +class PartitionKeyKind: + HASH: str = "Hash" + MULTI_HASH: str = "MultiHash" + +class PartitionKeyVersion: + V1: int = 1 + V2: int = 2 class NonePartitionKeyValue: """Represents None value for partitionKey when it's missing in a container. @@ -102,17 +109,22 @@ class PartitionKey(dict): """ @overload - def __init__(self, path: List[str], *, kind: Literal["MultiHash"] = "MultiHash", version: int = 2) -> None: + def __init__(self, path: List[str], *, kind: Literal["MultiHash"] = PartitionKeyKind.MULTI_HASH, + version: int = PartitionKeyVersion.V2 + ) -> None: ... @overload - def __init__(self, path: str, *, kind: Literal["Hash"] = "Hash", version: int = 2) -> None: + def __init__(self, path: str, *, kind: Literal["Hash"] = PartitionKeyKind.HASH, + version:int = PartitionKeyVersion.V2 + ) -> None: ... def __init__(self, *args, **kwargs): path = args[0] if args else kwargs['path'] - kind = args[1] if len(args) > 1 else kwargs.get('kind', 'Hash' if isinstance(path, str) else 'MultiHash') - version = args[2] if len(args) > 2 else kwargs.get('version', 2) + kind = args[1] if len(args) > 1 else kwargs.get('kind', PartitionKeyKind.HASH if isinstance(path, str) + else PartitionKeyKind.MULTI_HASH) + version = args[2] if len(args) > 2 else kwargs.get('version', PartitionKeyVersion.V2) super().__init__(paths=[path] if isinstance(path, str) else path, kind=kind, version=version) def __repr__(self) -> str: @@ -128,7 +140,7 @@ def kind(self, value: Literal["MultiHash", "Hash"]) -> None: @property def path(self) -> str: - if self.kind == "MultiHash": + if self.kind == PartitionKeyKind.MULTI_HASH: return ''.join(self["paths"]) return self["paths"][0] @@ -151,7 +163,7 @@ def _get_epk_range_for_prefix_partition_key( self, pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> _Range: - if self.kind != "MultiHash": + if self.kind != PartitionKeyKind.MULTI_HASH: raise ValueError( "Effective Partition Key Range for Prefix Partition Keys is only supported for Hierarchical Partition Keys.") # pylint: disable=line-too-long len_pk_value = len(pk_value) @@ -224,27 +236,33 @@ def _get_effective_partition_key_for_hash_partitioning( partition_key_components = [hash_value] + truncated_components return _to_hex_encoded_binary_string_v1(partition_key_components) - def _get_effective_partition_key_string( - self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + @staticmethod + def get_hashed_partition_key_string( + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]], + kind: str, + version: int = PartitionKeyVersion.V2, ) -> Union[int, str]: if not pk_value: return _MinimumInclusiveEffectivePartitionKey - if isinstance(self, _Infinity): - return _MaximumExclusiveEffectivePartitionKey - - kind = self.kind - if kind == 'Hash': - version = self.version or 2 - if version == 1: + if kind == PartitionKeyKind.HASH: + if version == PartitionKeyVersion.V1: return PartitionKey._get_effective_partition_key_for_hash_partitioning(pk_value) - if version == 2: + if version == PartitionKeyVersion.V2: return PartitionKey._get_effective_partition_key_for_hash_partitioning_v2(pk_value) - elif kind == 'MultiHash': - return self._get_effective_partition_key_for_multi_hash_partitioning_v2(pk_value) + elif kind == PartitionKeyKind.MULTI_HASH: + return PartitionKey._get_effective_partition_key_for_multi_hash_partitioning_v2(pk_value) return _to_hex_encoded_binary_string(pk_value) + def _get_effective_partition_key_string( + self, + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + ) -> Union[int, str]: + if isinstance(self, _Infinity): + return _MaximumExclusiveEffectivePartitionKey + + return PartitionKey.get_hashed_partition_key_string(pk_value=pk_value, kind=self.kind, version=self.version) + @staticmethod def _write_for_hashing( value: Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]], @@ -331,7 +349,7 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( def _is_prefix_partition_key( self, partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> bool: # pylint: disable=line-too-long - if self.kind != "MultiHash": + if self.kind != PartitionKeyKind.MULTI_HASH: return False ret = ((isinstance(partition_key, Sequence) and not isinstance(partition_key, str)) and len(self['paths']) != len(partition_key)) diff --git a/sdk/cosmos/azure-cosmos/tests/test_change_feed.py b/sdk/cosmos/azure-cosmos/tests/test_change_feed.py index 94713d543003..0d8d5b6ed312 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/tests/test_change_feed.py @@ -51,7 +51,6 @@ def test_get_feed_ranges(self, setup): assert len(result) == 1 @pytest.mark.parametrize("change_feed_filter_param", ["partitionKey", "partitionKeyRangeId", "feedRange"]) - # @pytest.mark.parametrize("change_feed_filter_param", ["partitionKeyRangeId"]) def test_query_change_feed_with_different_filter(self, change_feed_filter_param, setup): created_collection = setup["created_db"].create_container(f"change_feed_test_{change_feed_filter_param}_{str(uuid.uuid4())}", PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/tests/test_config.py b/sdk/cosmos/azure-cosmos/tests/test_config.py index f8c2f7832bdb..93589a35e148 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_config.py +++ b/sdk/cosmos/azure-cosmos/tests/test_config.py @@ -7,17 +7,21 @@ import unittest import uuid +from azure.core.exceptions import AzureError, ServiceRequestError, ServiceResponseError, ClientAuthenticationError +from azure.core.pipeline.policies import AsyncRetryPolicy, RetryPolicy +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._retry_utility import _has_database_account_header, _has_read_retryable_headers, _configure_timeout +from azure.cosmos._routing.routing_range import Range from azure.cosmos.cosmos_client import CosmosClient from azure.cosmos.exceptions import CosmosHttpResponseError from azure.cosmos.http_constants import StatusCodes -from azure.cosmos.partition_key import PartitionKey +from azure.cosmos.partition_key import (PartitionKey, PartitionKeyKind, PartitionKeyVersion, _Undefined, + NonePartitionKeyValue) from azure.cosmos import (ContainerProxy, DatabaseProxy, documents, exceptions, - http_constants, _retry_utility) -from azure.core.exceptions import AzureError, ServiceRequestError, ServiceResponseError, ClientAuthenticationError -from azure.core.pipeline.policies import AsyncRetryPolicy, RetryPolicy + http_constants) from devtools_testutils.azure_recorded_testcase import get_credential from devtools_testutils.helpers import is_live +from typing import Sequence, Type, Union try: import urllib3 @@ -525,3 +529,36 @@ async def send(self, request): self.update_context(response.context, retry_settings) return response + +def hash_partition_key_value( + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]], + kind: str = PartitionKeyKind.HASH, + version: int = PartitionKeyVersion.V2, + ): + return PartitionKey.get_hashed_partition_key_string( + pk_value=pk_value, + kind=kind, + version=version, + ) + +def create_range(range_min: str, range_max: str, is_min_inclusive: bool = True, is_max_inclusive: bool = False): + if range_max == range_min: + range_max += "FF" + return Range( + range_min=range_min, + range_max=range_max, + isMinInclusive=is_min_inclusive, + isMaxInclusive=is_max_inclusive, + ) + +def create_feed_range_in_dict(feed_range): + return FeedRangeInternalEpk(feed_range).to_dict() + +def create_feed_range_between_pk_values(pk1, pk2): + range_min = hash_partition_key_value([pk1]) + range_max = hash_partition_key_value([pk2]) + if range_min > range_max: + range_min, range_max = range_max, range_min + range_max += "FF" + new_range = create_range(range_min, range_max) + return FeedRangeInternalEpk(new_range).to_dict() \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py new file mode 100644 index 000000000000..52e3506b9b0b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -0,0 +1,128 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import os +import unittest +import uuid + +import pytest + +from azure.cosmos import CosmosClient +import test_config + +from itertools import combinations +from typing import List, Mapping, Set + +CONFIG = test_config.TestConfig() +HOST = CONFIG.host +KEY = CONFIG.credential +DATABASE_ID = CONFIG.TEST_DATABASE_ID +SINGLE_PARTITION_CONTAINER_ID = CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID +MULTI_PARTITION_CONTAINER_ID = CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID +TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] +PK_VALUES = ('pk1', 'pk2', 'pk3') +def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: + if len(items) == 0: + return + + pk_values = [item['pk'] for item in items] + pk_value_set.update(pk_values) + +@pytest.fixture(scope="class", autouse=True) +def setup_and_teardown(): + print("Setup: This runs before any tests") + document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] + + database = CosmosClient(HOST, KEY).get_database_client(DATABASE_ID) + for container_id in TEST_CONTAINERS_IDS: + container = database.get_container_client(container_id) + for document_definition in document_definitions: + container.upsert_item(body=document_definition) + yield + # Code to run after tests + print("Teardown: This runs after all tests") + +def get_container(container_id: str): + client = CosmosClient(HOST, KEY) + db = client.get_database_client(DATABASE_ID) + return db.get_container_client(container_id) + +@pytest.mark.cosmosQuery +class TestQueryFeedRange(): + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_all_partitions(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + expected_pk_values = set(PK_VALUES) + actual_pk_values = set() + iter_feed_ranges = list(container.read_feed_ranges()) + for feed_range in iter_feed_ranges: + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values == actual_pk_values + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_single_partition_key(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + for pk_value in PK_VALUES: + expected_pk_values = {pk_value} + actual_pk_values = set() + + feed_range = test_config.create_feed_range_between_pk_values(pk_value, pk_value) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values == actual_pk_values + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_multiple_partition_key(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + for pk_value1, pk_value2 in combinations(PK_VALUES, 2): + expected_pk_values = {pk_value1, pk_value2} + actual_pk_values = set() + + feed_range = test_config.create_feed_range_between_pk_values(pk_value1, pk_value2) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values.issubset(actual_pk_values) + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_a_full_range(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + expected_pk_values = set(PK_VALUES) + actual_pk_values = set() + new_range = test_config.create_range( + range_min="", + range_max="FF", + is_min_inclusive=True, + is_max_inclusive=False, + ) + feed_range = test_config.create_feed_range_in_dict(new_range) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values.issubset(actual_pk_values) + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py new file mode 100644 index 000000000000..52e3506b9b0b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py @@ -0,0 +1,128 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import os +import unittest +import uuid + +import pytest + +from azure.cosmos import CosmosClient +import test_config + +from itertools import combinations +from typing import List, Mapping, Set + +CONFIG = test_config.TestConfig() +HOST = CONFIG.host +KEY = CONFIG.credential +DATABASE_ID = CONFIG.TEST_DATABASE_ID +SINGLE_PARTITION_CONTAINER_ID = CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID +MULTI_PARTITION_CONTAINER_ID = CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID +TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] +PK_VALUES = ('pk1', 'pk2', 'pk3') +def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: + if len(items) == 0: + return + + pk_values = [item['pk'] for item in items] + pk_value_set.update(pk_values) + +@pytest.fixture(scope="class", autouse=True) +def setup_and_teardown(): + print("Setup: This runs before any tests") + document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] + + database = CosmosClient(HOST, KEY).get_database_client(DATABASE_ID) + for container_id in TEST_CONTAINERS_IDS: + container = database.get_container_client(container_id) + for document_definition in document_definitions: + container.upsert_item(body=document_definition) + yield + # Code to run after tests + print("Teardown: This runs after all tests") + +def get_container(container_id: str): + client = CosmosClient(HOST, KEY) + db = client.get_database_client(DATABASE_ID) + return db.get_container_client(container_id) + +@pytest.mark.cosmosQuery +class TestQueryFeedRange(): + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_all_partitions(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + expected_pk_values = set(PK_VALUES) + actual_pk_values = set() + iter_feed_ranges = list(container.read_feed_ranges()) + for feed_range in iter_feed_ranges: + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values == actual_pk_values + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_single_partition_key(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + for pk_value in PK_VALUES: + expected_pk_values = {pk_value} + actual_pk_values = set() + + feed_range = test_config.create_feed_range_between_pk_values(pk_value, pk_value) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values == actual_pk_values + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_multiple_partition_key(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + for pk_value1, pk_value2 in combinations(PK_VALUES, 2): + expected_pk_values = {pk_value1, pk_value2} + actual_pk_values = set() + + feed_range = test_config.create_feed_range_between_pk_values(pk_value1, pk_value2) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values.issubset(actual_pk_values) + + @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) + def test_query_with_feed_range_for_a_full_range(self, container_id): + container = get_container(container_id) + query = 'SELECT * from c' + + expected_pk_values = set(PK_VALUES) + actual_pk_values = set() + new_range = test_config.create_range( + range_min="", + range_max="FF", + is_min_inclusive=True, + is_max_inclusive=False, + ) + feed_range = test_config.create_feed_range_in_dict(new_range) + items = list(container.query_items( + query=query, + enable_cross_partition_query=True, + feed_range=feed_range + )) + add_all_pk_values_to_set(items, actual_pk_values) + assert expected_pk_values.issubset(actual_pk_values) + +if __name__ == "__main__": + unittest.main() From 40e61dd4dfa151e28dd18d4e6098826399c371c3 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Wed, 18 Jun 2025 14:19:18 -0700 Subject: [PATCH 02/20] Added overload methods for 'query_items' API --- .../azure-cosmos/azure/cosmos/_utils.py | 29 +- .../azure-cosmos/azure/cosmos/container.py | 299 +++++++++++++----- 2 files changed, 256 insertions(+), 72 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index f1899415af87..ddd36ebc0e63 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -27,7 +27,7 @@ import base64 import json import time -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Tuple from ._version import VERSION @@ -76,3 +76,30 @@ def get_index_metrics_info(delimited_string: Optional[str]) -> Dict[str, Any]: def current_time_millis() -> int: return int(round(time.time() * 1000)) + +def add_args_to_kwargs( + arg_names: Tuple[str, ...], + args: Tuple[Any, ...], + kwargs: Dict[str, Any] + ) -> None: + """Add positional arguments(args) to keyword argument dictionary(kwargs) using names in arg_names as keys. + """ + + if len(args) > len(arg_names): + raise IndexError(f"Positional argument is out of range. Expected {len(arg_names)} arguments, " + f"but got {len(args)} instead. Please review argument list in API documentation.") + + for name, arg in zip(arg_names, args): + if name in kwargs: + raise KeyError(f"{name} cannot be used as positional and keyword argument at the same time.") + + kwargs[name] = arg + +def verify_exclusive_arguments(exclusive_keys, **kwargs: Any) -> None: + """Verify if exclusive arguments are present in kwargs. + """ + count = sum(1 for key in exclusive_keys if key in kwargs and kwargs[key] is not None) + + if count > 1: + raise ValueError( + "partition_key_range_id, partition_key, feed_range are exclusive parameters, please only set one of them") \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 8819171cfac4..70a172d8bc7c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -31,6 +31,7 @@ from azure.core.tracing.decorator import distributed_trace from azure.cosmos._change_feed.change_feed_utils import add_args_to_kwargs, validate_kwargs +from . import _utils as utils from ._base import ( build_options, validate_cache_staleness_value, @@ -600,30 +601,27 @@ def query_items_change_feed( ) return result - # TODO: add override methods to give hint to users such as exclusive options: partition_key and feed_range - @distributed_trace - def query_items( # pylint:disable=docstring-missing-param - self, - query: str, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - enable_cross_partition_query: Optional[bool] = None, - max_item_count: Optional[int] = None, - enable_scan_in_query: Optional[bool] = None, - populate_query_metrics: Optional[bool] = None, - *, - continuation_token_limit: Optional[int] = None, - #TODO: consider adding `excluded_locations` as kwarg option here - feed_range: Optional[Dict[str, Any]] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - populate_index_metrics: Optional[bool] = None, - priority: Optional[Literal["High", "Low"]] = None, - response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, - session_token: Optional[str] = None, - throughput_bucket: Optional[int] = None, - **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: + @overload + def query_items( + self, + query: str, + *, + continuation_token_limit: Optional[int] = None, + enable_cross_partition_query: Optional[bool] = None, + enable_scan_in_query: Optional[bool] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + max_item_count: Optional[int] = None, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + populate_index_metrics: Optional[bool] = None, + populate_query_metrics: Optional[bool] = None, + priority: Optional[Literal["High", "Low"]] = None, + response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, + throughput_bucket: Optional[int] = None, + **kwargs: Any + ): """Return all results matching the given `query`. You can use any value for the container name in the FROM clause, but @@ -631,23 +629,166 @@ def query_items( # pylint:disable=docstring-missing-param name is "products," and is aliased as "p" for easier referencing in the WHERE clause. - :param str query: The Azure Cosmos DB SQL query to execute. - :param parameters: Optional array of parameters to the query. + :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query + response. Valid values are positive integers. + A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_cross_partition_query: Allows sending of more than one request to + execute the query in the Azure Cosmos DB service. + More than one request is necessary if the query is not scoped to single partition key value. + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + indexing was opted out on the requested paths. + :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations + in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + If all preferred locations were excluded, primary/hub location will be used. + This excluded_location will override existing excluded_locations in client level. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. + :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in + milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, + responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword parameters: Optional array of parameters to the query. Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. - :type parameters: [List[Dict[str, object]]] - :param partition_key: partition key at which the query request is targeted. - :type partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] - :param bool enable_cross_partition_query: Allows sending of more than one request to + :paramtype parameters: [List[Dict[str, object]]] + :keyword partition_key: partition key at which the query request is targeted. + :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] + :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used + existing indexes and how it could use potential new indexes. Please note that this options will incur + overhead, so it should be enabled only when debugging slow queries. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. + :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + request. Once the user has reached their provisioned throughput, low priority requests are throttled + before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword str query: The Azure Cosmos DB SQL query to execute. + :keyword response_hook: A callable invoked with the response metadata. + :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. + :keyword int throughput_bucket: The desired throughput bucket for the client + :returns: An Iterable of items (dicts). + :rtype: ItemPaged[Dict[str, Any]] + + .. admonition:: Example: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items] + :end-before: [END query_items] + :language: python + :dedent: 0 + :caption: Get all products that have not been discontinued: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items_param] + :end-before: [END query_items_param] + :language: python + :dedent: 0 + :caption: Parameterized query to get all products that have been discontinued: + """ + ... + + @overload + def query_items( + self, + query: str, + *, + continuation_token_limit: Optional[int] = None, + enable_cross_partition_query: Optional[bool] = None, + enable_scan_in_query: Optional[bool] = None, + feed_range: Optional[Dict[str, Any]] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + max_item_count: Optional[int] = None, + parameters: Optional[List[Dict[str, object]]] = None, + populate_index_metrics: Optional[bool] = None, + populate_query_metrics: Optional[bool] = None, + priority: Optional[Literal["High", "Low"]] = None, + response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, + throughput_bucket: Optional[int] = None, + **kwargs: Any + ): + """Return all results matching the given `query`. + + You can use any value for the container name in the FROM clause, but + often the container name is used. In the examples below, the container + name is "products," and is aliased as "p" for easier referencing in + the WHERE clause. + + :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query + response. Valid values are positive integers. + A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_cross_partition_query: Allows sending of more than one request to execute the query in the Azure Cosmos DB service. More than one request is necessary if the query is not scoped to single partition key value. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. - :param bool populate_query_metrics: Enable returning query metrics in response headers. + :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations + in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + If all preferred locations were excluded, primary/hub location will be used. + This excluded_location will override existing excluded_locations in client level. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. + :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in + milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, + responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword parameters: Optional array of parameters to the query. + Each parameter is a dict() with 'name' and 'value' keys. + Ignored if no query is provided. + :paramtype parameters: [List[Dict[str, object]]] + :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used + existing indexes and how it could use potential new indexes. Please note that this options will incur + overhead, so it should be enabled only when debugging slow queries. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. + :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + request. Once the user has reached their provisioned throughput, low priority requests are throttled + before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword str query: The Azure Cosmos DB SQL query to execute. + :keyword response_hook: A callable invoked with the response metadata. + :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. + :keyword int throughput_bucket: The desired throughput bucket for the client + :returns: An Iterable of items (dicts). + :rtype: ItemPaged[Dict[str, Any]] + + .. admonition:: Example: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items] + :end-before: [END query_items] + :language: python + :dedent: 0 + :caption: Get all products that have not been discontinued: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items_param] + :end-before: [END query_items_param] + :language: python + :dedent: 0 + :caption: Parameterized query to get all products that have been discontinued: + """ + ... + + @distributed_trace + def query_items( # pylint:disable=docstring-missing-param + self, + *args: Any, + **kwargs: Any + ) -> ItemPaged[Dict[str, Any]]: + """Return all results matching the given `query`. + + You can use any value for the container name in the FROM clause, but + often the container name is used. In the examples below, the container + name is "products," and is aliased as "p" for easier referencing in + the WHERE clause. + :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_cross_partition_query: Allows sending of more than one request to + execute the query in the Azure Cosmos DB service. + More than one request is necessary if the query is not scoped to single partition key value. + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. @@ -657,12 +798,21 @@ def query_items( # pylint:disable=docstring-missing-param :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword parameters: Optional array of parameters to the query. + Each parameter is a dict() with 'name' and 'value' keys. + Ignored if no query is provided. + :paramtype parameters: [List[Dict[str, object]]] + :keyword partition_key: partition key at which the query request is targeted. + :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used existing indexes and how it could use potential new indexes. Please note that this options will incur overhead, so it should be enabled only when debugging slow queries. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword str query: The Azure Cosmos DB SQL query to execute. :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. @@ -686,55 +836,62 @@ def query_items( # pylint:disable=docstring-missing-param :dedent: 0 :caption: Parameterized query to get all products that have been discontinued: """ - if feed_range is not None: - kwargs['feed_range'] = feed_range - if 'feed_range' in kwargs and partition_key is not None: - warnings.warn("'feed_range' and 'partition_key' are mutually exclusive, " - "if both were given, the 'feed_range' will be used automatically.") - partition_key = None - if initial_headers is not None: - kwargs['initial_headers'] = initial_headers - if priority is not None: - kwargs['priority'] = priority - if session_token is not None: - kwargs['session_token'] = session_token - if throughput_bucket is not None: - kwargs["throughputBucket"] = throughput_bucket + # Add positional arguments to keyword argument to support backward compatibility. + original_positional_arg_names = ("query", "parameters", "partition_key", "enable_cross_partition_query", + "max_item_count", "enable_scan_in_query", "populate_query_metrics") + utils.add_args_to_kwargs(original_positional_arg_names, args, kwargs) feed_options = build_options(kwargs) - if enable_cross_partition_query is not None: - feed_options["enableCrossPartitionQuery"] = enable_cross_partition_query - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - if populate_index_metrics is not None: - feed_options["populateIndexMetrics"] = populate_index_metrics - properties = self._get_properties_with_options(feed_options) + container_properties = self._get_properties_with_options(feed_options) + + # Update 'feed_options' from 'kwargs' + if "enable_cross_partition_query" in kwargs: + feed_options["enableCrossPartitionQuery"] = kwargs.pop("enable_cross_partition_query") + if "max_item_count" in kwargs: + feed_options["maxItemCount"] = kwargs.pop("max_item_count") + if "populate_query_metrics" in kwargs: + feed_options["populateQueryMetrics"] = kwargs.pop("populate_query_metrics") + if "populate_index_metrics" in kwargs: + feed_options["populateIndexMetrics"] = kwargs.pop("populate_index_metrics") + if "enable_scan_in_query" in kwargs: + feed_options["enableScanInQuery"] = kwargs.pop("enable_scan_in_query") + if "max_integrated_cache_staleness_in_ms" in kwargs: + max_integrated_cache_staleness_in_ms = kwargs.pop("max_integrated_cache_staleness_in_ms") + validate_cache_staleness_value(max_integrated_cache_staleness_in_ms) + feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms + if "continuation_token_limit" in kwargs: + feed_options["responseContinuationTokenLimitInKb"] = kwargs.pop("continuation_token_limit") + feed_options["correlatedActivityId"] = GenerateGuidId() + feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + + # Set query with 'query' and 'parameters' from kwargs + if "parameters" in kwargs: + query = {"query": kwargs.pop("query", None), "parameters": kwargs.pop("parameters", None)} + else: + query = kwargs.pop("query", None) + + # Set range filters for query. Options are either 'feed_range' or 'partition_key' + utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) + partition_key = None if "feed_range" in kwargs: - feed_options["feedRange"] = kwargs.pop('feed_range') - elif partition_key is not None: + feed_options["feedRange"] = kwargs.pop("feed_range", None) + elif "partition_key" in kwargs: + partition_key = kwargs.pop("partition_key") partition_key_value = self._set_partition_key(partition_key) - if is_prefix_partition_key(properties, partition_key): + if is_prefix_partition_key(container_properties, partition_key): kwargs["isPrefixPartitionQuery"] = True - kwargs["partitionKeyDefinition"] = properties["partitionKey"] + kwargs["partitionKeyDefinition"] = container_properties["partitionKey"] kwargs["partitionKeyDefinition"]["partition_key"] = partition_key_value else: feed_options["partitionKey"] = partition_key_value - if enable_scan_in_query is not None: - feed_options["enableScanInQuery"] = enable_scan_in_query - if max_integrated_cache_staleness_in_ms: - validate_cache_staleness_value(max_integrated_cache_staleness_in_ms) - feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms - correlated_activity_id = GenerateGuidId() - feed_options["correlatedActivityId"] = correlated_activity_id - if continuation_token_limit is not None: - feed_options["responseContinuationTokenLimitInKb"] = continuation_token_limit + + # Set 'response_hook' + response_hook = kwargs.pop("response_hook", None) if response_hook and hasattr(response_hook, "clear"): response_hook.clear() - feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + items = self.client_connection.QueryItems( database_or_container_link=self.container_link, - query=query if parameters is None else {"query": query, "parameters": parameters}, + query=query, options=feed_options, partition_key=partition_key, response_hook=response_hook, From 53cf35727571be371882377d577af12d798d5273 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Thu, 19 Jun 2025 12:06:45 -0700 Subject: [PATCH 03/20] Added Tests --- .../azure-cosmos/azure/cosmos/_utils.py | 28 ++++----- .../azure-cosmos/azure/cosmos/container.py | 4 +- sdk/cosmos/azure-cosmos/tests/test_query.py | 43 ++++++++++++++ sdk/cosmos/azure-cosmos/tests/test_utils.py | 59 +++++++++++++++++++ 4 files changed, 118 insertions(+), 16 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index ddd36ebc0e63..dfab4dbfedba 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -27,12 +27,12 @@ import base64 import json import time -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from ._version import VERSION -def get_user_agent(suffix: Optional[str]) -> str: +def get_user_agent(suffix: Optional[str] = None) -> str: os_name = safe_user_agent_header(platform.platform()) python_version = safe_user_agent_header(platform.python_version()) user_agent = "azsdk-python-cosmos/{} Python/{} ({})".format(VERSION, python_version, os_name) @@ -40,7 +40,7 @@ def get_user_agent(suffix: Optional[str]) -> str: user_agent += f" {suffix}" return user_agent -def get_user_agent_async(suffix: Optional[str]) -> str: +def get_user_agent_async(suffix: Optional[str] = None) -> str: os_name = safe_user_agent_header(platform.platform()) python_version = safe_user_agent_header(platform.python_version()) user_agent = "azsdk-python-cosmos-async/{} Python/{} ({})".format(VERSION, python_version, os_name) @@ -49,7 +49,7 @@ def get_user_agent_async(suffix: Optional[str]) -> str: return user_agent -def safe_user_agent_header(s: Optional[str]) -> str: +def safe_user_agent_header(s: Optional[str] = None) -> str: if s is None: s = "unknown" # remove all white spaces @@ -59,7 +59,7 @@ def safe_user_agent_header(s: Optional[str]) -> str: return s -def get_index_metrics_info(delimited_string: Optional[str]) -> Dict[str, Any]: +def get_index_metrics_info(delimited_string: Optional[str] = None) -> Dict[str, Any]: if delimited_string is None: return {} try: @@ -78,7 +78,7 @@ def current_time_millis() -> int: return int(round(time.time() * 1000)) def add_args_to_kwargs( - arg_names: Tuple[str, ...], + arg_names: List[str], args: Tuple[Any, ...], kwargs: Dict[str, Any] ) -> None: @@ -86,20 +86,20 @@ def add_args_to_kwargs( """ if len(args) > len(arg_names): - raise IndexError(f"Positional argument is out of range. Expected {len(arg_names)} arguments, " + raise ValueError(f"Positional argument is out of range. Expected {len(arg_names)} arguments, " f"but got {len(args)} instead. Please review argument list in API documentation.") for name, arg in zip(arg_names, args): if name in kwargs: - raise KeyError(f"{name} cannot be used as positional and keyword argument at the same time.") - + raise ValueError(f"{name} cannot be used as positional and keyword argument at the same time.") kwargs[name] = arg -def verify_exclusive_arguments(exclusive_keys, **kwargs: Any) -> None: +def verify_exclusive_arguments( + exclusive_keys: List[str], + **kwargs: Dict[str, Any]) -> None: """Verify if exclusive arguments are present in kwargs. """ - count = sum(1 for key in exclusive_keys if key in kwargs and kwargs[key] is not None) + keys_in_kwargs = [key for key in exclusive_keys if key in kwargs and kwargs[key] is not None] - if count > 1: - raise ValueError( - "partition_key_range_id, partition_key, feed_range are exclusive parameters, please only set one of them") \ No newline at end of file + if len(keys_in_kwargs) > 1: + raise ValueError(f"{', '.join(keys_in_kwargs)} are exclusive parameters, please only set one of them") \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 70a172d8bc7c..5a7ff1e190d3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -837,8 +837,8 @@ def query_items( # pylint:disable=docstring-missing-param :caption: Parameterized query to get all products that have been discontinued: """ # Add positional arguments to keyword argument to support backward compatibility. - original_positional_arg_names = ("query", "parameters", "partition_key", "enable_cross_partition_query", - "max_item_count", "enable_scan_in_query", "populate_query_metrics") + original_positional_arg_names = ["query", "parameters", "partition_key", "enable_cross_partition_query", + "max_item_count", "enable_scan_in_query", "populate_query_metrics"] utils.add_args_to_kwargs(original_positional_arg_names, args, kwargs) feed_options = build_options(kwargs) container_properties = self._get_properties_with_options(feed_options) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py index aa17116b2f39..c2889ad03dad 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query.py @@ -568,6 +568,49 @@ def test_query_request_params_none_retry_policy(self): retry_utility.ExecuteFunction = self.OriginalExecuteFunction self.created_db.delete_container(created_collection.id) + def test_query_positional_args(self): + container_id = "Multi Partition Test Container Query Positional Args " + str(uuid.uuid4()) + partition_key = "pk" + partition_key_value1 = "pk1" + partition_key_value2 = "pk2" + container = self.created_db.create_container_if_not_exists( + id=container_id, + partition_key=PartitionKey(path='/' + partition_key, kind='Hash'), + offer_throughput=self.config.THROUGHPUT_FOR_5_PARTITIONS, + ) + + num_items = 10 + new_items = [] + for pk_value in [partition_key_value1, partition_key_value2]: + for i in range(num_items): + new_items.append({'pk': pk_value, 'id': f"{pk_value}_{i}", 'name': 'sample name'}) + + for item in new_items: + container.upsert_item(body=item) + + query = "SELECT * FROM root r WHERE r.name=@name" + parameters = [{'name': '@name', 'value': 'sample name'}] + partition_key = partition_key_value2 + max_item_count = 3 + pager = container.query_items( + query, + parameters, + partition_key, + True, + max_item_count, + True, + True, + ).by_page() + + ids = [] + for page in pager: + items = list(page) + num_items = len(items) + for item in items: + assert item['pk'] == partition_key + ids.append(item['id']) + assert num_items <= max_item_count + assert ids == [item['id'] for item in new_items if item['pk'] == partition_key] def _MockExecuteFunctionSessionRetry(self, function, *args, **kwargs): if args: diff --git a/sdk/cosmos/azure-cosmos/tests/test_utils.py b/sdk/cosmos/azure-cosmos/tests/test_utils.py index 52e155748a77..9586b0548af4 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_utils.py +++ b/sdk/cosmos/azure-cosmos/tests/test_utils.py @@ -5,6 +5,8 @@ import unittest import uuid +import pytest + import azure.cosmos import azure.cosmos._utils as _utils import test_config @@ -34,6 +36,63 @@ def test_connection_string(self): self.assertTrue(db is not None) client.delete_database(db.id) + def test_add_args_to_kwargs(self): + arg_names = ["arg1", "arg2", "arg3", "arg4"] + args = ("arg1_val", "arg2_val", "arg3_val", "arg4_val") + kwargs = {} + + # Test any number of positional arguments less than or equals to the number of argument names + for num_args in range(len(arg_names)): + args = tuple(f"arg{i+1}_val" for i in range(num_args)) + kwargs = {} + _utils.add_args_to_kwargs(arg_names, args, kwargs) + + assert len(kwargs.keys()) == len(args) + for arg_name, arg in zip(arg_names, args): + assert arg_name in kwargs + assert kwargs[arg_name] == arg + + # test if arg_name already in kwargs + with pytest.raises(ValueError) as e: + _utils.add_args_to_kwargs(arg_names, args, kwargs) + assert str(e.value) == f"{arg_names[0]} cannot be used as positional and keyword argument at the same time." + + # Test if number of positional argument greater than expected argument names + args = ("arg1_val", "arg2_val", "arg3_val", "arg4_val", "arg5_val") + with pytest.raises(ValueError) as e: + _utils.add_args_to_kwargs(arg_names, args, kwargs) + assert str(e.value) == (f"Positional argument is out of range. Expected {len(arg_names)} arguments, " + f"but got {len(args)} instead. Please review argument list in API documentation.") + + def test_verify_exclusive_arguments(self): + exclusive_keys = ["key1", "key2", "key3", "key4"] + + ## Test valid cases + kwargs = {} + assert _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) is None + + kwargs = {"key1": "test_value"} + assert _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) is None + + kwargs = {"key1": "test_value", "key9": "test_value"} + assert _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) is None + + # Even if some keys are in exclusive_keys list, if the values were 'None' we ignore them + kwargs = {"key1": "test_value", "key2": None, "key3": None} + assert _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) is None + + ## Test invalid cases + kwargs = {"key1": "test_value", "key2": "test_value"} + expected_error_message = "key1, key2 are exclusive parameters, please only set one of them" + with pytest.raises(ValueError) as e: + _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) + assert str(e.value) == expected_error_message + + kwargs = {"key1": "test_value", "key2": "test_value", "key3": "test_value"} + expected_error_message = "key1, key2, key3 are exclusive parameters, please only set one of them" + with pytest.raises(ValueError) as e: + _utils.verify_exclusive_arguments(exclusive_keys, **kwargs) + assert str(e.value) == expected_error_message if __name__ == "__main__": unittest.main() From 609274fb8b030d66fa725ad3d95122bbac906a63 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Thu, 19 Jun 2025 15:32:28 -0700 Subject: [PATCH 04/20] Added feed_range to query_item for async --- .../azure/cosmos/_cosmos_client_connection.py | 40 +-- .../azure-cosmos/azure/cosmos/_utils.py | 13 +- .../azure/cosmos/aio/_container.py | 284 ++++++++++++++---- .../aio/_cosmos_client_connection_async.py | 47 ++- .../azure-cosmos/azure/cosmos/container.py | 44 ++- .../azure/cosmos/partition_key.py | 19 +- .../tests/test_query_feed_range.py | 7 +- .../tests/test_query_feed_range_async.py | 97 +++--- 8 files changed, 356 insertions(+), 195 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 3f1b795fa831..e26b96c50476 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -1310,7 +1310,7 @@ def UpsertItem( self, database_or_container_link: str, document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, + options: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> CosmosDict: """Upserts a document in a collection. @@ -3161,28 +3161,20 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: request_params = RequestObject(resource_type, documents._OperationType.SqlQuery, req_headers) request_params.set_excluded_location_from_options(options) - # check if query has prefix partition key - isPrefixPartitionQuery = kwargs.pop("isPrefixPartitionQuery", None) - if (("feedRange" in options) or - (isPrefixPartitionQuery and "partitionKeyDefinition" in kwargs)): - + # Check if the over lapping ranges can be populated + feed_range_epk = None + if "feed_range" in kwargs: + feed_range = kwargs.pop("feed_range") + feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range() + elif "prefix_partition_key_object" in kwargs and "prefix_partition_key_value" in kwargs: + prefix_partition_key_obj = kwargs.pop("prefix_partition_key_object") + prefix_partition_key_value = kwargs.pop("prefix_partition_key_value") + feed_range_epk = prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value) + + # If feed_range_epk exist, query with the range + if feed_range_epk is not None: last_response_headers = CaseInsensitiveDict() - # here get the over lapping ranges - if "feedRange" in options: - feedrangeEPK = FeedRangeInternalEpk.from_json(options.pop("feedRange")).get_normalized_range() - else: - # Default to empty Dictionary, but unlikely to be empty as we first check if we have it in kwargs - pk_properties: Union[PartitionKey, Dict] = kwargs.pop("partitionKeyDefinition", {}) - partition_key_definition = PartitionKey( - path=pk_properties["paths"], - kind=pk_properties["kind"], - version=pk_properties["version"]) - partition_key_value = pk_properties["partition_key"] - feedrangeEPK = partition_key_definition._get_epk_range_for_prefix_partition_key( - partition_key_value - ) # cspell:disable-line - - over_lapping_ranges = self._routing_map_provider.get_overlapping_ranges(resource_id, [feedrangeEPK], + over_lapping_ranges = self._routing_map_provider.get_overlapping_ranges(resource_id, [feed_range_epk], options) # It is possible to get more than one over lapping range. We need to get the query results for each one results: Dict[str, Any] = {} @@ -3199,8 +3191,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: single_range = routing_range.Range.PartitionKeyRangeToRange(over_lapping_range) # Since the range min and max are all Upper Cased string Hex Values, # we can compare the values lexicographically - EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feedrangeEPK.min), - range_max=min(single_range.max, feedrangeEPK.max), + EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feed_range_epk.min), + range_max=min(single_range.max, feed_range_epk.max), isMinInclusive=True, isMaxInclusive=False) if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max: # The Epk Sub Range spans exactly one physical partition diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index dfab4dbfedba..9d6f39e998fd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -83,6 +83,12 @@ def add_args_to_kwargs( kwargs: Dict[str, Any] ) -> None: """Add positional arguments(args) to keyword argument dictionary(kwargs) using names in arg_names as keys. + To be backward-compatible, some expected positional arguments has to be allowed. This method will verify number of + maximum positional arguments and add them to the keyword argument dictionary(kwargs) + + :param List[str] arg_names: The names of positional arguments. + :param Tuple[Any, ...] args: The tuple of positional arguments. + :param Dict[str, Any] kwargs: The dictionary of keyword arguments as reference. This dictionary will be updated. """ if len(args) > len(arg_names): @@ -98,8 +104,13 @@ def verify_exclusive_arguments( exclusive_keys: List[str], **kwargs: Dict[str, Any]) -> None: """Verify if exclusive arguments are present in kwargs. + For some Cosmos SDK APIs, some arguments are exclusive, or cannot be used at the same time. This method will verify + that and raise an error if exclusive arguments are present. + + :param List[str] exclusive_keys: The names of exclusive arguments. + :param Dict[str, Any] kwargs: The dictionary of keyword arguments. """ keys_in_kwargs = [key for key in exclusive_keys if key in kwargs and kwargs[key] is not None] if len(keys_in_kwargs) > 1: - raise ValueError(f"{', '.join(keys_in_kwargs)} are exclusive parameters, please only set one of them") \ No newline at end of file + raise ValueError(f"{', '.join(keys_in_kwargs)} are exclusive parameters, please only set one of them") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index c31d138bb261..02704403d062 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -34,6 +34,7 @@ from ._cosmos_client_connection_async import CosmosClientConnection from ._scripts import ScriptsProxy +from .. import _utils as utils from .._base import ( build_options as _build_options, validate_cache_staleness_value, @@ -425,27 +426,26 @@ def read_all_items( ) return items - @distributed_trace + @overload def query_items( - self, - query: str, - *, - continuation_token_limit: Optional[int] = None, - enable_scan_in_query: Optional[bool] = None, - feed_range: Optional[Dict[str, Any]] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - max_item_count: Optional[int] = None, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - populate_index_metrics: Optional[bool] = None, - populate_query_metrics: Optional[bool] = None, - priority: Optional[Literal["High", "Low"]] = None, - response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, - session_token: Optional[str] = None, - throughput_bucket: Optional[int] = None, - **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: + self, + query: str, + *, + continuation_token_limit: Optional[int] = None, + enable_scan_in_query: Optional[bool] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + max_item_count: Optional[int] = None, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + populate_index_metrics: Optional[bool] = None, + populate_query_metrics: Optional[bool] = None, + priority: Optional[Literal["High", "Low"]] = None, + response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, + throughput_bucket: Optional[int] = None, + **kwargs: Any + ): """Return all results matching the given `query`. You can use any value for the container name in the FROM clause, but @@ -454,98 +454,252 @@ def query_items( the WHERE clause. :param str query: The Azure Cosmos DB SQL query to execute. + :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query + response. Valid values are positive integers. + A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + indexing was opted out on the requested paths. + :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations + in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + If all preferred locations were excluded, primary/hub location will be used. + This excluded_location will override existing excluded_locations in client level. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. + :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in + milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, + responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword parameters: Optional array of parameters to the query. Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. - :paramtype parameters: List[Dict[str, Any]] - :keyword partition_key: Specifies the partition key value for the item. If none is provided, - a cross-partition query will be executed. + :paramtype parameters: [List[Dict[str, object]]] + :keyword partition_key: partition key at which the query request is targeted. :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] - :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as - indexing was opted out on the requested paths. - :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used existing indexes and how it could use potential new indexes. Please note that this options will incur overhead, so it should be enabled only when debugging slow queries. - :keyword str session_token: Token for use with Session consistency. - :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. + :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + request. Once the user has reached their provisioned throughput, low priority requests are throttled + before high priority requests start getting throttled. Feature must first be enabled at the account level. :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. + :keyword int throughput_bucket: The desired throughput bucket for the client + :returns: An Iterable of items (dicts). + :rtype: ItemPaged[Dict[str, Any]] + + .. admonition:: Example: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items] + :end-before: [END query_items] + :language: python + :dedent: 0 + :caption: Get all products that have not been discontinued: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items_param] + :end-before: [END query_items_param] + :language: python + :dedent: 0 + :caption: Parameterized query to get all products that have been discontinued: + """ + ... + + @overload + def query_items( + self, + query: str, + *, + continuation_token_limit: Optional[int] = None, + enable_scan_in_query: Optional[bool] = None, + feed_range: Optional[Dict[str, Any]] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + max_item_count: Optional[int] = None, + parameters: Optional[List[Dict[str, object]]] = None, + populate_index_metrics: Optional[bool] = None, + populate_query_metrics: Optional[bool] = None, + priority: Optional[Literal["High", "Low"]] = None, + response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, + session_token: Optional[str] = None, + throughput_bucket: Optional[int] = None, + **kwargs: Any + ): + """Return all results matching the given `query`. + + You can use any value for the container name in the FROM clause, but + often the container name is used. In the examples below, the container + name is "products," and is aliased as "p" for easier referencing in + the WHERE clause. + + :param str query: The Azure Cosmos DB SQL query to execute. :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + indexing was opted out on the requested paths. + :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations + in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + If all preferred locations were excluded, primary/hub location will be used. + This excluded_location will override existing excluded_locations in client level. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword parameters: Optional array of parameters to the query. + Each parameter is a dict() with 'name' and 'value' keys. + Ignored if no query is provided. + :paramtype parameters: [List[Dict[str, object]]] + :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used + existing indexes and how it could use potential new indexes. Please note that this options will incur + overhead, so it should be enabled only when debugging slow queries. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword response_hook: A callable invoked with the response metadata. + :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. :keyword int throughput_bucket: The desired throughput bucket for the client + :returns: An Iterable of items (dicts). + :rtype: ItemPaged[Dict[str, Any]] + + .. admonition:: Example: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items] + :end-before: [END query_items] + :language: python + :dedent: 0 + :caption: Get all products that have not been discontinued: + + .. literalinclude:: ../samples/examples.py + :start-after: [START query_items_param] + :end-before: [END query_items_param] + :language: python + :dedent: 0 + :caption: Parameterized query to get all products that have been discontinued: + """ + ... + + @distributed_trace + def query_items( + self, + *args: Any, + **kwargs: Any + ) -> AsyncItemPaged[Dict[str, Any]]: + """Return all results matching the given `query`. + + You can use any value for the container name in the FROM clause, but + often the container name is used. In the examples below, the container + name is "products," and is aliased as "p" for easier referencing in + the WHERE clause. + + :param Any args: args + :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query + response. Valid values are positive integers. + A value of 0 is the same as not passing a value (default no limit). + :keyword bool enable_cross_partition_query: Allows sending of more than one request to + execute the query in the Azure Cosmos DB service. + More than one request is necessary if the query is not scoped to single partition key value. + :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as + indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. - :returns: An AsyncItemPaged of items (dicts). - :rtype: AsyncItemPaged[Dict[str, Any]] + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. + :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. + :keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in + milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency, + responses are guaranteed to be no staler than this value. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword parameters: Optional array of parameters to the query. + Each parameter is a dict() with 'name' and 'value' keys. + Ignored if no query is provided. + :paramtype parameters: [List[Dict[str, object]]] + :keyword partition_key: partition key at which the query request is targeted. + :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] + :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used + existing indexes and how it could use potential new indexes. Please note that this options will incur + overhead, so it should be enabled only when debugging slow queries. + :keyword bool populate_query_metrics: Enable returning query metrics in response headers. + :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + request. Once the user has reached their provisioned throughput, low priority requests are throttled + before high priority requests start getting throttled. Feature must first be enabled at the account level. + :keyword str query: The Azure Cosmos DB SQL query to execute. + :keyword response_hook: A callable invoked with the response metadata. + :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] + :keyword str session_token: Token for use with Session consistency. + :keyword int throughput_bucket: The desired throughput bucket for the client + :returns: An Iterable of items (dicts). + :rtype: ItemPaged[Dict[str, Any]] .. admonition:: Example: - .. literalinclude:: ../samples/examples_async.py + .. literalinclude:: ../samples/examples.py :start-after: [START query_items] :end-before: [END query_items] :language: python :dedent: 0 :caption: Get all products that have not been discontinued: - :name: query_items - .. literalinclude:: ../samples/examples_async.py + .. literalinclude:: ../samples/examples.py :start-after: [START query_items_param] :end-before: [END query_items_param] :language: python :dedent: 0 :caption: Parameterized query to get all products that have been discontinued: - :name: query_items_param """ - if feed_range is not None: - kwargs['feed_range'] = feed_range - if initial_headers is not None: - kwargs['initial_headers'] = initial_headers - if priority is not None: - kwargs['priority'] = priority - if session_token is not None: - kwargs['session_token'] = session_token - if throughput_bucket is not None: - kwargs["throughput_bucket"] = throughput_bucket + original_positional_arg_names = ["query"] + utils.add_args_to_kwargs(original_positional_arg_names, args, kwargs) feed_options = _build_options(kwargs) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if populate_query_metrics is not None: - feed_options["populateQueryMetrics"] = populate_query_metrics - if populate_index_metrics is not None: - feed_options["populateIndexMetrics"] = populate_index_metrics - if enable_scan_in_query is not None: - feed_options["enableScanInQuery"] = enable_scan_in_query + + # Update 'feed_options' from 'kwargs' + if "max_item_count" in kwargs: + feed_options["maxItemCount"] = kwargs.pop("max_item_count") + if "populate_query_metrics" in kwargs: + feed_options["populateQueryMetrics"] = kwargs.pop("populate_query_metrics") + if "populate_index_metrics" in kwargs: + feed_options["populateIndexMetrics"] = kwargs.pop("populate_index_metrics") + if "enable_scan_in_query" in kwargs: + feed_options["enableScanInQuery"] = kwargs.pop("enable_scan_in_query") + if "max_integrated_cache_staleness_in_ms" in kwargs: + max_integrated_cache_staleness_in_ms = kwargs.pop("max_integrated_cache_staleness_in_ms") + validate_cache_staleness_value(max_integrated_cache_staleness_in_ms) + feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms + if "continuation_token_limit" in kwargs: + feed_options["responseContinuationTokenLimitInKb"] = kwargs.pop("continuation_token_limit") + feed_options["correlatedActivityId"] = GenerateGuidId() + + # Set query with 'query' and 'parameters' from kwargs + if "parameters" in kwargs: + query = {"query": kwargs.pop("query", None), "parameters": kwargs.pop("parameters", None)} + else: + query = kwargs.pop("query", None) + + # Set method to get/cache container properties kwargs["containerProperties"] = self._get_properties_with_options - if partition_key is not None: - feed_options["partitionKey"] = self._set_partition_key(partition_key) + + utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) + partition_key = None + if "feed_range" not in kwargs and "partition_key" in kwargs: + partition_key_value = kwargs.pop("partition_key") + feed_options["partitionKey"] = self._set_partition_key(partition_key_value) else: feed_options["enableCrossPartitionQuery"] = True - if max_integrated_cache_staleness_in_ms: - validate_cache_staleness_value(max_integrated_cache_staleness_in_ms) - feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms - correlated_activity_id = GenerateGuidId() - feed_options["correlatedActivityId"] = correlated_activity_id - if continuation_token_limit is not None: - feed_options["responseContinuationTokenLimitInKb"] = continuation_token_limit + + # Set 'response_hook' + response_hook = kwargs.pop("response_hook", None) if response_hook and hasattr(response_hook, "clear"): response_hook.clear() - if self.container_link in self.__get_client_container_caches(): - feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] items = self.client_connection.QueryItems( database_or_container_link=self.container_link, - query=query if parameters is None else {"query": query, "parameters": parameters}, + query=query, options=feed_options, partition_key=partition_key, response_hook=response_hook, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 4670e1edb06d..4769ddac22a3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -56,6 +56,7 @@ from .. import documents from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable from .._change_feed.change_feed_state import ChangeFeedState +from .._change_feed.feed_range_internal import FeedRangeInternalEpk from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants @@ -74,7 +75,8 @@ PartitionKey, PartitionKeyKind, _return_undefined_or_empty_partition_key, - NonePartitionKeyValue, _Empty + NonePartitionKeyValue, _Empty, + build_partition_key_from_properties, ) from ._auth_policy_async import AsyncCosmosBearerTokenCredentialPolicy from .._cosmos_http_logging_policy import CosmosHttpLoggingPolicy @@ -2162,7 +2164,7 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Ca def ReadContainers( self, database_link: str, - options: Optional[Mapping[str, Any]] = None, + options: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all collections in a database. @@ -2898,10 +2900,10 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: return [] initial_headers = self.default_headers.copy() - cont_prop_func = kwargs.pop("containerProperties", None) - cont_prop = None - if cont_prop_func: - cont_prop = await cont_prop_func(options) # get properties with feed options + container_property_func = kwargs.pop("containerProperties", None) + container_property = None + if container_property_func: + container_property = await container_property_func(options) # get properties with feed options # Copy to make sure that default_headers won't be changed. if query is None: @@ -2951,24 +2953,21 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: request_params = _request_object.RequestObject(typ, documents._OperationType.SqlQuery, req_headers) request_params.set_excluded_location_from_options(options) - # check if query has prefix partition key - partition_key_value = options.get("partitionKey", None) - is_prefix_partition_query = False - partition_key_obj = None - if cont_prop and partition_key_value is not None: - partition_key_definition = cont_prop["partitionKey"] - partition_key_obj = PartitionKey(path=partition_key_definition["paths"], - kind=partition_key_definition["kind"], - version=partition_key_definition["version"]) - is_prefix_partition_query = partition_key_obj._is_prefix_partition_key(partition_key_value) - - if is_prefix_partition_query and partition_key_obj: - # here get the overlapping ranges - req_headers.pop(http_constants.HttpHeaders.PartitionKey, None) - feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key( - partition_key_value) # cspell:disable-line - over_lapping_ranges = await self._routing_map_provider.get_overlapping_ranges(id_, - [feed_range_epk], + # Check if the over lapping ranges can be populated + feed_range_epk = None + if "feed_range" in kwargs: + feed_range = kwargs.pop("feed_range") + feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range() + elif options.get("partitionKey") is not None and container_property is not None: + # check if query has prefix partition key + partition_key_value = options["partitionKey"] + partition_key_obj = build_partition_key_from_properties(container_property) + if partition_key_obj._is_prefix_partition_key(partition_key_value): + req_headers.pop(http_constants.HttpHeaders.PartitionKey, None) + feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key(partition_key_value) + + if feed_range_epk is not None: + over_lapping_ranges = await self._routing_map_provider.get_overlapping_ranges(id_, [feed_range_epk], options) results: Dict[str, Any] = {} # For each over lapping range we will take a sub range of the feed range EPK that overlaps with the over diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 5a7ff1e190d3..c07560279092 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -49,6 +49,8 @@ from .partition_key import ( NonePartitionKeyValue, PartitionKey, + PartitionKeyType, + build_partition_key_from_properties, _Empty, _Undefined, _return_undefined_or_empty_partition_key @@ -61,23 +63,14 @@ # pylint: disable=missing-client-constructor-parameter-credential,missing-client-constructor-parameter-kwargs # pylint: disable=docstring-keyword-should-match-keyword-only -PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long - -def get_partition_key_from_properties(container_properties: Dict[str, Any]) -> PartitionKey: - partition_key_definition = container_properties["partitionKey"] - return PartitionKey( - path=partition_key_definition["paths"], - kind=partition_key_definition["kind"], - version=partition_key_definition["version"]) - def is_prefix_partition_key(container_properties: Dict[str, Any], partition_key: PartitionKeyType) -> bool: - partition_key_obj: PartitionKey = get_partition_key_from_properties(container_properties) + partition_key_obj: PartitionKey = build_partition_key_from_properties(container_properties) return partition_key_obj._is_prefix_partition_key(partition_key) def get_epk_range_for_partition_key( container_properties: Dict[str, Any], partition_key_value: PartitionKeyType) -> Range: - partition_key_obj: PartitionKey = get_partition_key_from_properties(container_properties) + partition_key_obj: PartitionKey = build_partition_key_from_properties(container_properties) return partition_key_obj._get_epk_range_for_partition_key(partition_key_value) class ContainerProxy: # pylint: disable=too-many-public-methods @@ -154,7 +147,7 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> def _set_partition_key( self, partition_key: PartitionKeyType - ) -> Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]: + ) -> PartitionKeyType: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) @@ -629,6 +622,7 @@ def query_items( name is "products," and is aliased as "p" for easier referencing in the WHERE clause. + :param str query: The Azure Cosmos DB SQL query to execute. :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). @@ -659,7 +653,6 @@ def query_items( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :keyword str query: The Azure Cosmos DB SQL query to execute. :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. @@ -713,6 +706,7 @@ def query_items( name is "products," and is aliased as "p" for easier referencing in the WHERE clause. + :param str query: The Azure Cosmos DB SQL query to execute. :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). @@ -742,7 +736,6 @@ def query_items( :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each request. Once the user has reached their provisioned throughput, low priority requests are throttled before high priority requests start getting throttled. Feature must first be enabled at the account level. - :keyword str query: The Azure Cosmos DB SQL query to execute. :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. @@ -781,6 +774,7 @@ def query_items( # pylint:disable=docstring-missing-param name is "products," and is aliased as "p" for easier referencing in the WHERE clause. + :param Any args: args :keyword int continuation_token_limit: The size limit in kb of the response continuation token in the query response. Valid values are positive integers. A value of 0 is the same as not passing a value (default no limit). @@ -841,6 +835,8 @@ def query_items( # pylint:disable=docstring-missing-param "max_item_count", "enable_scan_in_query", "populate_query_metrics"] utils.add_args_to_kwargs(original_positional_arg_names, args, kwargs) feed_options = build_options(kwargs) + + # Get container property and init client container caches container_properties = self._get_properties_with_options(feed_options) # Update 'feed_options' from 'kwargs' @@ -871,19 +867,19 @@ def query_items( # pylint:disable=docstring-missing-param # Set range filters for query. Options are either 'feed_range' or 'partition_key' utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) - partition_key = None - if "feed_range" in kwargs: - feed_options["feedRange"] = kwargs.pop("feed_range", None) - elif "partition_key" in kwargs: - partition_key = kwargs.pop("partition_key") - partition_key_value = self._set_partition_key(partition_key) - if is_prefix_partition_key(container_properties, partition_key): - kwargs["isPrefixPartitionQuery"] = True - kwargs["partitionKeyDefinition"] = container_properties["partitionKey"] - kwargs["partitionKeyDefinition"]["partition_key"] = partition_key_value + if "feed_range" not in kwargs and "partition_key" in kwargs: + partition_key_value = self._set_partition_key(kwargs.pop("partition_key")) + partition_key_obj = build_partition_key_from_properties(container_properties) + if partition_key_obj._is_prefix_partition_key(partition_key_value): + kwargs["prefix_partition_key_object"] = partition_key_obj + kwargs["prefix_partition_key_value"] = partition_key_value else: + # Add to feed_options, only when feed_range not given and partition_key was not prefixed partition_key feed_options["partitionKey"] = partition_key_value + # Set 'partition_key' for QueryItems method. This can be 'None' if feed range or prefix partition key was set + partition_key = feed_options.get("partitionKey") + # Set 'response_hook' response_hook = kwargs.pop("response_hook", None) if response_hook and hasattr(response_hook, "clear"): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index b17e9a89ef0a..a7b4bfb136b1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -23,7 +23,7 @@ from io import BytesIO import binascii import struct -from typing import IO, Sequence, Type, Union, overload, List, cast +from typing import Any, IO, Sequence, Type, Union, overload, List, cast, Dict from typing_extensions import Literal from ._cosmos_integers import _UInt32, _UInt64, _UInt128 @@ -96,6 +96,8 @@ class _Undefined: class _Infinity: """Represents infinity value for partitionKey.""" +PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], + Type[NonePartitionKeyValue], _Empty, _Undefined] class PartitionKey(dict): """Key used to partition a container into logical partitions. @@ -109,13 +111,13 @@ class PartitionKey(dict): """ @overload - def __init__(self, path: List[str], *, kind: Literal["MultiHash"] = PartitionKeyKind.MULTI_HASH, + def __init__(self, path: List[str], *, kind: Literal["MultiHash"] = "MultiHash", version: int = PartitionKeyVersion.V2 ) -> None: ... @overload - def __init__(self, path: str, *, kind: Literal["Hash"] = PartitionKeyKind.HASH, + def __init__(self, path: str, *, kind: Literal["Hash"] = "Hash", version:int = PartitionKeyVersion.V2 ) -> None: ... @@ -187,7 +189,7 @@ def _get_epk_range_for_prefix_partition_key( def _get_epk_range_for_partition_key( self, - pk_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long + pk_value: PartitionKeyType # pylint: disable=line-too-long ) -> _Range: if self._is_prefix_partition_key(pk_value): return self._get_epk_range_for_prefix_partition_key( @@ -348,7 +350,7 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( def _is_prefix_partition_key( self, - partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> bool: # pylint: disable=line-too-long + partition_key: PartitionKeyType) -> bool: # pylint: disable=line-too-long if self.kind != PartitionKeyKind.MULTI_HASH: return False ret = ((isinstance(partition_key, Sequence) and @@ -490,3 +492,10 @@ def _write_for_binary_encoding( elif isinstance(value, _Undefined): binary_writer.write(bytes([_PartitionKeyComponentType.Undefined])) + +def build_partition_key_from_properties(container_properties: Dict[str, Any]) -> PartitionKey: + partition_key_definition = container_properties["partitionKey"] + return PartitionKey( + path=partition_key_definition["paths"], + kind=partition_key_definition["kind"], + version=partition_key_definition["version"]) \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 52e3506b9b0b..9a2d0c074a7b 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -1,15 +1,12 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import os +import pytest +import test_config import unittest import uuid -import pytest - from azure.cosmos import CosmosClient -import test_config - from itertools import combinations from typing import List, Mapping, Set diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py index 52e3506b9b0b..bd37c7b765a6 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py @@ -1,15 +1,13 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import os -import unittest -import uuid - import pytest - -from azure.cosmos import CosmosClient +import pytest_asyncio import test_config +import unittest +import uuid +from azure.cosmos.aio import CosmosClient from itertools import combinations from typing import List, Mapping, Set @@ -21,15 +19,15 @@ MULTI_PARTITION_CONTAINER_ID = CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] PK_VALUES = ('pk1', 'pk2', 'pk3') -def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: +async def add_all_pk_values_to_set_async(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: if len(items) == 0: return pk_values = [item['pk'] for item in items] pk_value_set.update(pk_values) -@pytest.fixture(scope="class", autouse=True) -def setup_and_teardown(): +@pytest_asyncio.fixture(scope="class", autouse=True) +async def setup_and_teardown_async(): print("Setup: This runs before any tests") document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] @@ -37,38 +35,40 @@ def setup_and_teardown(): for container_id in TEST_CONTAINERS_IDS: container = database.get_container_client(container_id) for document_definition in document_definitions: - container.upsert_item(body=document_definition) + await container.upsert_item(body=document_definition) yield # Code to run after tests print("Teardown: This runs after all tests") -def get_container(container_id: str): +async def get_container(container_id: str): client = CosmosClient(HOST, KEY) db = client.get_database_client(DATABASE_ID) return db.get_container_client(container_id) -@pytest.mark.cosmosQuery -class TestQueryFeedRange(): +@pytest.mark.cosmosMultiRegion +@pytest.mark.asyncio +@pytest.mark.usefixtures("setup_and_teardown_async") +class TestQueryFeedRangeAsync: @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_all_partitions(self, container_id): - container = get_container(container_id) + async def test_query_with_feed_range_for_all_partitions(self, container_id): + container = await get_container(container_id) query = 'SELECT * from c' expected_pk_values = set(PK_VALUES) actual_pk_values = set() - iter_feed_ranges = list(container.read_feed_ranges()) - for feed_range in iter_feed_ranges: - items = list(container.query_items( - query=query, - enable_cross_partition_query=True, - feed_range=feed_range - )) - add_all_pk_values_to_set(items, actual_pk_values) + async for feed_range in container.read_feed_ranges(): + items = [item async for item in + (container.query_items( + query=query, + feed_range=feed_range + ) + )] + await add_all_pk_values_to_set_async(items, actual_pk_values) assert expected_pk_values == actual_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_single_partition_key(self, container_id): - container = get_container(container_id) + async def test_query_with_feed_range_for_single_partition_key(self, container_id): + container = await get_container(container_id) query = 'SELECT * from c' for pk_value in PK_VALUES: @@ -76,17 +76,18 @@ def test_query_with_feed_range_for_single_partition_key(self, container_id): actual_pk_values = set() feed_range = test_config.create_feed_range_between_pk_values(pk_value, pk_value) - items = list(container.query_items( - query=query, - enable_cross_partition_query=True, - feed_range=feed_range - )) - add_all_pk_values_to_set(items, actual_pk_values) + items = [item async for item in + (container.query_items( + query=query, + feed_range=feed_range + ) + )] + await add_all_pk_values_to_set_async(items, actual_pk_values) assert expected_pk_values == actual_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_multiple_partition_key(self, container_id): - container = get_container(container_id) + async def test_query_with_feed_range_for_multiple_partition_key(self, container_id): + container = await get_container(container_id) query = 'SELECT * from c' for pk_value1, pk_value2 in combinations(PK_VALUES, 2): @@ -94,17 +95,18 @@ def test_query_with_feed_range_for_multiple_partition_key(self, container_id): actual_pk_values = set() feed_range = test_config.create_feed_range_between_pk_values(pk_value1, pk_value2) - items = list(container.query_items( - query=query, - enable_cross_partition_query=True, - feed_range=feed_range - )) - add_all_pk_values_to_set(items, actual_pk_values) + items = [item async for item in + (container.query_items( + query=query, + feed_range=feed_range + ) + )] + await add_all_pk_values_to_set_async(items, actual_pk_values) assert expected_pk_values.issubset(actual_pk_values) @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_a_full_range(self, container_id): - container = get_container(container_id) + async def test_query_with_feed_range_for_a_full_range(self, container_id): + container = await get_container(container_id) query = 'SELECT * from c' expected_pk_values = set(PK_VALUES) @@ -116,12 +118,13 @@ def test_query_with_feed_range_for_a_full_range(self, container_id): is_max_inclusive=False, ) feed_range = test_config.create_feed_range_in_dict(new_range) - items = list(container.query_items( - query=query, - enable_cross_partition_query=True, - feed_range=feed_range - )) - add_all_pk_values_to_set(items, actual_pk_values) + items = [item async for item in + (container.query_items( + query=query, + feed_range=feed_range + ) + )] + await add_all_pk_values_to_set_async(items, actual_pk_values) assert expected_pk_values.issubset(actual_pk_values) if __name__ == "__main__": From 9dadc7df5938c94c55385e68fb65aed52dbdd7d8 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 23 Jun 2025 10:59:17 -0700 Subject: [PATCH 05/20] Added samples for query_items with feed_range --- sdk/cosmos/azure-cosmos/samples/examples.py | 14 ++++++++++++++ sdk/cosmos/azure-cosmos/samples/examples_async.py | 12 ++++++++++++ 2 files changed, 26 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/samples/examples.py b/sdk/cosmos/azure-cosmos/samples/examples.py index 6f51ae61bba8..b0d1597a69a9 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples.py +++ b/sdk/cosmos/azure-cosmos/samples/examples.py @@ -137,6 +137,20 @@ container.delete_item(queried_item, partition_key="Widget") # [END delete_items] +# Query items with feed range is also supported. This example +# gets all items within the feed range. +# [START query_items_feed_range] +import json + +for feed_range in container.read_feed_ranges(): + for queried_item in container.query_items( + query='SELECT * FROM c', + enable_cross_partition_query=True, + feed_range=feed_range, + ): + print(json.dumps(queried_item, indent=True)) +# [END query_items_param] + # Retrieve the properties of a database # [START get_database_properties] properties = database.read() diff --git a/sdk/cosmos/azure-cosmos/samples/examples_async.py b/sdk/cosmos/azure-cosmos/samples/examples_async.py index d67230df7238..70010c171a76 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples_async.py +++ b/sdk/cosmos/azure-cosmos/samples/examples_async.py @@ -262,6 +262,18 @@ async def examples_async(): await container.delete_item(item_dict, partition_key=["GA", "Atlanta", 30363]) # [END delete_items] + # Query items with feed range is also supported. This example + # gets all items within the feed range. + # [START query_items_feed_range] + import json + + async for feed_range in container.read_feed_ranges(): + async for queried_item in container.query_items( + query='SELECT * from c', + feed_range=feed_range): + print(json.dumps(queried_item, indent=True)) + # [END query_items_param] + # Get the feed ranges list from container. # [START read_feed_ranges] feed_ranges = [feed_range async for feed_range in container.read_feed_ranges()] From 9958b70489719dd9014f7cca553c125dce17e3ec Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 23 Jun 2025 15:33:00 -0700 Subject: [PATCH 06/20] Fix pylint error --- .../azure-cosmos/azure/cosmos/_cosmos_client_connection.py | 2 +- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 971eb60de2c3..bfcd235adce3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -26,7 +26,7 @@ import os import urllib.parse import uuid -from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast, Type +from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast from typing_extensions import TypedDict from urllib3.util.retry import Retry diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 5f8658cecd93..71b69eafa413 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -23,7 +23,7 @@ """ import warnings from datetime import datetime -from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload, Iterable, Callable +from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, cast, overload, Iterable, Callable from typing_extensions import Literal from azure.core import MatchConditions From c64f832586c6f49318e2a0c306ac7323683d523d Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 23 Jun 2025 15:34:39 -0700 Subject: [PATCH 07/20] Updated CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 1949488630a3..2328fb11e6e3 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.13.0b3 (Unreleased) #### Features Added +* Added feed range support in `query_items`. See [PR 41722](https://github.com/Azure/azure-sdk-for-python/pull/41722). #### Breaking Changes From de1446f7949b80f696c1df2e45fba78639bde9bb Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 23 Jun 2025 16:29:18 -0700 Subject: [PATCH 08/20] Fix test error --- .../tests/test_changefeed_partition_key_variation_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_changefeed_partition_key_variation_async.py b/sdk/cosmos/azure-cosmos/tests/test_changefeed_partition_key_variation_async.py index d68a7a27048b..5e3b2bc728f8 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_changefeed_partition_key_variation_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_changefeed_partition_key_variation_async.py @@ -262,7 +262,7 @@ async def _get_properties_override(): for item in items: try: - epk_range = container.get_epk_range_for_partition_key(container_properties, item["pk"]) + epk_range = container._get_epk_range_for_partition_key(container_properties, item["pk"]) assert epk_range is not None, f"EPK range should not be None for partition key {item['pk']}." except Exception as e: assert False, f"Failed to get EPK range for partition key {item['pk']}: {str(e)}" From fed65bf0cb606fe9d4184f8e2ac0a47e4e42f36a Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Tue, 24 Jun 2025 13:20:21 -0700 Subject: [PATCH 09/20] Fix test error --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | 4 ++-- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index f6d98c3c98a3..5fad012c269f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -920,8 +920,8 @@ def query_items_change_feed( # pylint: disable=unused-argument cast(PartitionKeyType, partition_key_value)) change_feed_state_context["partitionKeyFeedRange"] = self._get_epk_range_for_partition_key( partition_key_value, feed_options) - if "feedRange" in kwargs: - change_feed_state_context["feedRange"] = feed_options.pop('feedRange') + if "feed_range" in kwargs: + change_feed_state_context["feedRange"] = kwargs.pop('feed_range') if "continuation" in feed_options: change_feed_state_context["continuation"] = feed_options.pop("continuation") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 71b69eafa413..3e2b439d4899 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -566,7 +566,6 @@ def query_items_change_feed( change_feed_state_context["startTime"] = kwargs.pop("start_time") container_properties = self._get_properties_with_options(feed_options) - # TODO: validate partition_key and feed_range exclusive check here to avoid any extra API calls if "partition_key" in kwargs: partition_key = kwargs.pop("partition_key") change_feed_state_context["partitionKey"] = self._set_partition_key(cast(PartitionKeyType, partition_key)) From 643b16f39c6ff6c9e35031fe8e037044a3a9706e Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Tue, 24 Jun 2025 14:08:39 -0700 Subject: [PATCH 10/20] Changed to run feed_range async tests on emulator only --- sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py index bd37c7b765a6..8cf477cee319 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py @@ -45,7 +45,7 @@ async def get_container(container_id: str): db = client.get_database_client(DATABASE_ID) return db.get_container_client(container_id) -@pytest.mark.cosmosMultiRegion +@pytest.mark.cosmosQuery @pytest.mark.asyncio @pytest.mark.usefixtures("setup_and_teardown_async") class TestQueryFeedRangeAsync: From 5940578f1f7891b24be92d3140d1951979d8bbdc Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Tue, 24 Jun 2025 17:13:03 -0700 Subject: [PATCH 11/20] Fix tests --- sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 9a2d0c074a7b..8e45dd3d9d82 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -22,7 +22,7 @@ def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[s if len(items) == 0: return - pk_values = [item['pk'] for item in items] + pk_values = [item['pk'] for item in items if 'pk' in item] pk_value_set.update(pk_values) @pytest.fixture(scope="class", autouse=True) From a06526b641e46c0cb19d66881f038401561ddfc4 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Wed, 25 Jun 2025 11:17:56 -0700 Subject: [PATCH 12/20] Fix tests --- .../tests/test_query_feed_range.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 8e45dd3d9d82..434945cd18bc 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -8,15 +8,18 @@ from azure.cosmos import CosmosClient from itertools import combinations +from azure.cosmos.partition_key import PartitionKey from typing import List, Mapping, Set CONFIG = test_config.TestConfig() HOST = CONFIG.host KEY = CONFIG.credential DATABASE_ID = CONFIG.TEST_DATABASE_ID -SINGLE_PARTITION_CONTAINER_ID = CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID -MULTI_PARTITION_CONTAINER_ID = CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID +TEST_NAME = "Query FeedRange " +SINGLE_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID +MULTI_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] +TEST_OFFER_THROUGHPUTS = [CONFIG.THROUGHPUT_FOR_1_PARTITION, CONFIG.THROUGHPUT_FOR_5_PARTITIONS] PK_VALUES = ('pk1', 'pk2', 'pk3') def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: if len(items) == 0: @@ -29,10 +32,13 @@ def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[s def setup_and_teardown(): print("Setup: This runs before any tests") document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] - database = CosmosClient(HOST, KEY).get_database_client(DATABASE_ID) - for container_id in TEST_CONTAINERS_IDS: - container = database.get_container_client(container_id) + + for container_id, offer_throughput in zip(TEST_CONTAINERS_IDS, TEST_OFFER_THROUGHPUTS): + container = database.create_container_if_not_exists( + id=container_id, + partition_key=PartitionKey(path='/' + CONFIG.TEST_CONTAINER_PARTITION_KEY, kind='Hash'), + offer_throughput=offer_throughput) for document_definition in document_definitions: container.upsert_item(body=document_definition) yield @@ -61,7 +67,7 @@ def test_query_with_feed_range_for_all_partitions(self, container_id): feed_range=feed_range )) add_all_pk_values_to_set(items, actual_pk_values) - assert expected_pk_values == actual_pk_values + assert actual_pk_values == expected_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) def test_query_with_feed_range_for_single_partition_key(self, container_id): @@ -79,7 +85,7 @@ def test_query_with_feed_range_for_single_partition_key(self, container_id): feed_range=feed_range )) add_all_pk_values_to_set(items, actual_pk_values) - assert expected_pk_values == actual_pk_values + assert actual_pk_values == expected_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) def test_query_with_feed_range_for_multiple_partition_key(self, container_id): From 973c4a8298ed2d250d5865bed1904b9f334b3d02 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Wed, 25 Jun 2025 13:08:02 -0700 Subject: [PATCH 13/20] Fix tests --- .../tests/test_query_feed_range.py | 7 +++--- .../tests/test_query_feed_range_async.py | 22 +++++++++++++------ 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 434945cd18bc..419da0e13881 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -20,24 +20,25 @@ MULTI_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] TEST_OFFER_THROUGHPUTS = [CONFIG.THROUGHPUT_FOR_1_PARTITION, CONFIG.THROUGHPUT_FOR_5_PARTITIONS] +PARTITION_KEY = CONFIG.TEST_CONTAINER_PARTITION_KEY PK_VALUES = ('pk1', 'pk2', 'pk3') def add_all_pk_values_to_set(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: if len(items) == 0: return - pk_values = [item['pk'] for item in items if 'pk' in item] + pk_values = [item[PARTITION_KEY] for item in items if PARTITION_KEY in item] pk_value_set.update(pk_values) @pytest.fixture(scope="class", autouse=True) def setup_and_teardown(): print("Setup: This runs before any tests") - document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] + document_definitions = [{PARTITION_KEY: pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] database = CosmosClient(HOST, KEY).get_database_client(DATABASE_ID) for container_id, offer_throughput in zip(TEST_CONTAINERS_IDS, TEST_OFFER_THROUGHPUTS): container = database.create_container_if_not_exists( id=container_id, - partition_key=PartitionKey(path='/' + CONFIG.TEST_CONTAINER_PARTITION_KEY, kind='Hash'), + partition_key=PartitionKey(path='/' + PARTITION_KEY, kind='Hash'), offer_throughput=offer_throughput) for document_definition in document_definitions: container.upsert_item(body=document_definition) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py index 8cf477cee319..57bb719a40a5 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py @@ -9,33 +9,41 @@ from azure.cosmos.aio import CosmosClient from itertools import combinations +from azure.cosmos.partition_key import PartitionKey from typing import List, Mapping, Set CONFIG = test_config.TestConfig() HOST = CONFIG.host KEY = CONFIG.credential DATABASE_ID = CONFIG.TEST_DATABASE_ID -SINGLE_PARTITION_CONTAINER_ID = CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID -MULTI_PARTITION_CONTAINER_ID = CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID +TEST_NAME = "Query FeedRange " +SINGLE_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID +MULTI_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_MULTI_PARTITION_CONTAINER_ID TEST_CONTAINERS_IDS = [SINGLE_PARTITION_CONTAINER_ID, MULTI_PARTITION_CONTAINER_ID] +TEST_OFFER_THROUGHPUTS = [CONFIG.THROUGHPUT_FOR_1_PARTITION, CONFIG.THROUGHPUT_FOR_5_PARTITIONS] +PARTITION_KEY = CONFIG.TEST_CONTAINER_PARTITION_KEY PK_VALUES = ('pk1', 'pk2', 'pk3') async def add_all_pk_values_to_set_async(items: List[Mapping[str, str]], pk_value_set: Set[str]) -> None: if len(items) == 0: return - pk_values = [item['pk'] for item in items] + pk_values = [item[PARTITION_KEY] for item in items if PARTITION_KEY in item] pk_value_set.update(pk_values) @pytest_asyncio.fixture(scope="class", autouse=True) async def setup_and_teardown_async(): print("Setup: This runs before any tests") - document_definitions = [{'pk': pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] - + document_definitions = [{PARTITION_KEY: pk, 'id': str(uuid.uuid4())} for pk in PK_VALUES] database = CosmosClient(HOST, KEY).get_database_client(DATABASE_ID) - for container_id in TEST_CONTAINERS_IDS: - container = database.get_container_client(container_id) + + for container_id, offer_throughput in zip(TEST_CONTAINERS_IDS, TEST_OFFER_THROUGHPUTS): + container = await database.create_container_if_not_exists( + id=container_id, + partition_key=PartitionKey(path='/' + PARTITION_KEY, kind='Hash'), + offer_throughput=offer_throughput) for document_definition in document_definitions: await container.upsert_item(body=document_definition) + yield # Code to run after tests print("Teardown: This runs after all tests") From 2a71b63a24f4650d2fbb51cb97f18f522f8fd507 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Thu, 26 Jun 2025 11:25:55 -0700 Subject: [PATCH 14/20] Fix tests with positional_args --- sdk/cosmos/azure-cosmos/tests/test_query.py | 32 ++++++++++--------- .../tests/test_query_feed_range.py | 2 +- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py index c2889ad03dad..1dd5cdf76d38 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query.py @@ -569,37 +569,39 @@ def test_query_request_params_none_retry_policy(self): self.created_db.delete_container(created_collection.id) def test_query_positional_args(self): - container_id = "Multi Partition Test Container Query Positional Args " + str(uuid.uuid4()) - partition_key = "pk" + container = self.created_db.get_container_client(self.config.TEST_MULTI_PARTITION_CONTAINER_ID) partition_key_value1 = "pk1" partition_key_value2 = "pk2" - container = self.created_db.create_container_if_not_exists( - id=container_id, - partition_key=PartitionKey(path='/' + partition_key, kind='Hash'), - offer_throughput=self.config.THROUGHPUT_FOR_5_PARTITIONS, - ) num_items = 10 new_items = [] for pk_value in [partition_key_value1, partition_key_value2]: for i in range(num_items): - new_items.append({'pk': pk_value, 'id': f"{pk_value}_{i}", 'name': 'sample name'}) + item = { + self.config.TEST_CONTAINER_PARTITION_KEY: pk_value, + 'id': f"{pk_value}_{i}", + 'name': 'sample name' + } + new_items.append(item) for item in new_items: container.upsert_item(body=item) query = "SELECT * FROM root r WHERE r.name=@name" parameters = [{'name': '@name', 'value': 'sample name'}] - partition_key = partition_key_value2 + partition_key_value = partition_key_value2 + enable_cross_partition_query = True max_item_count = 3 + enable_scan_in_query = True + populate_query_metrics = True pager = container.query_items( query, parameters, - partition_key, - True, + partition_key_value, + enable_cross_partition_query, max_item_count, - True, - True, + enable_scan_in_query, + populate_query_metrics, ).by_page() ids = [] @@ -607,10 +609,10 @@ def test_query_positional_args(self): items = list(page) num_items = len(items) for item in items: - assert item['pk'] == partition_key + assert item['pk'] == partition_key_value ids.append(item['id']) assert num_items <= max_item_count - assert ids == [item['id'] for item in new_items if item['pk'] == partition_key] + assert ids == [item['id'] for item in new_items if item['pk'] == partition_key_value] def _MockExecuteFunctionSessionRetry(self, function, *args, **kwargs): if args: diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 419da0e13881..5d18d0330ad9 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -13,7 +13,7 @@ CONFIG = test_config.TestConfig() HOST = CONFIG.host -KEY = CONFIG.credential +KEY = CONFIG.masterKey DATABASE_ID = CONFIG.TEST_DATABASE_ID TEST_NAME = "Query FeedRange " SINGLE_PARTITION_CONTAINER_ID = TEST_NAME + CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID From 1b02d46a46ca323bff1e4d64411a0ee9e165d795 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Thu, 3 Jul 2025 11:17:04 -0700 Subject: [PATCH 15/20] Addressing comments - Doc string updates - Tests were updated to use existing helper methods - Added tests with query_items with feed_range and partition_key --- .../azure-cosmos/azure/cosmos/_utils.py | 20 ++++++++++- .../azure/cosmos/aio/_container.py | 34 +++++++++---------- .../azure-cosmos/azure/cosmos/container.py | 22 ++++++------ .../azure-cosmos/samples/examples_async.py | 7 +--- sdk/cosmos/azure-cosmos/tests/test_config.py | 9 ----- .../tests/test_query_feed_range.py | 25 +++++++------- .../tests/test_query_feed_range_async.py | 29 ++++++++-------- 7 files changed, 74 insertions(+), 72 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 3a9aa86c9534..e5c9bd4aea14 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -100,6 +100,24 @@ def add_args_to_kwargs( raise ValueError(f"{name} cannot be used as positional and keyword argument at the same time.") kwargs[name] = arg + +def format_list_with_and(items: List[str]) -> str: + """Format a list of items into a string with commas and 'and' for the last item. + + :param List[str] items: The list of items to format. + :return: A formatted string with items separated by commas and 'and' before the last item. + """ + quoted = [f"'{item}'" for item in items] + if len(quoted) > 2: + return ", ".join(quoted[:-1]) + ", and " + quoted[-1] + elif len(quoted) == 2: + return " and ".join(quoted) + elif quoted: + return quoted[0] + else: + return "" + + def verify_exclusive_arguments( exclusive_keys: List[str], **kwargs: Dict[str, Any]) -> None: @@ -112,4 +130,4 @@ def verify_exclusive_arguments( keys_in_kwargs = [key for key in exclusive_keys if key in kwargs and kwargs[key] is not None] if len(keys_in_kwargs) > 1: - raise ValueError(f"{', '.join(keys_in_kwargs)} are exclusive parameters, please only set one of them") + raise ValueError(f"{format_list_with_and(keys_in_kwargs)} are exclusive parameters, please only set one of them.") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 5fad012c269f..9acbaddea20d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -458,7 +458,7 @@ def query_items( :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. @@ -470,10 +470,10 @@ def query_items( Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] - :keyword partition_key: partition key at which the query request is targeted. + :keyword partition_key: Partition key at which the query request is targeted. :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -482,20 +482,20 @@ def query_items( :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] .. admonition:: Example: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items] :end-before: [END query_items] :language: python :dedent: 0 :caption: Get all products that have not been discontinued: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items_param] :end-before: [END query_items_param] :language: python @@ -538,7 +538,7 @@ def query_items( :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. @@ -552,7 +552,7 @@ def query_items( Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -561,20 +561,20 @@ def query_items( :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] .. admonition:: Example: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items] :end-before: [END query_items] :language: python :dedent: 0 :caption: Get all products that have not been discontinued: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items_param] :end-before: [END query_items_param] :language: python @@ -606,7 +606,7 @@ def query_items( :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. @@ -619,10 +619,10 @@ def query_items( Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] - :keyword partition_key: partition key at which the query request is targeted. + :keyword partition_key: Partition key at which the query request is targeted. :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -632,20 +632,20 @@ def query_items( :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] .. admonition:: Example: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items] :end-before: [END query_items] :language: python :dedent: 0 :caption: Get all products that have not been discontinued: - .. literalinclude:: ../samples/examples.py + .. literalinclude:: ../samples/examples_async.py :start-after: [START query_items_param] :end-before: [END query_items_param] :language: python diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 3e2b439d4899..d95f42003984 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -626,7 +626,7 @@ def query_items( :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, str] initial_headers: Initial headers to be sent as part of the request. @@ -638,10 +638,10 @@ def query_items( Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] - :keyword partition_key: partition key at which the query request is targeted. + :keyword partition_key: Partition key at which the query request is targeted. :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -650,7 +650,7 @@ def query_items( :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] @@ -710,7 +710,7 @@ def query_items( :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. @@ -724,7 +724,7 @@ def query_items( Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -733,7 +733,7 @@ def query_items( :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] @@ -778,7 +778,7 @@ def query_items( # pylint:disable=docstring-missing-param :keyword bool enable_scan_in_query: Allow scan on the queries which couldn't be served as indexing was opted out on the requested paths. :keyword list[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations - in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on. + in this list are specified as the names of the Azure Cosmos locations like, 'West US', 'East US' and so on. If all preferred locations were excluded, primary/hub location will be used. This excluded_location will override existing excluded_locations in client level. :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. @@ -791,10 +791,10 @@ def query_items( # pylint:disable=docstring-missing-param Each parameter is a dict() with 'name' and 'value' keys. Ignored if no query is provided. :paramtype parameters: [List[Dict[str, object]]] - :keyword partition_key: partition key at which the query request is targeted. + :keyword partition_key: Partition key at which the query request is targeted. :paramtype partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]]] :keyword bool populate_index_metrics: Used to obtain the index metrics to understand how the query engine used - existing indexes and how it could use potential new indexes. Please note that this options will incur + existing indexes and how it could use potential new indexes. Please note that this option will incur overhead, so it should be enabled only when debugging slow queries. :keyword bool populate_query_metrics: Enable returning query metrics in response headers. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each @@ -804,7 +804,7 @@ def query_items( # pylint:disable=docstring-missing-param :keyword response_hook: A callable invoked with the response metadata. :paramtype response_hook: Callable[[Mapping[str, str], Dict[str, Any]], None] :keyword str session_token: Token for use with Session consistency. - :keyword int throughput_bucket: The desired throughput bucket for the client + :keyword int throughput_bucket: The desired throughput bucket for the client. :returns: An Iterable of items (dicts). :rtype: ItemPaged[Dict[str, Any]] diff --git a/sdk/cosmos/azure-cosmos/samples/examples_async.py b/sdk/cosmos/azure-cosmos/samples/examples_async.py index 70010c171a76..c37a2809f35e 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples_async.py +++ b/sdk/cosmos/azure-cosmos/samples/examples_async.py @@ -7,6 +7,7 @@ from azure.cosmos import exceptions, PartitionKey from azure.cosmos.aio import CosmosClient +import json import os @@ -97,8 +98,6 @@ async def examples_async(): # The asynchronous client returns asynchronous iterators for its query methods; # as such, we iterate over it by using an async for loop # [START query_items] - import json - async for queried_item in container.query_items( query='SELECT * FROM products p WHERE p.productModel <> "DISCONTINUED"' ): @@ -247,8 +246,6 @@ async def examples_async(): # Query the items in a container using SQL-like syntax. This example # gets all items whose product model hasn't been discontinued. # [START query_items] - import json - async for queried_item in container.query_items( query='SELECT * FROM location l WHERE l.state = "WA"' ): @@ -265,8 +262,6 @@ async def examples_async(): # Query items with feed range is also supported. This example # gets all items within the feed range. # [START query_items_feed_range] - import json - async for feed_range in container.read_feed_ranges(): async for queried_item in container.query_items( query='SELECT * from c', diff --git a/sdk/cosmos/azure-cosmos/tests/test_config.py b/sdk/cosmos/azure-cosmos/tests/test_config.py index 93589a35e148..cc2e78b8e9d2 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_config.py +++ b/sdk/cosmos/azure-cosmos/tests/test_config.py @@ -553,12 +553,3 @@ def create_range(range_min: str, range_max: str, is_min_inclusive: bool = True, def create_feed_range_in_dict(feed_range): return FeedRangeInternalEpk(feed_range).to_dict() - -def create_feed_range_between_pk_values(pk1, pk2): - range_min = hash_partition_key_value([pk1]) - range_max = hash_partition_key_value([pk2]) - if range_min > range_max: - range_min, range_max = range_max, range_min - range_max += "FF" - new_range = create_range(range_min, range_max) - return FeedRangeInternalEpk(new_range).to_dict() \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py index 5d18d0330ad9..83f6bce0d603 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py @@ -71,7 +71,7 @@ def test_query_with_feed_range_for_all_partitions(self, container_id): assert actual_pk_values == expected_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_single_partition_key(self, container_id): + def test_query_with_feed_range_for_partition_key(self, container_id): container = get_container(container_id) query = 'SELECT * from c' @@ -79,7 +79,7 @@ def test_query_with_feed_range_for_single_partition_key(self, container_id): expected_pk_values = {pk_value} actual_pk_values = set() - feed_range = test_config.create_feed_range_between_pk_values(pk_value, pk_value) + feed_range = container.feed_range_from_partition_key(pk_value) items = list(container.query_items( query=query, enable_cross_partition_query=True, @@ -89,22 +89,21 @@ def test_query_with_feed_range_for_single_partition_key(self, container_id): assert actual_pk_values == expected_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - def test_query_with_feed_range_for_multiple_partition_key(self, container_id): + def test_query_with_both_feed_range_and_partition_key(self, container_id): container = get_container(container_id) - query = 'SELECT * from c' - - for pk_value1, pk_value2 in combinations(PK_VALUES, 2): - expected_pk_values = {pk_value1, pk_value2} - actual_pk_values = set() - feed_range = test_config.create_feed_range_between_pk_values(pk_value1, pk_value2) - items = list(container.query_items( + expected_error_message = "'feed_range' and 'partition_key' are exclusive parameters, please only set one of them." + query = 'SELECT * from c' + partition_key = PK_VALUES[0] + feed_range = container.feed_range_from_partition_key(partition_key) + with pytest.raises(ValueError) as e: + list(container.query_items( query=query, enable_cross_partition_query=True, - feed_range=feed_range + feed_range=feed_range, + partition_key=partition_key )) - add_all_pk_values_to_set(items, actual_pk_values) - assert expected_pk_values.issubset(actual_pk_values) + assert str(e.value) == expected_error_message @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) def test_query_with_feed_range_for_a_full_range(self, container_id): diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py index 57bb719a40a5..21388291523b 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py @@ -75,7 +75,7 @@ async def test_query_with_feed_range_for_all_partitions(self, container_id): assert expected_pk_values == actual_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - async def test_query_with_feed_range_for_single_partition_key(self, container_id): + async def test_query_with_feed_range_for_partition_key(self, container_id): container = await get_container(container_id) query = 'SELECT * from c' @@ -83,7 +83,7 @@ async def test_query_with_feed_range_for_single_partition_key(self, container_id expected_pk_values = {pk_value} actual_pk_values = set() - feed_range = test_config.create_feed_range_between_pk_values(pk_value, pk_value) + feed_range = await container.feed_range_from_partition_key(pk_value) items = [item async for item in (container.query_items( query=query, @@ -94,23 +94,22 @@ async def test_query_with_feed_range_for_single_partition_key(self, container_id assert expected_pk_values == actual_pk_values @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) - async def test_query_with_feed_range_for_multiple_partition_key(self, container_id): + async def test_query_with_both_feed_range_and_partition_key(self, container_id): container = await get_container(container_id) - query = 'SELECT * from c' - - for pk_value1, pk_value2 in combinations(PK_VALUES, 2): - expected_pk_values = {pk_value1, pk_value2} - actual_pk_values = set() - feed_range = test_config.create_feed_range_between_pk_values(pk_value1, pk_value2) + expected_error_message = "'feed_range' and 'partition_key' are exclusive parameters, please only set one of them." + query = 'SELECT * from c' + partition_key = PK_VALUES[0] + feed_range = await container.feed_range_from_partition_key(partition_key) + with pytest.raises(ValueError) as e: items = [item async for item in - (container.query_items( - query=query, - feed_range=feed_range - ) + (container.query_items( + query=query, + feed_range=feed_range, + partition_key=partition_key + ) )] - await add_all_pk_values_to_set_async(items, actual_pk_values) - assert expected_pk_values.issubset(actual_pk_values) + assert str(e.value) == expected_error_message @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS) async def test_query_with_feed_range_for_a_full_range(self, container_id): From da8ac148349eb933ffbfa7d92775b51aa8954ac3 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Mon, 7 Jul 2025 11:43:16 -0700 Subject: [PATCH 16/20] Fix pylint error --- sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index e5c9bd4aea14..9aa4209422cc 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -106,17 +106,17 @@ def format_list_with_and(items: List[str]) -> str: :param List[str] items: The list of items to format. :return: A formatted string with items separated by commas and 'and' before the last item. + :rtype: str """ + formatted_items = "" quoted = [f"'{item}'" for item in items] if len(quoted) > 2: - return ", ".join(quoted[:-1]) + ", and " + quoted[-1] + formatted_items = ", ".join(quoted[:-1]) + ", and " + quoted[-1] elif len(quoted) == 2: - return " and ".join(quoted) + formatted_items = " and ".join(quoted) elif quoted: - return quoted[0] - else: - return "" - + formatted_items = quoted[0] + return formatted_items def verify_exclusive_arguments( exclusive_keys: List[str], @@ -130,4 +130,5 @@ def verify_exclusive_arguments( keys_in_kwargs = [key for key in exclusive_keys if key in kwargs and kwargs[key] is not None] if len(keys_in_kwargs) > 1: - raise ValueError(f"{format_list_with_and(keys_in_kwargs)} are exclusive parameters, please only set one of them.") + raise ValueError(f"{format_list_with_and(keys_in_kwargs)} are exclusive parameters, " + f"please only set one of them.") From db5b0c5f8dd4d8755831fe7b85b50161ecd8d080 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Fri, 11 Jul 2025 13:04:53 -0700 Subject: [PATCH 17/20] Changed to non-public APIs for internal classes/methods --- .../azure/cosmos/_cosmos_client_connection.py | 2 +- ...bal_partition_endpoint_manager_circuit_breaker.py | 2 +- .../azure-cosmos/azure/cosmos/aio/_container.py | 2 +- .../cosmos/aio/_cosmos_client_connection_async.py | 6 +++--- ...rtition_endpoint_manager_circuit_breaker_async.py | 2 +- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 8 ++++---- .../azure-cosmos/azure/cosmos/partition_key.py | 12 ++++++------ 7 files changed, 17 insertions(+), 17 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index bfcd235adce3..07783f6a7bc2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3166,7 +3166,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: elif "prefix_partition_key_object" in kwargs and "prefix_partition_key_value" in kwargs: prefix_partition_key_obj = kwargs.pop("prefix_partition_key_object") prefix_partition_key_value: SequentialPartitionKeyType = kwargs.pop("prefix_partition_key_value") - feed_range_epk = prefix_partition_key_obj.get_epk_range_for_prefix_partition_key(prefix_partition_key_value) + feed_range_epk = prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value) # If feed_range_epk exist, query with the range if feed_range_epk is not None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker.py index f58119245639..87df417b4f93 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_global_partition_endpoint_manager_circuit_breaker.py @@ -67,7 +67,7 @@ def create_pk_range_wrapper(self, request: RequestObject) -> Optional[PartitionK 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 + 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)) partition_range = Range.PartitionKeyRangeToRange(partition_ranges[0]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 9acbaddea20d..762729b883e1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -157,7 +157,7 @@ async def _get_epk_range_for_partition_key( partition_key_definition = container_properties["partitionKey"] partition_key = get_partition_key_from_definition(partition_key_definition) - return partition_key.get_epk_range_for_partition_key(partition_key_value) + return partition_key._get_epk_range_for_partition_key(partition_key_value) @distributed_trace_async async def read( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 2f3462763a31..f3bc8b556842 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -76,7 +76,7 @@ SequentialPartitionKeyType, _return_undefined_or_empty_partition_key, NonePartitionKeyValue, _Empty, - build_partition_key_from_properties, + _build_partition_key_from_properties, ) from ._auth_policy_async import AsyncCosmosBearerTokenCredentialPolicy from .._cosmos_http_logging_policy import CosmosHttpLoggingPolicy @@ -2961,11 +2961,11 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: elif options.get("partitionKey") is not None and container_property is not None: # check if query has prefix partition key partition_key_value = options["partitionKey"] - partition_key_obj = build_partition_key_from_properties(container_property) + partition_key_obj = _build_partition_key_from_properties(container_property) if partition_key_obj.is_prefix_partition_key(partition_key_value): req_headers.pop(http_constants.HttpHeaders.PartitionKey, None) partition_key_value = cast(SequentialPartitionKeyType, partition_key_value) - feed_range_epk = partition_key_obj.get_epk_range_for_prefix_partition_key(partition_key_value) + feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key(partition_key_value) if feed_range_epk is not None: over_lapping_ranges = await self._routing_map_provider.get_overlapping_ranges(id_, [feed_range_epk], diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_partition_endpoint_manager_circuit_breaker_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_partition_endpoint_manager_circuit_breaker_async.py index b3364912e1ef..9737b34fdd06 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_partition_endpoint_manager_circuit_breaker_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_partition_endpoint_manager_circuit_breaker_async.py @@ -65,7 +65,7 @@ async def create_pk_range_wrapper(self, request: RequestObject) -> Optional[Part 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)] + epk_range = [partition_key._get_epk_range_for_partition_key(partition_key_value)] partition_ranges = await (self.client._routing_map_provider .get_overlapping_ranges(container_link, epk_range)) partition_range = Range.PartitionKeyRangeToRange(partition_ranges[0]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index d95f42003984..956c52d3a2b8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -51,7 +51,7 @@ PartitionKey, PartitionKeyType, SequentialPartitionKeyType, - build_partition_key_from_properties, + _build_partition_key_from_properties, _return_undefined_or_empty_partition_key, ) from .scripts import ScriptsProxy @@ -65,8 +65,8 @@ def get_epk_range_for_partition_key( container_properties: Dict[str, Any], partition_key_value: PartitionKeyType) -> Range: - partition_key_obj: PartitionKey = build_partition_key_from_properties(container_properties) - return partition_key_obj.get_epk_range_for_partition_key(partition_key_value) + partition_key_obj: PartitionKey = _build_partition_key_from_properties(container_properties) + return partition_key_obj._get_epk_range_for_partition_key(partition_key_value) class ContainerProxy: # pylint: disable=too-many-public-methods """An interface to interact with a specific DB Container. @@ -863,7 +863,7 @@ def query_items( # pylint:disable=docstring-missing-param utils.verify_exclusive_arguments(["feed_range", "partition_key"], **kwargs) if "feed_range" not in kwargs and "partition_key" in kwargs: partition_key_value = self._set_partition_key(kwargs.pop("partition_key")) - partition_key_obj = build_partition_key_from_properties(container_properties) + partition_key_obj = _build_partition_key_from_properties(container_properties) if partition_key_obj.is_prefix_partition_key(partition_key_value): kwargs["prefix_partition_key_object"] = partition_key_obj kwargs["prefix_partition_key_value"] = cast(SequentialPartitionKeyType, partition_key_value) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index a6a5590a66a2..d2e49b6fd607 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -184,7 +184,7 @@ def version(self) -> int: def version(self, value: int) -> None: self["version"] = value - def get_epk_range_for_prefix_partition_key( + def _get_epk_range_for_prefix_partition_key( self, pk_value: SequentialPartitionKeyType ) -> _Range: @@ -210,12 +210,12 @@ def get_epk_range_for_prefix_partition_key( max_epk = str(min_epk) + "FF" return _Range(min_epk, max_epk, True, False) - def get_epk_range_for_partition_key( + def _get_epk_range_for_partition_key( self, pk_value: PartitionKeyType ) -> _Range: if self.is_prefix_partition_key(pk_value): - return self.get_epk_range_for_prefix_partition_key( + return self._get_epk_range_for_prefix_partition_key( cast(SequentialPartitionKeyType, pk_value)) # else return point range @@ -260,7 +260,7 @@ def _get_effective_partition_key_for_hash_partitioning( return _to_hex_encoded_binary_string_v1(partition_key_components) @staticmethod - def get_hashed_partition_key_string( + def _get_hashed_partition_key_string( pk_value: SequentialPartitionKeyType, kind: str, version: int = PartitionKeyVersion.V2, @@ -284,7 +284,7 @@ def _get_effective_partition_key_string( if isinstance(self, _Infinity): return _MaximumExclusiveEffectivePartitionKey - return PartitionKey.get_hashed_partition_key_string(pk_value=pk_value, kind=self.kind, version=self.version) + return PartitionKey._get_hashed_partition_key_string(pk_value=pk_value, kind=self.kind, version=self.version) @staticmethod def _write_for_hashing( @@ -531,6 +531,6 @@ def get_partition_key_from_definition( version: int = partition_key_definition.get("version", 1) # Default to version 1 if not provided return PartitionKey(path=path, kind=kind, version=version) -def build_partition_key_from_properties(container_properties: Dict[str, Any]) -> PartitionKey: +def _build_partition_key_from_properties(container_properties: Dict[str, Any]) -> PartitionKey: partition_key_definition = container_properties["partitionKey"] return get_partition_key_from_definition(partition_key_definition) From ba078705eb1268a5663adaddc61691fe689253fa Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Fri, 11 Jul 2025 13:07:45 -0700 Subject: [PATCH 18/20] Changed to non-public APIs for internal classes/methods --- .../azure/cosmos/_cosmos_client_connection.py | 8 ++-- .../aio/_cosmos_client_connection_async.py | 4 +- .../azure-cosmos/azure/cosmos/container.py | 36 ++++++++-------- .../azure/cosmos/partition_key.py | 42 +++++++++---------- 4 files changed, 45 insertions(+), 45 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 07783f6a7bc2..dc9500d27fcc 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -73,8 +73,8 @@ _Undefined, _Empty, PartitionKeyKind, - PartitionKeyType, - SequentialPartitionKeyType, + _PartitionKeyType, + _SequentialPartitionKeyType, _return_undefined_or_empty_partition_key, ) @@ -1060,7 +1060,7 @@ def QueryItems( database_or_container_link: str, query: Optional[Union[str, Dict[str, Any]]], options: Optional[Mapping[str, Any]] = None, - partition_key: Optional[PartitionKeyType] = None, + partition_key: Optional[_PartitionKeyType] = None, response_hook: Optional[Callable[[Mapping[str, Any], Dict[str, Any]], None]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: @@ -3165,7 +3165,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range() elif "prefix_partition_key_object" in kwargs and "prefix_partition_key_value" in kwargs: prefix_partition_key_obj = kwargs.pop("prefix_partition_key_object") - prefix_partition_key_value: SequentialPartitionKeyType = kwargs.pop("prefix_partition_key_value") + prefix_partition_key_value: _SequentialPartitionKeyType = kwargs.pop("prefix_partition_key_value") feed_range_epk = prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value) # If feed_range_epk exist, query with the range diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index f3bc8b556842..cd5caf4da16c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -73,7 +73,7 @@ from ..partition_key import ( _Undefined, PartitionKeyKind, - SequentialPartitionKeyType, + _SequentialPartitionKeyType, _return_undefined_or_empty_partition_key, NonePartitionKeyValue, _Empty, _build_partition_key_from_properties, @@ -2964,7 +2964,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: partition_key_obj = _build_partition_key_from_properties(container_property) if partition_key_obj.is_prefix_partition_key(partition_key_value): req_headers.pop(http_constants.HttpHeaders.PartitionKey, None) - partition_key_value = cast(SequentialPartitionKeyType, partition_key_value) + partition_key_value = cast(_SequentialPartitionKeyType, partition_key_value) feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key(partition_key_value) if feed_range_epk is not None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 956c52d3a2b8..bc16316086b8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -49,8 +49,8 @@ from .partition_key import ( NonePartitionKeyValue, PartitionKey, - PartitionKeyType, - SequentialPartitionKeyType, + _PartitionKeyType, + _SequentialPartitionKeyType, _build_partition_key_from_properties, _return_undefined_or_empty_partition_key, ) @@ -64,7 +64,7 @@ def get_epk_range_for_partition_key( container_properties: Dict[str, Any], - partition_key_value: PartitionKeyType) -> Range: + partition_key_value: _PartitionKeyType) -> Range: partition_key_obj: PartitionKey = _build_partition_key_from_properties(container_properties) return partition_key_obj._get_epk_range_for_partition_key(partition_key_value) @@ -141,8 +141,8 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> def _set_partition_key( self, - partition_key: PartitionKeyType - ) -> PartitionKeyType: + partition_key: _PartitionKeyType + ) -> _PartitionKeyType: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) @@ -211,7 +211,7 @@ def read( # pylint:disable=docstring-missing-param def read_item( # pylint:disable=docstring-missing-param self, item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, populate_query_metrics: Optional[bool] = None, post_trigger_include: Optional[str] = None, *, @@ -359,7 +359,7 @@ def query_items_change_feed( *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, priority: Optional[Literal["High", "Low"]] = None, mode: Optional[Literal["LatestVersion", "AllVersionsAndDeletes"]] = None, response_hook: Optional[Callable[[Mapping[str, str], Dict[str, Any]], None]] = None, @@ -568,7 +568,7 @@ def query_items_change_feed( container_properties = self._get_properties_with_options(feed_options) if "partition_key" in kwargs: partition_key = kwargs.pop("partition_key") - change_feed_state_context["partitionKey"] = self._set_partition_key(cast(PartitionKeyType, partition_key)) + change_feed_state_context["partitionKey"] = self._set_partition_key(cast(_PartitionKeyType, partition_key)) change_feed_state_context["partitionKeyFeedRange"] = \ get_epk_range_for_partition_key(container_properties, partition_key) if "feed_range" in kwargs: @@ -600,7 +600,7 @@ def query_items( max_integrated_cache_staleness_in_ms: Optional[int] = None, max_item_count: Optional[int] = None, parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, + partition_key: Optional[_PartitionKeyType] = None, populate_index_metrics: Optional[bool] = None, populate_query_metrics: Optional[bool] = None, priority: Optional[Literal["High", "Low"]] = None, @@ -866,7 +866,7 @@ def query_items( # pylint:disable=docstring-missing-param partition_key_obj = _build_partition_key_from_properties(container_properties) if partition_key_obj.is_prefix_partition_key(partition_key_value): kwargs["prefix_partition_key_object"] = partition_key_obj - kwargs["prefix_partition_key_value"] = cast(SequentialPartitionKeyType, partition_key_value) + kwargs["prefix_partition_key_value"] = cast(_SequentialPartitionKeyType, partition_key_value) else: # Add to feed_options, only when feed_range not given and partition_key was not prefixed partition_key feed_options["partitionKey"] = partition_key_value @@ -1168,7 +1168,7 @@ def create_item( # pylint:disable=docstring-missing-param def patch_item( self, item: Union[str, Dict[str, Any]], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, patch_operations: List[Dict[str, Any]], *, filter_predicate: Optional[str] = None, @@ -1257,7 +1257,7 @@ def patch_item( def execute_item_batch( self, batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, *, pre_trigger_include: Optional[str] = None, post_trigger_include: Optional[str] = None, @@ -1329,7 +1329,7 @@ def execute_item_batch( def delete_item( # pylint:disable=docstring-missing-param self, item: Union[Mapping[str, Any], str], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, populate_query_metrics: Optional[bool] = None, pre_trigger_include: Optional[str] = None, post_trigger_include: Optional[str] = None, @@ -1521,7 +1521,7 @@ def query_conflicts( query: str, parameters: Optional[List[Dict[str, object]]] = None, enable_cross_partition_query: Optional[bool] = None, - partition_key: Optional[PartitionKeyType] = None, + partition_key: Optional[_PartitionKeyType] = None, max_item_count: Optional[int] = None, *, response_hook: Optional[Callable[[Mapping[str, Any], ItemPaged[Dict[str, Any]]], None]] = None, @@ -1567,7 +1567,7 @@ def query_conflicts( def get_conflict( self, conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, **kwargs: Any ) -> Dict[str, Any]: """Get the conflict identified by `conflict`. @@ -1595,7 +1595,7 @@ def get_conflict( def delete_conflict( self, conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, **kwargs: Any ) -> None: """Delete a specified conflict from the container. @@ -1624,7 +1624,7 @@ def delete_conflict( @distributed_trace def delete_all_items_by_partition_key( self, - partition_key: PartitionKeyType, + partition_key: _PartitionKeyType, *, pre_trigger_include: Optional[str] = None, post_trigger_include: Optional[str] = None, @@ -1745,7 +1745,7 @@ def get_latest_session_token( """ return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) - def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict[str, Any]: + def feed_range_from_partition_key(self, partition_key: _PartitionKeyType) -> Dict[str, Any]: """ Gets the feed range for a given partition key. :param partition_key: partition key to get feed range. :type partition_key: PartitionKeyType diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index d2e49b6fd607..c1d46108c020 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -96,9 +96,9 @@ class _Undefined: class _Infinity: """Represents infinity value for partitionKey.""" -SingularPartitionKeyType = Union[None, bool, float, int, str, Type[NonePartitionKeyValue], _Empty, _Undefined] -SequentialPartitionKeyType = Sequence[SingularPartitionKeyType] -PartitionKeyType = Union[SingularPartitionKeyType, SequentialPartitionKeyType] +_SingularPartitionKeyType = Union[None, bool, float, int, str, Type[NonePartitionKeyValue], _Empty, _Undefined] +_SequentialPartitionKeyType = Sequence[_SingularPartitionKeyType] +_PartitionKeyType = Union[_SingularPartitionKeyType, _SequentialPartitionKeyType] class PartitionKey(dict): """Key used to partition a container into logical partitions. @@ -186,7 +186,7 @@ def version(self, value: int) -> None: def _get_epk_range_for_prefix_partition_key( self, - pk_value: SequentialPartitionKeyType + pk_value: _SequentialPartitionKeyType ) -> _Range: if self.kind != PartitionKeyKind.MULTI_HASH: raise ValueError( @@ -212,11 +212,11 @@ def _get_epk_range_for_prefix_partition_key( def _get_epk_range_for_partition_key( self, - pk_value: PartitionKeyType + pk_value: _PartitionKeyType ) -> _Range: if self.is_prefix_partition_key(pk_value): return self._get_epk_range_for_prefix_partition_key( - cast(SequentialPartitionKeyType, pk_value)) + cast(_SequentialPartitionKeyType, pk_value)) # else return point range if isinstance(pk_value, (list, tuple)) or (isinstance(pk_value, Sequence) and not isinstance(pk_value, str)): @@ -228,15 +228,15 @@ def _get_epk_range_for_partition_key( @staticmethod def _truncate_for_v1_hashing( - value: SingularPartitionKeyType - ) -> SingularPartitionKeyType: + value: _SingularPartitionKeyType + ) -> _SingularPartitionKeyType: if isinstance(value, str): return value[:100] return value @staticmethod def _get_effective_partition_key_for_hash_partitioning( - pk_value: Union[str, SequentialPartitionKeyType] + pk_value: Union[str, _SequentialPartitionKeyType] ) -> str: truncated_components = [] # In Python, Strings are sequences, so we make sure we instead hash the entire string instead of each character @@ -261,7 +261,7 @@ def _get_effective_partition_key_for_hash_partitioning( @staticmethod def _get_hashed_partition_key_string( - pk_value: SequentialPartitionKeyType, + pk_value: _SequentialPartitionKeyType, kind: str, version: int = PartitionKeyVersion.V2, ) -> Union[int, str]: @@ -279,7 +279,7 @@ def _get_hashed_partition_key_string( def _get_effective_partition_key_string( self, - pk_value: SequentialPartitionKeyType + pk_value: _SequentialPartitionKeyType ) -> Union[int, str]: if isinstance(self, _Infinity): return _MaximumExclusiveEffectivePartitionKey @@ -288,21 +288,21 @@ def _get_effective_partition_key_string( @staticmethod def _write_for_hashing( - value: SingularPartitionKeyType, + value: _SingularPartitionKeyType, writer: IO[bytes] ) -> None: PartitionKey._write_for_hashing_core(value, bytes([0]), writer) @staticmethod def _write_for_hashing_v2( - value: SingularPartitionKeyType, + value: _SingularPartitionKeyType, writer: IO[bytes] ) -> None: PartitionKey._write_for_hashing_core(value, bytes([0xFF]), writer) @staticmethod def _write_for_hashing_core( - value: SingularPartitionKeyType, + value: _SingularPartitionKeyType, string_suffix: bytes, writer: IO[bytes] ) -> None: @@ -328,7 +328,7 @@ def _write_for_hashing_core( @staticmethod def _get_effective_partition_key_for_hash_partitioning_v2( - pk_value: SequentialPartitionKeyType + pk_value: _SequentialPartitionKeyType ) -> str: with BytesIO() as ms: for component in pk_value: @@ -347,7 +347,7 @@ def _get_effective_partition_key_for_hash_partitioning_v2( @staticmethod def _get_effective_partition_key_for_multi_hash_partitioning_v2( - pk_value: SequentialPartitionKeyType + pk_value: _SequentialPartitionKeyType ) -> str: sb = [] for value in pk_value: @@ -371,7 +371,7 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( def is_prefix_partition_key( self, - partition_key: PartitionKeyType) -> bool: # pylint: disable=line-too-long + partition_key: _PartitionKeyType) -> bool: # pylint: disable=line-too-long if self.kind != PartitionKeyKind.MULTI_HASH: return False ret = ((isinstance(partition_key, Sequence) and @@ -395,7 +395,7 @@ def _to_hex_encoded_binary_string(components: Sequence[object]) -> str: for component in components: if isinstance(component, (bool, int, float, str, _Infinity, _Undefined)): - component = cast(SingularPartitionKeyType, component) + component = cast(_SingularPartitionKeyType, component) _write_for_binary_encoding(component, ms) else: raise TypeError(f"Unexpected type for PK component: {type(component)}") @@ -406,7 +406,7 @@ def _to_hex_encoded_binary_string_v1(components: Sequence[object]) -> str: ms = BytesIO() for component in components: if isinstance(component, (bool, int, float, str, _Infinity, _Undefined)): - component = cast(SingularPartitionKeyType, component) + component = cast(_SingularPartitionKeyType, component) _write_for_binary_encoding_v1(component, ms) else: raise TypeError(f"Unexpected type for PK component: {type(component)}") @@ -414,7 +414,7 @@ def _to_hex_encoded_binary_string_v1(components: Sequence[object]) -> str: return _to_hex(bytearray(ms.getvalue()), 0, ms.tell()) def _write_for_binary_encoding_v1( - value: SingularPartitionKeyType, + value: _SingularPartitionKeyType, binary_writer: IO[bytes] ) -> None: if isinstance(value, bool): @@ -467,7 +467,7 @@ def _write_for_binary_encoding_v1( binary_writer.write(bytes([_PartitionKeyComponentType.Undefined])) def _write_for_binary_encoding( - value: SingularPartitionKeyType, + value: _SingularPartitionKeyType, binary_writer: IO[bytes] ) -> None: if isinstance(value, bool): From 28f151815b86f90e5712f4b4fabe0a86ae324279 Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Fri, 11 Jul 2025 13:09:53 -0700 Subject: [PATCH 19/20] Changed to non-public APIs for internal classes/methods --- .../_container_recreate_retry_policy.py | 8 ++--- .../azure/cosmos/_cosmos_client_connection.py | 4 +-- .../aio/_cosmos_client_connection_async.py | 4 +-- .../azure/cosmos/partition_key.py | 30 +++++++++---------- 4 files changed, 23 insertions(+), 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py index d7574e9aa79e..e36adfa5f323 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_container_recreate_retry_policy.py @@ -28,7 +28,7 @@ from azure.core.pipeline.transport._base import HttpRequest from . import http_constants -from .partition_key import _Empty, _Undefined, PartitionKeyKind +from .partition_key import _Empty, _Undefined, _PartitionKeyKind # pylint: disable=protected-access @@ -88,7 +88,7 @@ def should_extract_partition_key(self, container_cache: Optional[Dict[str, Any]] if self._headers and http_constants.HttpHeaders.PartitionKey in self._headers: current_partition_key = self._headers[http_constants.HttpHeaders.PartitionKey] partition_key_definition = container_cache["partitionKey"] if container_cache else None - if partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: + if partition_key_definition and partition_key_definition["kind"] == _PartitionKeyKind.MULTI_HASH: # A null in the multihash partition key indicates a failure in extracting partition keys # from the document definition return 'null' in current_partition_key @@ -110,7 +110,7 @@ def _extract_partition_key(self, client: Optional[Any], container_cache: Optiona elif options and isinstance(options["partitionKey"], _Empty): new_partition_key = [] # else serialize using json dumps method which apart from regular values will serialize None into null - elif partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: + elif partition_key_definition and partition_key_definition["kind"] == _PartitionKeyKind.MULTI_HASH: new_partition_key = json.dumps(options["partitionKey"], separators=(',', ':')) else: new_partition_key = json.dumps([options["partitionKey"]]) @@ -131,7 +131,7 @@ async def _extract_partition_key_async(self, client: Optional[Any], elif isinstance(options["partitionKey"], _Empty): new_partition_key = [] # else serialize using json dumps method which apart from regular values will serialize None into null - elif partition_key_definition and partition_key_definition["kind"] == PartitionKeyKind.MULTI_HASH: + elif partition_key_definition and partition_key_definition["kind"] == _PartitionKeyKind.MULTI_HASH: new_partition_key = json.dumps(options["partitionKey"], separators=(',', ':')) else: new_partition_key = json.dumps([options["partitionKey"]]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index dc9500d27fcc..8d638b69b6be 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -72,7 +72,7 @@ from .partition_key import ( _Undefined, _Empty, - PartitionKeyKind, + _PartitionKeyKind, _PartitionKeyType, _SequentialPartitionKeyType, _return_undefined_or_empty_partition_key, @@ -3334,7 +3334,7 @@ def _ExtractPartitionKey( partitionKeyDefinition: Mapping[str, Any], document: Mapping[str, Any] ) -> Union[List[Optional[Union[str, float, bool]]], str, float, bool, _Empty, _Undefined]: - if partitionKeyDefinition["kind"] == PartitionKeyKind.MULTI_HASH: + if partitionKeyDefinition["kind"] == _PartitionKeyKind.MULTI_HASH: ret: List[Optional[Union[str, float, bool]]] = [] for partition_key_level in partitionKeyDefinition["paths"]: # Parses the paths into a list of token each representing a property diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index cd5caf4da16c..b43155853c07 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -72,7 +72,7 @@ from .. import _utils from ..partition_key import ( _Undefined, - PartitionKeyKind, + _PartitionKeyKind, _SequentialPartitionKeyType, _return_undefined_or_empty_partition_key, NonePartitionKeyValue, _Empty, @@ -3183,7 +3183,7 @@ async def _AddPartitionKey(self, collection_link, document, options): # Extracts the partition key from the document using the partitionKey definition def _ExtractPartitionKey(self, partitionKeyDefinition, document): - if partitionKeyDefinition["kind"] == PartitionKeyKind.MULTI_HASH: + if partitionKeyDefinition["kind"] == _PartitionKeyKind.MULTI_HASH: ret = [] for partition_key_level in partitionKeyDefinition.get("paths"): # Parses the paths into a list of token each representing a property diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index c1d46108c020..e2c287367b25 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -68,11 +68,11 @@ class _PartitionKeyComponentType: Float = 0x14 Infinity = 0xFF -class PartitionKeyKind: +class _PartitionKeyKind: HASH: str = "Hash" MULTI_HASH: str = "MultiHash" -class PartitionKeyVersion: +class _PartitionKeyVersion: V1: int = 1 V2: int = 2 @@ -135,21 +135,21 @@ class PartitionKey(dict): @overload def __init__(self, path: List[str], *, kind: Literal["MultiHash"] = "MultiHash", - version: int = PartitionKeyVersion.V2 + version: int = _PartitionKeyVersion.V2 ) -> None: ... @overload def __init__(self, path: str, *, kind: Literal["Hash"] = "Hash", - version:int = PartitionKeyVersion.V2 + version:int = _PartitionKeyVersion.V2 ) -> None: ... def __init__(self, *args, **kwargs): path = args[0] if args else kwargs['path'] - kind = args[1] if len(args) > 1 else kwargs.get('kind', PartitionKeyKind.HASH if isinstance(path, str) - else PartitionKeyKind.MULTI_HASH) - version = args[2] if len(args) > 2 else kwargs.get('version', PartitionKeyVersion.V2) + kind = args[1] if len(args) > 1 else kwargs.get('kind', _PartitionKeyKind.HASH if isinstance(path, str) + else _PartitionKeyKind.MULTI_HASH) + version = args[2] if len(args) > 2 else kwargs.get('version', _PartitionKeyVersion.V2) super().__init__(paths=[path] if isinstance(path, str) else path, kind=kind, version=version) def __repr__(self) -> str: @@ -165,7 +165,7 @@ def kind(self, value: Literal["MultiHash", "Hash"]) -> None: @property def path(self) -> str: - if self.kind == PartitionKeyKind.MULTI_HASH: + if self.kind == _PartitionKeyKind.MULTI_HASH: return ''.join(self["paths"]) return self["paths"][0] @@ -188,7 +188,7 @@ def _get_epk_range_for_prefix_partition_key( self, pk_value: _SequentialPartitionKeyType ) -> _Range: - if self.kind != PartitionKeyKind.MULTI_HASH: + if self.kind != _PartitionKeyKind.MULTI_HASH: raise ValueError( "Effective Partition Key Range for Prefix Partition Keys is only supported for Hierarchical Partition Keys.") # pylint: disable=line-too-long len_pk_value = len(pk_value) @@ -263,17 +263,17 @@ def _get_effective_partition_key_for_hash_partitioning( def _get_hashed_partition_key_string( pk_value: _SequentialPartitionKeyType, kind: str, - version: int = PartitionKeyVersion.V2, + version: int = _PartitionKeyVersion.V2, ) -> Union[int, str]: if not pk_value: return _MinimumInclusiveEffectivePartitionKey - if kind == PartitionKeyKind.HASH: - if version == PartitionKeyVersion.V1: + if kind == _PartitionKeyKind.HASH: + if version == _PartitionKeyVersion.V1: return PartitionKey._get_effective_partition_key_for_hash_partitioning(pk_value) - if version == PartitionKeyVersion.V2: + if version == _PartitionKeyVersion.V2: return PartitionKey._get_effective_partition_key_for_hash_partitioning_v2(pk_value) - elif kind == PartitionKeyKind.MULTI_HASH: + elif kind == _PartitionKeyKind.MULTI_HASH: return PartitionKey._get_effective_partition_key_for_multi_hash_partitioning_v2(pk_value) return _to_hex_encoded_binary_string(pk_value) @@ -372,7 +372,7 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( def is_prefix_partition_key( self, partition_key: _PartitionKeyType) -> bool: # pylint: disable=line-too-long - if self.kind != PartitionKeyKind.MULTI_HASH: + if self.kind != _PartitionKeyKind.MULTI_HASH: return False ret = ((isinstance(partition_key, Sequence) and not isinstance(partition_key, str)) and len(self['paths']) != len(partition_key)) From 4a4a7bcc7ca84876d79eea0026e34fe4ce909ccf Mon Sep 17 00:00:00 2001 From: Allen Kim Date: Fri, 11 Jul 2025 14:03:56 -0700 Subject: [PATCH 20/20] Pylint error --- .../azure-cosmos/azure/cosmos/_cosmos_client_connection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 8d638b69b6be..69ada1d563d6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3166,7 +3166,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: elif "prefix_partition_key_object" in kwargs and "prefix_partition_key_value" in kwargs: prefix_partition_key_obj = kwargs.pop("prefix_partition_key_object") prefix_partition_key_value: _SequentialPartitionKeyType = kwargs.pop("prefix_partition_key_value") - feed_range_epk = prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value) + feed_range_epk = ( + prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value)) # If feed_range_epk exist, query with the range if feed_range_epk is not None: