From 2950e205df044d9b907a813da946ff940c46e315 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 14 Aug 2024 15:09:54 -0700 Subject: [PATCH 01/59] merge from main and resolve conflicts --- sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 19 +- .../azure/cosmos/_change_feed/__init__.py | 20 ++ .../azure/cosmos/_change_feed/aio/__init__.py | 20 ++ .../_change_feed/aio/change_feed_fetcher.py | 182 ++++++++++ .../_change_feed/aio/change_feed_iterable.py | 118 +++++++ .../aio/change_feed_start_from.py | 189 ++++++++++ .../_change_feed/aio/change_feed_state.py | 279 +++++++++++++++ .../aio/composite_continuation_token.py | 70 ++++ ...feed_range_composite_continuation_token.py | 134 +++++++ .../_change_feed/change_feed_fetcher.py | 181 ++++++++++ .../_change_feed/change_feed_iterable.py | 118 +++++++ .../_change_feed/change_feed_start_from.py | 189 ++++++++++ .../cosmos/_change_feed/change_feed_state.py | 279 +++++++++++++++ .../composite_continuation_token.py | 72 ++++ ...feed_range_composite_continuation_token.py | 134 +++++++ .../azure/cosmos/_cosmos_client_connection.py | 46 +-- .../azure/cosmos/_routing/routing_range.py | 73 ++++ .../azure-cosmos/azure/cosmos/_utils.py | 14 + .../azure/cosmos/aio/_container.py | 265 +++++++++++--- .../aio/_cosmos_client_connection_async.py | 10 +- .../azure-cosmos/azure/cosmos/container.py | 243 ++++++++++--- .../azure-cosmos/azure/cosmos/exceptions.py | 19 +- .../azure/cosmos/partition_key.py | 14 + .../azure-cosmos/test/test_change_feed.py | 295 ++++++++++++++++ .../test/test_change_feed_async.py | 322 +++++++++++++++++ sdk/cosmos/azure-cosmos/test/test_query.py | 290 +--------------- .../azure-cosmos/test/test_query_async.py | 328 +----------------- 27 files changed, 3177 insertions(+), 746 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/__init__.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/__init__.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_change_feed.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_change_feed_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index bcf5d95e0ef0..ab305f03b020 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -284,23 +284,8 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches if options.get("disableRUPerMinuteUsage"): headers[http_constants.HttpHeaders.DisableRUPerMinuteUsage] = options["disableRUPerMinuteUsage"] - if options.get("changeFeed") is True: - # On REST level, change feed is using IfNoneMatch/ETag instead of continuation. - if_none_match_value = None - if options.get("continuation"): - if_none_match_value = options["continuation"] - elif options.get("isStartFromBeginning") and not options["isStartFromBeginning"]: - if_none_match_value = "*" - elif options.get("startTime"): - start_time = options.get("startTime") - headers[http_constants.HttpHeaders.IfModified_since] = start_time - if if_none_match_value: - headers[http_constants.HttpHeaders.IfNoneMatch] = if_none_match_value - - headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - else: - if options.get("continuation"): - headers[http_constants.HttpHeaders.Continuation] = options["continuation"] + if options.get("continuation"): + headers[http_constants.HttpHeaders.Continuation] = options["continuation"] if options.get("populatePartitionKeyRangeStatistics"): headers[http_constants.HttpHeaders.PopulatePartitionKeyRangeStatistics] = options[ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/__init__.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/__init__.py new file mode 100644 index 000000000000..f5373937e446 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/__init__.py @@ -0,0 +1,20 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/__init__.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/__init__.py new file mode 100644 index 000000000000..f5373937e446 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/__init__.py @@ -0,0 +1,20 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py new file mode 100644 index 000000000000..83ca3025ee07 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -0,0 +1,182 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for processing change feed implementation in the Azure Cosmos +database service. +""" +import base64 +import copy +import json +from abc import ABC, abstractmethod + +from azure.cosmos import http_constants, exceptions +from azure.cosmos._change_feed.aio.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 +from azure.cosmos.aio import _retry_utility_async +from azure.cosmos.exceptions import CosmosHttpResponseError + + +class ChangeFeedFetcher(ABC): + + @abstractmethod + async def fetch_next_block(self): + pass + +class ChangeFeedFetcherV1(ChangeFeedFetcher): + """Internal class for change feed fetch v1 implementation. + This is used when partition key range id is used or when the supplied continuation token is in just simple etag. + Please note v1 does not support split or merge. + + """ + def __init__( + self, + client, + resource_link: str, + feed_options: dict[str, any], + fetch_function): + + self._client = client + self._feed_options = feed_options + + self._change_feed_state = self._feed_options.pop("changeFeedState") + if not isinstance(self._change_feed_state, ChangeFeedStateV1): + raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version {type(self._change_feed_state)}") + self._change_feed_state.__class__ = ChangeFeedStateV1 + + self._resource_link = resource_link + self._fetch_function = fetch_function + + async def fetch_next_block(self): + """Returns a block of results. + + :return: List of results. + :rtype: list + """ + async def callback(): + return await self.fetch_change_feed_items(self._fetch_function) + + return await _retry_utility_async.ExecuteAsync(self._client, self._client._global_endpoint_manager, callback) + + async def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + new_options = copy.deepcopy(self._feed_options) + new_options["changeFeedState"] = self._change_feed_state + + self._change_feed_state.populate_feed_options(new_options) + is_s_time_first_fetch = True + while True: + (fetched_items, response_headers) = await fetch_function(new_options) + continuation_key = http_constants.HttpHeaders.ETag + # In change feed queries, the continuation token is always populated. The hasNext() test is whether + # there is any items in the response or not. + # For start time however we get no initial results, so we need to pass continuation token? Is this true? + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + + if fetched_items: + break + elif is_s_time_first_fetch: + is_s_time_first_fetch = False + else: + break + return fetched_items + + +class ChangeFeedFetcherV2(object): + """Internal class for change feed fetch v2 implementation. + """ + + def __init__( + self, + client, + resource_link: str, + feed_options: dict[str, any], + fetch_function): + + self._client = client + self._feed_options = feed_options + + self._change_feed_state = self._feed_options.pop("changeFeedState") + if not isinstance(self._change_feed_state, ChangeFeedStateV2): + raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") + self._change_feed_state.__class__ = ChangeFeedStateV2 + + self._resource_link = resource_link + self._fetch_function = fetch_function + + async def fetch_next_block(self): + """Returns a block of results. + + :return: List of results. + :rtype: list + """ + + async def callback(): + return await self.fetch_change_feed_items(self._fetch_function) + + try: + return await _retry_utility_async.ExecuteAsync(self._client, self._client._global_endpoint_manager, callback) + except CosmosHttpResponseError as e: + if exceptions._partition_range_is_gone(e) or exceptions._is_partition_split_or_merge(e): + # refresh change feed state + await self._change_feed_state.handle_feed_range_gone(self._client._routing_map_provider, self._resource_link) + else: + raise e + + return await self.fetch_next_block() + + async def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + new_options = copy.deepcopy(self._feed_options) + new_options["changeFeedState"] = self._change_feed_state + + self._change_feed_state.populate_feed_options(new_options) + + is_s_time_first_fetch = True + while True: + (fetched_items, response_headers) = await fetch_function(new_options) + + continuation_key = http_constants.HttpHeaders.ETag + # In change feed queries, the continuation token is always populated. The hasNext() test is whether + # there is any items in the response or not. + # For start time however we get no initial results, so we need to pass continuation token? Is this true? + if fetched_items: + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + response_headers[continuation_key] = self._get_base64_encoded_continuation() + break + else: + self._change_feed_state.apply_not_modified_response() + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + response_headers[continuation_key] = self._get_base64_encoded_continuation() + should_retry = self._change_feed_state.should_retry_on_not_modified_response() or is_s_time_first_fetch + is_s_time_first_fetch = False + if not should_retry: + break + + return fetched_items + + def _get_base64_encoded_continuation(self) -> str: + continuation_json = json.dumps(self._change_feed_state.to_dict()) + json_bytes = continuation_json.encode('utf-8') + # Encode the bytes to a Base64 string + base64_bytes = base64.b64encode(json_bytes) + # Convert the Base64 bytes to a string + return base64_bytes.decode('utf-8') + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py new file mode 100644 index 000000000000..501f3a7e4150 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -0,0 +1,118 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Iterable change feed results in the Azure Cosmos database service. +""" +from azure.core.async_paging import AsyncPageIterator + +from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 +from azure.cosmos._change_feed.aio.change_feed_state import ChangeFeedStateV1, ChangeFeedState +from azure.cosmos._utils import is_base64_encoded + + +class ChangeFeedIterable(AsyncPageIterator): + """Represents an iterable object of the change feed results. + + ChangeFeedIterable is a wrapper for change feed execution. + """ + + def __init__( + self, + client, + options, + fetch_function=None, + collection_link=None, + continuation_token=None, + ): + """Instantiates a ChangeFeedIterable for non-client side partitioning queries. + + ChangeFeedFetcher will be used as the internal query execution + context. + + :param CosmosClient client: Instance of document client. + :param dict options: The request options for the request. + :param method fetch_function: + + """ + self._client = client + self.retry_options = client.connection_policy.RetryOptions + self._options = options + self._fetch_function = fetch_function + self._collection_link = collection_link + + change_feed_state = self._options.get("changeFeedState") + if not change_feed_state: + raise ValueError("Missing changeFeedState in feed options") + + if isinstance(change_feed_state, ChangeFeedStateV1): + if continuation_token: + if is_base64_encoded(continuation_token): + raise ValueError("Incompatible continuation token") + else: + change_feed_state.apply_server_response_continuation(continuation_token) + + self._change_feed_fetcher = ChangeFeedFetcherV1( + self._client, + self._collection_link, + self._options, + fetch_function + ) + else: + if continuation_token: + if not is_base64_encoded(continuation_token): + raise ValueError("Incompatible continuation token") + + effective_change_feed_context = {"continuationFeedRange": continuation_token} + effective_change_feed_state = ChangeFeedState.from_json(change_feed_state.container_rid, effective_change_feed_context) + # replace with the effective change feed state + self._options["continuationFeedRange"] = effective_change_feed_state + + self._change_feed_fetcher = ChangeFeedFetcherV2( + self._client, + self._collection_link, + self._options, + fetch_function + ) + super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) + + async def _unpack(self, block): + continuation = None + if self._client.last_response_headers: + continuation = self._client.last_response_headers.get('etag') + + if block: + self._did_a_call_already = False + return continuation, block + + async def _fetch_next(self, *args): # pylint: disable=unused-argument + """Return a block of results with respecting retry policy. + + This method only exists for backward compatibility reasons. (Because + QueryIterable has exposed fetch_next_block api). + + :param Any args: + :return: List of results. + :rtype: list + """ + block = await self._change_feed_fetcher.fetch_next_block() + if not block: + raise StopAsyncIteration + return block diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py new file mode 100644 index 000000000000..99aeeb6eb914 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py @@ -0,0 +1,189 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed start from implementation in the Azure Cosmos database service. +""" + +from abc import ABC, abstractmethod +from datetime import datetime, timezone +from enum import Enum +from typing import Optional, Union, Literal, Any + +from azure.cosmos import http_constants +from azure.cosmos._routing.routing_range import Range + +class ChangeFeedStartFromType(Enum): + BEGINNING = "Beginning" + NOW = "Now" + LEASE = "Lease" + POINT_IN_TIME = "PointInTime" + +class ChangeFeedStartFromInternal(ABC): + """Abstract class for change feed start from implementation in the Azure Cosmos database service. + """ + + _type_property_name = "Type" + + @abstractmethod + def to_dict(self) -> dict[str, Any]: + pass + + @staticmethod + def from_start_time(start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]]) -> 'ChangeFeedStartFromInternal': + if start_time is None: + return ChangeFeedStartFromNow() + elif isinstance(start_time, datetime): + return ChangeFeedStartFromPointInTime(start_time) + elif start_time.lower() == ChangeFeedStartFromType.NOW.value.lower(): + return ChangeFeedStartFromNow() + elif start_time.lower() == ChangeFeedStartFromType.BEGINNING.value.lower(): + return ChangeFeedStartFromBeginning() + else: + raise ValueError(f"Invalid start_time '{start_time}'") + + @staticmethod + def from_json(data: dict[str, any]) -> 'ChangeFeedStartFromInternal': + change_feed_start_from_type = data.get(ChangeFeedStartFromInternal._type_property_name) + if change_feed_start_from_type is None: + raise ValueError(f"Invalid start from json [Missing {ChangeFeedStartFromInternal._type_property_name}]") + + if change_feed_start_from_type == ChangeFeedStartFromType.BEGINNING.value: + return ChangeFeedStartFromBeginning.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.LEASE.value: + return ChangeFeedStartFromETagAndFeedRange.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.NOW.value: + return ChangeFeedStartFromNow.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.POINT_IN_TIME.value: + return ChangeFeedStartFromPointInTime.from_json(data) + else: + raise ValueError(f"Can not process changeFeedStartFrom for type {change_feed_start_from_type}") + + @abstractmethod + def populate_request_headers(self, request_headers) -> None: + pass + + +class ChangeFeedStartFromBeginning(ChangeFeedStartFromInternal): + """Class for change feed start from beginning implementation in the Azure Cosmos database service. + """ + + def to_dict(self) -> dict[str, Any]: + return { + self._type_property_name: ChangeFeedStartFromType.BEGINNING.value + } + + def populate_request_headers(self, request_headers) -> None: + pass # there is no headers need to be set for start from beginning + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromBeginning': + return ChangeFeedStartFromBeginning() + + +class ChangeFeedStartFromETagAndFeedRange(ChangeFeedStartFromInternal): + """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. + """ + + _etag_property_name = "Etag" + _feed_range_property_name = "FeedRange" + + def __init__(self, etag, feed_range): + if feed_range is None: + raise ValueError("feed_range is missing") + + self._etag = etag + self._feed_range = feed_range + + def to_dict(self) -> dict[str, Any]: + return { + self._type_property_name: ChangeFeedStartFromType.LEASE.value, + self._etag_property_name: self._etag, + self._feed_range_property_name: self._feed_range.to_dict() + } + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromETagAndFeedRange': + etag = data.get(cls._etag_property_name) + if etag is None: + raise ValueError(f"Invalid change feed start from [Missing {cls._etag_property_name}]") + + feed_range_data = data.get(cls._feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid change feed start from [Missing {cls._feed_range_property_name}]") + feed_range = Range.ParseFromDict(feed_range_data) + return cls(etag, feed_range) + + def populate_request_headers(self, request_headers) -> None: + # change feed uses etag as the continuationToken + if self._etag: + request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._etag + + +class ChangeFeedStartFromNow(ChangeFeedStartFromInternal): + """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. + """ + + def to_dict(self) -> dict[str, Any]: + return { + self._type_property_name: ChangeFeedStartFromType.NOW.value + } + + def populate_request_headers(self, request_headers) -> None: + request_headers[http_constants.HttpHeaders.IfNoneMatch] = "*" + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromNow': + return ChangeFeedStartFromNow() + + +class ChangeFeedStartFromPointInTime(ChangeFeedStartFromInternal): + """Class for change feed start from point in time implementation in the Azure Cosmos database service. + """ + + _point_in_time_ms_property_name = "PointInTimeMs" + + def __init__(self, start_time: datetime): + if start_time is None: + raise ValueError("start_time is missing") + + self._start_time = start_time + + def to_dict(self) -> dict[str, Any]: + return { + self._type_property_name: ChangeFeedStartFromType.POINT_IN_TIME.value, + self._point_in_time_ms_property_name: + int(self._start_time.astimezone(timezone.utc).timestamp() * 1000) + } + + def populate_request_headers(self, request_headers) -> None: + request_headers[http_constants.HttpHeaders.IfModified_since] =\ + self._start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': + point_in_time_ms = data.get(cls._point_in_time_ms_property_name) + if point_in_time_ms is None: + raise ValueError(f"Invalid change feed start from {cls._point_in_time_ms_property_name} ") + + point_in_time = datetime.fromtimestamp(point_in_time_ms).astimezone(timezone.utc) + return ChangeFeedStartFromPointInTime(point_in_time) + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py new file mode 100644 index 000000000000..ae2e37568bd4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py @@ -0,0 +1,279 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed state implementation in the Azure Cosmos +database service. +""" + +import base64 +import collections +import json +from abc import ABC, abstractmethod +from typing import Optional, Union, List, Any + +from azure.cosmos import http_constants +from azure.cosmos._change_feed.aio.change_feed_start_from import ChangeFeedStartFromETagAndFeedRange, \ + ChangeFeedStartFromInternal +from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._change_feed.aio.feed_range_composite_continuation_token import FeedRangeCompositeContinuation +from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider +from azure.cosmos._routing.routing_range import Range +from azure.cosmos._utils import is_key_exists_and_not_none +from azure.cosmos.exceptions import CosmosFeedRangeGoneError +from azure.cosmos.partition_key import _Empty, _Undefined + + +class ChangeFeedState(ABC): + version_property_name = "v" + + @abstractmethod + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + pass + + @abstractmethod + async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, request_headers: dict[str, any]) -> None: + pass + + @abstractmethod + def apply_server_response_continuation(self, continuation: str) -> None: + pass + + @staticmethod + def from_json(container_link: str, container_rid: str, data: dict[str, Any]): + if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): + return ChangeFeedStateV1.from_json(container_link, container_rid, data) + else: + if is_key_exists_and_not_none(data, "continuationFeedRange"): + # get changeFeedState from continuation + continuation_json_str = base64.b64decode(data["continuationFeedRange"]).decode('utf-8') + continuation_json = json.loads(continuation_json_str) + version = continuation_json.get(ChangeFeedState.version_property_name) + if version is None: + raise ValueError("Invalid base64 encoded continuation string [Missing version]") + elif version == "V2": + return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) + else: + raise ValueError("Invalid base64 encoded continuation string [Invalid version]") + # when there is no continuation token, by default construct ChangeFeedStateV2 + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) + +class ChangeFeedStateV1(ChangeFeedState): + """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag + """ + + def __init__( + self, + container_link: str, + container_rid: str, + change_feed_start_from: ChangeFeedStartFromInternal, + partition_key_range_id: Optional[str] = None, + partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, + continuation: Optional[str] = None): + + self._container_link = container_link + self._container_rid = container_rid + self._change_feed_start_from = change_feed_start_from + self._partition_key_range_id = partition_key_range_id + self._partition_key = partition_key + self._continuation = continuation + + @property + def container_rid(self): + return self._container_rid + + @classmethod + def from_json(cls, container_link: str, container_rid: str, data: dict[str, Any]) -> 'ChangeFeedStateV1': + return cls( + container_link, + container_rid, + ChangeFeedStartFromInternal.from_start_time(data.get("startTime")), + data.get("partitionKeyRangeId"), + data.get("partitionKey"), + data.get("continuationPkRangeId") + ) + + async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, Any]) -> None: + headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(headers) + if self._continuation: + headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation + + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + if self._partition_key_range_id is not None: + feed_options["partitionKeyRangeId"] = self._partition_key_range_id + if self._partition_key is not None: + feed_options["partitionKey"] = self._partition_key + + def apply_server_response_continuation(self, continuation: str) -> None: + self._continuation = continuation + +class ChangeFeedStateV2(ChangeFeedState): + container_rid_property_name = "containerRid" + change_feed_mode_property_name = "mode" + change_feed_start_from_property_name = "startFrom" + continuation_property_name = "continuation" + + # TODO: adding change feed mode + def __init__( + self, + container_link: str, + container_rid: str, + feed_range: Range, + change_feed_start_from: ChangeFeedStartFromInternal, + continuation: Optional[FeedRangeCompositeContinuation] = None): + + self._container_link = container_link + self._container_rid = container_rid + self._feed_range = feed_range + self._change_feed_start_from = change_feed_start_from + self._continuation = continuation + if self._continuation is None: + composite_continuation_token_queue = collections.deque() + composite_continuation_token_queue.append(CompositeContinuationToken(self._feed_range, None)) + self._continuation =\ + FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) + + @property + def container_rid(self) -> str : + return self._container_rid + + def to_dict(self) -> dict[str, Any]: + return { + self.version_property_name: "V2", + self.container_rid_property_name: self._container_rid, + self.change_feed_mode_property_name: "Incremental", + self.change_feed_start_from_property_name: self._change_feed_start_from.to_dict(), + self.continuation_property_name: self._continuation.to_dict() + } + + async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, any]) -> None: + headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(headers) + + if self._continuation.current_token is not None and self._continuation.current_token.token is not None: + change_feed_start_from_feed_range_and_etag =\ + ChangeFeedStartFromETagAndFeedRange(self._continuation.current_token.token, self._continuation.current_token.feed_range) + change_feed_start_from_feed_range_and_etag.populate_request_headers(headers) + + # based on the feed range to find the overlapping partition key range id + over_lapping_ranges =\ + await routing_provider.get_overlapping_ranges( + self._container_link, + [self._continuation.current_token.feed_range]) + + if len(over_lapping_ranges) > 1: + raise CosmosFeedRangeGoneError(message= + f"Range {self._continuation.current_token.feed_range}" + f" spans {len(over_lapping_ranges)}" + f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + else: + overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) + if overlapping_feed_range == self._continuation.current_token.feed_range: + # exactly mapping to one physical partition, only need to set the partitionKeyRangeId + headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + else: + # the current token feed range spans less than single physical partition + # for this case, need to set both the partition key range id and epk filter headers + headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + headers[http_constants.HttpHeaders.StartEpkString] = self._continuation.current_token.feed_range.min + headers[http_constants.HttpHeaders.EndEpkString] = self._continuation.current_token.feed_range.max + + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + pass + + async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, resource_link: str) -> None: + await self._continuation.handle_feed_range_gone(routing_provider, resource_link) + + def apply_server_response_continuation(self, continuation: str) -> None: + self._continuation.apply_server_response_continuation(continuation) + + def should_retry_on_not_modified_response(self): + self._continuation.should_retry_on_not_modified_response() + + def apply_not_modified_response(self) -> None: + self._continuation.apply_not_modified_response() + + @classmethod + def from_continuation( + cls, + container_link: str, + container_rid: str, + continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': + + container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) + if container_rid_from_continuation is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.container_rid_property_name}]") + elif container_rid_from_continuation != container_rid: + raise ValueError("Invalid continuation: [Mismatch collection rid]") + + change_feed_start_from_data = continuation_json.get(ChangeFeedStateV2.change_feed_start_from_property_name) + if change_feed_start_from_data is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.change_feed_start_from_property_name}]") + change_feed_start_from = ChangeFeedStartFromInternal.from_json(change_feed_start_from_data) + + continuation_data = continuation_json.get(ChangeFeedStateV2.continuation_property_name) + if continuation_data is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.continuation_property_name}]") + continuation = FeedRangeCompositeContinuation.from_json(continuation_data) + return ChangeFeedStateV2( + container_link=container_link, + container_rid=container_rid, + feed_range=continuation.feed_range, + change_feed_start_from=change_feed_start_from, + continuation=continuation) + + @classmethod + def from_initial_state( + cls, + container_link: str, + collection_rid: str, + data: dict[str, Any]) -> 'ChangeFeedStateV2': + + if is_key_exists_and_not_none(data, "feedRange"): + feed_range_str = base64.b64decode(data["feedRange"]).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range = Range.ParseFromDict(feed_range_json) + elif is_key_exists_and_not_none(data, "partitionKeyFeedRange"): + feed_range = data["partitionKeyFeedRange"] + else: + # default to full range + feed_range = Range( + "", + "FF", + True, + False) + + change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(data.get("startTime")) + return cls( + container_link=container_link, + container_rid=collection_rid, + feed_range=feed_range, + change_feed_start_from=change_feed_start_from, + continuation=None) + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py new file mode 100644 index 000000000000..6d779fed1037 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py @@ -0,0 +1,70 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed composite continuation token in the Azure Cosmos +database service. +""" +from azure.cosmos._routing.routing_range import Range + + +class CompositeContinuationToken(object): + _token_property_name = "token" + _feed_range_property_name = "range" + + def __init__(self, feed_range: Range, token): + if range is None: + raise ValueError("range is missing") + + self._token = token + self._feed_range = feed_range + + def to_dict(self): + return { + self._token_property_name: self._token, + self._feed_range_property_name: self._feed_range.to_dict() + } + + @property + def feed_range(self): + return self._feed_range + + @property + def token(self): + return self._token + + def update_token(self, etag): + self._token = etag + + @classmethod + def from_json(cls, data): + token = data.get(cls._token_property_name) + if token is None: + raise ValueError(f"Invalid composite token [Missing {cls._token_property_name}]") + + feed_range_data = data.get(cls._feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid composite token [Missing {cls._feed_range_property_name}]") + + feed_range = Range.ParseFromDict(feed_range_data) + return cls(feed_range=feed_range, token=token) + + def __repr__(self): + return f"CompositeContinuationToken(token={self.token}, range={self._feed_range})" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py new file mode 100644 index 000000000000..6e1b8f974eea --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py @@ -0,0 +1,134 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed continuation token by feed range in the Azure Cosmos +database service. +""" +import collections +from collections import deque +from typing import Any + +from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider +from azure.cosmos._routing.routing_range import Range + + +class FeedRangeCompositeContinuation(object): + _version_property_name = "V" + _container_rid_property_name = "Rid" + _continuation_property_name = "Continuation" + _feed_range_property_name = "Range" + + def __init__( + self, + container_rid: str, + feed_range: Range, + continuation: collections.deque[CompositeContinuationToken]): + if container_rid is None: + raise ValueError("container_rid is missing") + + self._container_rid = container_rid + self._feed_range = feed_range + self._continuation = continuation + self._current_token = self._continuation[0] + self._initial_no_result_range = None + + @property + def current_token(self): + return self._current_token + + def to_dict(self) -> dict[str, Any]: + return { + self._version_property_name: "v1", #TODO: should this start from v2 + self._container_rid_property_name: self._container_rid, + self._continuation_property_name: [childToken.to_dict() for childToken in self._continuation], + self._feed_range_property_name: self._feed_range.to_dict() + } + + @classmethod + def from_json(cls, data) -> 'FeedRangeCompositeContinuation': + version = data.get(cls._version_property_name) + if version is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._version_property_name}]") + if version != "v1": + raise ValueError("Invalid feed range composite continuation token [Invalid version]") + + container_rid = data.get(cls._container_rid_property_name) + if container_rid is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") + + feed_range_data = data.get(cls._feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._feed_range_property_name}]") + feed_range = Range.ParseFromDict(feed_range_data) + + continuation_data = data.get(cls._continuation_property_name) + if continuation_data is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") + if not isinstance(continuation_data, list) or len(continuation_data) == 0: + raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") + continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] + + return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) + + async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: + overlapping_ranges = await routing_provider.get_overlapping_ranges(collection_link, self._current_token.feed_range) + + if len(overlapping_ranges) == 1: + # merge,reusing the existing the feedRange and continuationToken + pass + else: + # split, remove the parent range and then add new child ranges. + # For each new child range, using the continuation token from the parent + self._continuation.popleft() + for child_range in overlapping_ranges: + self._continuation.append(CompositeContinuationToken(Range.PartitionKeyRangeToRange(child_range), self._current_token.token)) + + self._current_token = self._continuation[0] + + def should_retry_on_not_modified_response(self) -> bool: + # when getting 304(Not Modified) response from one sub feed range, we will try to fetch for the next sub feed range + # we will repeat the above logic until we have looped through all sub feed ranges + + # TODO: validate the response headers, can we get the status code + if len(self._continuation) > 1: + return self._current_token.feed_range != self._initial_no_result_range + + else: + return False + + def _move_to_next_token(self) -> None: + first_composition_token = self._continuation.popleft() + # add the composition token to the end of the list + self._continuation.append(first_composition_token) + self._current_token = self._continuation[0] + + def apply_server_response_continuation(self, etag) -> None: + self._current_token.update_token(etag) + self._move_to_next_token() + + def apply_not_modified_response(self) -> None: + if self._initial_no_result_range is None: + self._initial_no_result_range = self._current_token.feed_range + + @property + def feed_range(self) -> Range: + return self._feed_range diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py new file mode 100644 index 000000000000..fd8ac2787a8f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -0,0 +1,181 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for processing change feed implementation in the Azure Cosmos +database service. +""" +import base64 +import copy +import json +from abc import ABC, abstractmethod + +from azure.cosmos import _retry_utility, http_constants, exceptions +from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 +from azure.cosmos.exceptions import CosmosHttpResponseError + + +class ChangeFeedFetcher(ABC): + + @abstractmethod + def fetch_next_block(self): + pass + +class ChangeFeedFetcherV1(ChangeFeedFetcher): + """Internal class for change feed fetch v1 implementation. + This is used when partition key range id is used or when the supplied continuation token is in just simple etag. + Please note v1 does not support split or merge. + + """ + def __init__( + self, + client, + resource_link: str, + feed_options: dict[str, any], + fetch_function): + + self._client = client + self._feed_options = feed_options + + self._change_feed_state = self._feed_options.pop("changeFeedState") + if not isinstance(self._change_feed_state, ChangeFeedStateV1): + raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version {type(self._change_feed_state)}") + self._change_feed_state.__class__ = ChangeFeedStateV1 + + self._resource_link = resource_link + self._fetch_function = fetch_function + + def fetch_next_block(self): + """Returns a block of results. + + :return: List of results. + :rtype: list + """ + def callback(): + return self.fetch_change_feed_items(self._fetch_function) + + return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback) + + def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + new_options = copy.deepcopy(self._feed_options) + new_options["changeFeedState"] = self._change_feed_state + + self._change_feed_state.populate_feed_options(new_options) + is_s_time_first_fetch = True + while True: + (fetched_items, response_headers) = fetch_function(new_options) + continuation_key = http_constants.HttpHeaders.ETag + # In change feed queries, the continuation token is always populated. The hasNext() test is whether + # there is any items in the response or not. + # For start time however we get no initial results, so we need to pass continuation token? Is this true? + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + + if fetched_items: + break + elif is_s_time_first_fetch: + is_s_time_first_fetch = False + else: + break + + return fetched_items + + +class ChangeFeedFetcherV2(object): + """Internal class for change feed fetch v2 implementation. + """ + + def __init__( + self, + client, + resource_link: str, + feed_options: dict[str, any], + fetch_function): + + self._client = client + self._feed_options = feed_options + + self._change_feed_state = self._feed_options.pop("changeFeedState") + if not isinstance(self._change_feed_state, ChangeFeedStateV2): + raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") + self._change_feed_state.__class__ = ChangeFeedStateV2 + + self._resource_link = resource_link + self._fetch_function = fetch_function + + def fetch_next_block(self): + """Returns a block of results. + + :return: List of results. + :rtype: list + """ + + def callback(): + return self.fetch_change_feed_items(self._fetch_function) + + try: + return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback) + except CosmosHttpResponseError as e: + if exceptions._partition_range_is_gone(e) or exceptions._is_partition_split_or_merge(e): + # refresh change feed state + self._change_feed_state.handle_feed_range_gone(self._client._routing_map_provider, self._resource_link) + else: + raise e + + return self.fetch_next_block() + + def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + new_options = copy.deepcopy(self._feed_options) + new_options["changeFeedState"] = self._change_feed_state + + self._change_feed_state.populate_feed_options(new_options) + + is_s_time_first_fetch = True + while True: + (fetched_items, response_headers) = fetch_function(new_options) + + continuation_key = http_constants.HttpHeaders.ETag + # In change feed queries, the continuation token is always populated. The hasNext() test is whether + # there is any items in the response or not. + # For start time however we get no initial results, so we need to pass continuation token? Is this true? + if fetched_items: + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + response_headers[continuation_key] = self._get_base64_encoded_continuation() + break + else: + self._change_feed_state.apply_not_modified_response() + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + response_headers[continuation_key] = self._get_base64_encoded_continuation() + should_retry = self._change_feed_state.should_retry_on_not_modified_response() or is_s_time_first_fetch + is_s_time_first_fetch = False + if not should_retry: + break + + return fetched_items + + def _get_base64_encoded_continuation(self) -> str: + continuation_json = json.dumps(self._change_feed_state.to_dict()) + json_bytes = continuation_json.encode('utf-8') + # Encode the bytes to a Base64 string + base64_bytes = base64.b64encode(json_bytes) + # Convert the Base64 bytes to a string + return base64_bytes.decode('utf-8') diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py new file mode 100644 index 000000000000..676036180d29 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -0,0 +1,118 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Iterable change feed results in the Azure Cosmos database service. +""" + +from azure.core.paging import PageIterator + +from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 +from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedState +from azure.cosmos._utils import is_base64_encoded + + +class ChangeFeedIterable(PageIterator): + """Represents an iterable object of the change feed results. + + ChangeFeedIterable is a wrapper for change feed execution. + """ + + def __init__( + self, + client, + options, + fetch_function=None, + collection_link=None, + continuation_token=None, + ): + """Instantiates a ChangeFeedIterable for non-client side partitioning queries. + + :param CosmosClient client: Instance of document client. + :param dict options: The request options for the request. + :param fetch_function: The fetch function. + :param collection_link: The collection resource link. + :param continuation_token: The continuation token passed in from by_page + """ + + self._client = client + self.retry_options = client.connection_policy.RetryOptions + self._options = options + self._fetch_function = fetch_function + self._collection_link = collection_link + + change_feed_state = self._options.get("changeFeedState") + if not change_feed_state: + raise ValueError("Missing changeFeedState in feed options") + + if isinstance(change_feed_state, ChangeFeedStateV1): + if continuation_token: + if is_base64_encoded(continuation_token): + raise ValueError("Incompatible continuation token") + else: + change_feed_state.apply_server_response_continuation(continuation_token) + + self._change_feed_fetcher = ChangeFeedFetcherV1( + self._client, + self._collection_link, + self._options, + fetch_function + ) + else: + if continuation_token: + if not is_base64_encoded(continuation_token): + raise ValueError("Incompatible continuation token") + + effective_change_feed_context = {"continuationFeedRange": continuation_token} + effective_change_feed_state = ChangeFeedState.from_json(change_feed_state.container_rid, effective_change_feed_context) + # replace with the effective change feed state + self._options["continuationFeedRange"] = effective_change_feed_state + + self._change_feed_fetcher = ChangeFeedFetcherV2( + self._client, + self._collection_link, + self._options, + fetch_function + ) + super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) + + def _unpack(self, block): + continuation = None + if self._client.last_response_headers: + continuation = self._client.last_response_headers.get('etag') + + if block: + self._did_a_call_already = False + return continuation, block + + def _fetch_next(self, *args): # pylint: disable=unused-argument + """Return a block of results with respecting retry policy. + + This method only exists for backward compatibility reasons. (Because + QueryIterable has exposed fetch_next_block api). + + :param Any args: + :return: List of results. + :rtype: list + """ + block = self._change_feed_fetcher.fetch_next_block() + if not block: + raise StopIteration + return block diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py new file mode 100644 index 000000000000..76a4d6b56803 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py @@ -0,0 +1,189 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed start from implementation in the Azure Cosmos database service. +""" + +from abc import ABC, abstractmethod +from datetime import datetime, timezone +from enum import Enum +from typing import Optional, Union, Literal, Any + +from azure.cosmos import http_constants +from azure.cosmos._routing.routing_range import Range + +class ChangeFeedStartFromType(Enum): + BEGINNING = "Beginning" + NOW = "Now" + LEASE = "Lease" + POINT_IN_TIME = "PointInTime" + +class ChangeFeedStartFromInternal(ABC): + """Abstract class for change feed start from implementation in the Azure Cosmos database service. + """ + + type_property_name = "Type" + + @abstractmethod + def to_dict(self) -> dict[str, Any]: + pass + + @staticmethod + def from_start_time(start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]]) -> 'ChangeFeedStartFromInternal': + if start_time is None: + return ChangeFeedStartFromNow() + elif isinstance(start_time, datetime): + return ChangeFeedStartFromPointInTime(start_time) + elif start_time.lower() == ChangeFeedStartFromType.NOW.value.lower(): + return ChangeFeedStartFromNow() + elif start_time.lower() == ChangeFeedStartFromType.BEGINNING.value.lower(): + return ChangeFeedStartFromBeginning() + else: + raise ValueError(f"Invalid start_time '{start_time}'") + + @staticmethod + def from_json(data: dict[str, any]) -> 'ChangeFeedStartFromInternal': + change_feed_start_from_type = data.get(ChangeFeedStartFromInternal.type_property_name) + if change_feed_start_from_type is None: + raise ValueError(f"Invalid start from json [Missing {ChangeFeedStartFromInternal.type_property_name}]") + + if change_feed_start_from_type == ChangeFeedStartFromType.BEGINNING.value: + return ChangeFeedStartFromBeginning.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.LEASE.value: + return ChangeFeedStartFromETagAndFeedRange.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.NOW.value: + return ChangeFeedStartFromNow.from_json(data) + elif change_feed_start_from_type == ChangeFeedStartFromType.POINT_IN_TIME.value: + return ChangeFeedStartFromPointInTime.from_json(data) + else: + raise ValueError(f"Can not process changeFeedStartFrom for type {change_feed_start_from_type}") + + @abstractmethod + def populate_request_headers(self, request_headers) -> None: + pass + + +class ChangeFeedStartFromBeginning(ChangeFeedStartFromInternal): + """Class for change feed start from beginning implementation in the Azure Cosmos database service. + """ + + def to_dict(self) -> dict[str, Any]: + return { + self.type_property_name: ChangeFeedStartFromType.BEGINNING.value + } + + def populate_request_headers(self, request_headers) -> None: + pass # there is no headers need to be set for start from beginning + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromBeginning': + return ChangeFeedStartFromBeginning() + + +class ChangeFeedStartFromETagAndFeedRange(ChangeFeedStartFromInternal): + """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. + """ + + _etag_property_name = "Etag" + _feed_range_property_name = "FeedRange" + + def __init__(self, etag, feed_range): + if feed_range is None: + raise ValueError("feed_range is missing") + + self._etag = etag + self._feed_range = feed_range + + def to_dict(self) -> dict[str, Any]: + return { + self.type_property_name: ChangeFeedStartFromType.LEASE.value, + self._etag_property_name: self._etag, + self._feed_range_property_name: self._feed_range.to_dict() + } + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromETagAndFeedRange': + etag = data.get(cls._etag_property_name) + if etag is None: + raise ValueError(f"Invalid change feed start from [Missing {cls._etag_property_name}]") + + feed_range_data = data.get(cls._feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid change feed start from [Missing {cls._feed_range_property_name}]") + feed_range = Range.ParseFromDict(feed_range_data) + return cls(etag, feed_range) + + def populate_request_headers(self, request_headers) -> None: + # change feed uses etag as the continuationToken + if self._etag: + request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._etag + + +class ChangeFeedStartFromNow(ChangeFeedStartFromInternal): + """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. + """ + + def to_dict(self) -> dict[str, Any]: + return { + self.type_property_name: ChangeFeedStartFromType.NOW.value + } + + def populate_request_headers(self, request_headers) -> None: + request_headers[http_constants.HttpHeaders.IfNoneMatch] = "*" + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromNow': + return ChangeFeedStartFromNow() + + +class ChangeFeedStartFromPointInTime(ChangeFeedStartFromInternal): + """Class for change feed start from point in time implementation in the Azure Cosmos database service. + """ + + _point_in_time_ms_property_name = "PointInTimeMs" + + def __init__(self, start_time: datetime): + if start_time is None: + raise ValueError("start_time is missing") + + self._start_time = start_time + + def to_dict(self) -> dict[str, Any]: + return { + self.type_property_name: ChangeFeedStartFromType.POINT_IN_TIME.value, + self._point_in_time_ms_property_name: + int(self._start_time.astimezone(timezone.utc).timestamp() * 1000) + } + + def populate_request_headers(self, request_headers) -> None: + request_headers[http_constants.HttpHeaders.IfModified_since] =\ + self._start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') + + @classmethod + def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': + point_in_time_ms = data.get(cls._point_in_time_ms_property_name) + if point_in_time_ms is None: + raise ValueError(f"Invalid change feed start from {cls._point_in_time_ms_property_name} ") + + point_in_time = datetime.fromtimestamp(point_in_time_ms).astimezone(timezone.utc) + return ChangeFeedStartFromPointInTime(point_in_time) + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py new file mode 100644 index 000000000000..8c61c306b94e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -0,0 +1,279 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed state implementation in the Azure Cosmos +database service. +""" + +import base64 +import collections +import json +from abc import ABC, abstractmethod +from typing import Optional, Union, List, Any + +from azure.cosmos import http_constants +from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromInternal, \ + ChangeFeedStartFromETagAndFeedRange +from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._change_feed.feed_range_composite_continuation_token import FeedRangeCompositeContinuation +from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider +from azure.cosmos._routing.routing_range import Range +from azure.cosmos._utils import is_key_exists_and_not_none +from azure.cosmos.exceptions import CosmosFeedRangeGoneError +from azure.cosmos.partition_key import _Empty, _Undefined + + +class ChangeFeedState(ABC): + version_property_name = "v" + + @abstractmethod + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + pass + + @abstractmethod + def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, request_headers: dict[str, any]) -> None: + pass + + @abstractmethod + def apply_server_response_continuation(self, continuation: str) -> None: + pass + + @staticmethod + def from_json(container_link: str, container_rid: str, data: dict[str, Any]): + if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): + return ChangeFeedStateV1.from_json(container_link, container_rid, data) + else: + if is_key_exists_and_not_none(data, "continuationFeedRange"): + # get changeFeedState from continuation + continuation_json_str = base64.b64decode(data["continuationFeedRange"]).decode('utf-8') + continuation_json = json.loads(continuation_json_str) + version = continuation_json.get(ChangeFeedState.version_property_name) + if version is None: + raise ValueError("Invalid base64 encoded continuation string [Missing version]") + elif version == "V2": + return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) + else: + raise ValueError("Invalid base64 encoded continuation string [Invalid version]") + # when there is no continuation token, by default construct ChangeFeedStateV2 + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) + +class ChangeFeedStateV1(ChangeFeedState): + """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag + """ + + def __init__( + self, + container_link: str, + container_rid: str, + change_feed_start_from: ChangeFeedStartFromInternal, + partition_key_range_id: Optional[str] = None, + partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, + continuation: Optional[str] = None): + + self._container_link = container_link + self._container_rid = container_rid + self._change_feed_start_from = change_feed_start_from + self._partition_key_range_id = partition_key_range_id + self._partition_key = partition_key + self._continuation = continuation + + @property + def container_rid(self): + return self._container_rid + + @classmethod + def from_json(cls, container_link: str, container_rid: str, data: dict[str, Any]) -> 'ChangeFeedStateV1': + return cls( + container_link, + container_rid, + ChangeFeedStartFromInternal.from_start_time(data.get("startTime")), + data.get("partitionKeyRangeId"), + data.get("partitionKey"), + data.get("continuationPkRangeId") + ) + + def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, Any]) -> None: + headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(headers) + if self._continuation: + headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation + + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + if self._partition_key_range_id is not None: + feed_options["partitionKeyRangeId"] = self._partition_key_range_id + if self._partition_key is not None: + feed_options["partitionKey"] = self._partition_key + + def apply_server_response_continuation(self, continuation: str) -> None: + self._continuation = continuation + +class ChangeFeedStateV2(ChangeFeedState): + container_rid_property_name = "containerRid" + change_feed_mode_property_name = "mode" + change_feed_start_from_property_name = "startFrom" + continuation_property_name = "continuation" + + # TODO: adding change feed mode + def __init__( + self, + container_link: str, + container_rid: str, + feed_range: Range, + change_feed_start_from: ChangeFeedStartFromInternal, + continuation: Optional[FeedRangeCompositeContinuation] = None): + + self._container_link = container_link + self._container_rid = container_rid + self._feed_range = feed_range + self._change_feed_start_from = change_feed_start_from + self._continuation = continuation + if self._continuation is None: + composite_continuation_token_queue = collections.deque() + composite_continuation_token_queue.append(CompositeContinuationToken(self._feed_range, None)) + self._continuation =\ + FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) + + @property + def container_rid(self) -> str : + return self._container_rid + + def to_dict(self) -> dict[str, Any]: + return { + self.version_property_name: "V2", + self.container_rid_property_name: self._container_rid, + self.change_feed_mode_property_name: "Incremental", + self.change_feed_start_from_property_name: self._change_feed_start_from.to_dict(), + self.continuation_property_name: self._continuation.to_dict() + } + + def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, any]) -> None: + headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(headers) + + if self._continuation.current_token is not None and self._continuation.current_token.token is not None: + change_feed_start_from_feed_range_and_etag =\ + ChangeFeedStartFromETagAndFeedRange(self._continuation.current_token.token, self._continuation.current_token.feed_range) + change_feed_start_from_feed_range_and_etag.populate_request_headers(headers) + + # based on the feed range to find the overlapping partition key range id + over_lapping_ranges =\ + routing_provider.get_overlapping_ranges( + self._container_link, + [self._continuation.current_token.feed_range]) + + if len(over_lapping_ranges) > 1: + raise CosmosFeedRangeGoneError(message= + f"Range {self._continuation.current_token.feed_range}" + f" spans {len(over_lapping_ranges)}" + f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + else: + overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) + if overlapping_feed_range == self._continuation.current_token.feed_range: + # exactly mapping to one physical partition, only need to set the partitionKeyRangeId + headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + else: + # the current token feed range spans less than single physical partition + # for this case, need to set both the partition key range id and epk filter headers + headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + headers[http_constants.HttpHeaders.StartEpkString] = self._continuation.current_token.feed_range.min + headers[http_constants.HttpHeaders.EndEpkString] = self._continuation.current_token.feed_range.max + + def populate_feed_options(self, feed_options: dict[str, any]) -> None: + pass + + def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, resource_link: str) -> None: + self._continuation.handle_feed_range_gone(routing_provider, resource_link) + + def apply_server_response_continuation(self, continuation: str) -> None: + self._continuation.apply_server_response_continuation(continuation) + + def should_retry_on_not_modified_response(self): + self._continuation.should_retry_on_not_modified_response() + + def apply_not_modified_response(self) -> None: + self._continuation.apply_not_modified_response() + + @classmethod + def from_continuation( + cls, + container_link: str, + container_rid: str, + continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': + + container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) + if container_rid_from_continuation is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.container_rid_property_name}]") + elif container_rid_from_continuation != container_rid: + raise ValueError("Invalid continuation: [Mismatch collection rid]") + + change_feed_start_from_data = continuation_json.get(ChangeFeedStateV2.change_feed_start_from_property_name) + if change_feed_start_from_data is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.change_feed_start_from_property_name}]") + change_feed_start_from = ChangeFeedStartFromInternal.from_json(change_feed_start_from_data) + + continuation_data = continuation_json.get(ChangeFeedStateV2.continuation_property_name) + if continuation_data is None: + raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.continuation_property_name}]") + continuation = FeedRangeCompositeContinuation.from_json(continuation_data) + return ChangeFeedStateV2( + container_link=container_link, + container_rid=container_rid, + feed_range=continuation.feed_range, + change_feed_start_from=change_feed_start_from, + continuation=continuation) + + @classmethod + def from_initial_state( + cls, + container_link: str, + collection_rid: str, + data: dict[str, Any]) -> 'ChangeFeedStateV2': + + if is_key_exists_and_not_none(data, "feedRange"): + feed_range_str = base64.b64decode(data["feedRange"]).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range = Range.ParseFromDict(feed_range_json) + elif is_key_exists_and_not_none(data, "partitionKeyFeedRange"): + feed_range = data["partitionKeyFeedRange"] + else: + # default to full range + feed_range = Range( + "", + "FF", + True, + False) + + change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(data.get("startTime")) + return cls( + container_link=container_link, + container_rid=collection_rid, + feed_range=feed_range, + change_feed_start_from=change_feed_start_from, + continuation=None) + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py new file mode 100644 index 000000000000..9945405e4b57 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py @@ -0,0 +1,72 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed composite continuation token in the Azure Cosmos +database service. +""" +from typing import Optional + +from azure.cosmos._routing.routing_range import Range + + +class CompositeContinuationToken(object): + token_property_name = "token" + feed_range_property_name = "range" + + def __init__(self, feed_range: Range, token: Optional[str] = None): + if feed_range is None: + raise ValueError("Missing required parameter feed_range") + + self._token = token + self._feed_range = feed_range + + def to_dict(self): + return { + self.token_property_name: self._token, + self.feed_range_property_name: self._feed_range.to_dict() + } + + @property + def feed_range(self): + return self._feed_range + + @property + def token(self): + return self._token + + def update_token(self, etag): + self._token = etag + + @classmethod + def from_json(cls, data): + token = data.get(cls.token_property_name) + if token is None: + raise ValueError(f"Invalid composite token [Missing {cls.token_property_name}]") + + feed_range_data = data.get(cls.feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid composite token [Missing {cls.feed_range_property_name}]") + + feed_range = Range.ParseFromDict(feed_range_data) + return cls(feed_range=feed_range, token=token) + + def __repr__(self): + return f"CompositeContinuationToken(token={self.token}, range={self._feed_range})" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py new file mode 100644 index 000000000000..2461436924aa --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -0,0 +1,134 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for change feed continuation token by feed range in the Azure Cosmos +database service. +""" +import collections +from collections import deque +from typing import Any + +from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider +from azure.cosmos._routing.routing_range import Range + + +class FeedRangeCompositeContinuation(object): + _version_property_name = "V" + _container_rid_property_name = "Rid" + _continuation_property_name = "Continuation" + _feed_range_property_name = "Range" + + def __init__( + self, + container_rid: str, + feed_range: Range, + continuation: collections.deque[CompositeContinuationToken]): + if container_rid is None: + raise ValueError("container_rid is missing") + + self._container_rid = container_rid + self._feed_range = feed_range + self._continuation = continuation + self._current_token = self._continuation[0] + self._initial_no_result_range = None + + @property + def current_token(self): + return self._current_token + + def to_dict(self) -> dict[str, Any]: + return { + self._version_property_name: "v1", #TODO: should this start from v2 + self._container_rid_property_name: self._container_rid, + self._continuation_property_name: [childToken.to_dict() for childToken in self._continuation], + self._feed_range_property_name: self._feed_range.to_dict() + } + + @classmethod + def from_json(cls, data) -> 'FeedRangeCompositeContinuation': + version = data.get(cls._version_property_name) + if version is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._version_property_name}]") + if version != "v1": + raise ValueError("Invalid feed range composite continuation token [Invalid version]") + + container_rid = data.get(cls._container_rid_property_name) + if container_rid is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") + + feed_range_data = data.get(cls._feed_range_property_name) + if feed_range_data is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._feed_range_property_name}]") + feed_range = Range.ParseFromDict(feed_range_data) + + continuation_data = data.get(cls._continuation_property_name) + if continuation_data is None: + raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") + if not isinstance(continuation_data, list) or len(continuation_data) == 0: + raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") + continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] + + return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) + + def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: + overlapping_ranges = routing_provider.get_overlapping_ranges(collection_link, self._current_token.feed_range) + + if len(overlapping_ranges) == 1: + # merge,reusing the existing the feedRange and continuationToken + pass + else: + # split, remove the parent range and then add new child ranges. + # For each new child range, using the continuation token from the parent + self._continuation.popleft() + for child_range in overlapping_ranges: + self._continuation.append(CompositeContinuationToken(Range.PartitionKeyRangeToRange(child_range), self._current_token.token)) + + self._current_token = self._continuation[0] + + def should_retry_on_not_modified_response(self) -> bool: + # when getting 304(Not Modified) response from one sub feed range, we will try to fetch for the next sub feed range + # we will repeat the above logic until we have looped through all sub feed ranges + + # TODO: validate the response headers, can we get the status code + if len(self._continuation) > 1: + return self._current_token.feed_range != self._initial_no_result_range + + else: + return False + + def _move_to_next_token(self) -> None: + first_composition_token = self._continuation.popleft() + # add the composition token to the end of the list + self._continuation.append(first_composition_token) + self._current_token = self._continuation[0] + + def apply_server_response_continuation(self, etag) -> None: + self._current_token.update_token(etag) + self._move_to_next_token() + + def apply_not_modified_response(self) -> None: + if self._initial_no_result_range is None: + self._initial_no_result_range = self._current_token.feed_range + + @property + def feed_range(self) -> Range: + return self._feed_range 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 1288e7a4e66e..a81ab438cbf2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -26,14 +26,10 @@ import os import urllib.parse from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast, Type -from typing_extensions import TypedDict -from urllib3.util.retry import Retry +from azure.core import PipelineClient from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged -from azure.core import PipelineClient -from azure.core.pipeline.transport import HttpRequest, \ - HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import from azure.core.pipeline.policies import ( HTTPPolicy, ContentDecodePolicy, @@ -44,22 +40,31 @@ DistributedTracingPolicy, ProxyPolicy ) +from azure.core.pipeline.transport import HttpRequest, \ + HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import +from typing_extensions import TypedDict +from urllib3.util.retry import Retry from . import _base as base -from ._base import _set_properties_cache -from . import documents -from .documents import ConnectionPolicy, DatabaseAccount -from ._constants import _Constants as Constants -from . import http_constants, exceptions +from . import _global_endpoint_manager as global_endpoint_manager from . import _query_iterable as query_iterable from . import _runtime_constants as runtime_constants -from ._request_object import RequestObject -from . import _synchronized_request as synchronized_request -from . import _global_endpoint_manager as global_endpoint_manager -from ._routing import routing_map_provider, routing_range -from ._retry_utility import ConnectionRetryPolicy from . import _session +from . import _synchronized_request as synchronized_request from . import _utils +from . import documents +from . import http_constants, exceptions +from ._auth_policy import CosmosBearerTokenCredentialPolicy +from ._base import _set_properties_cache +from ._change_feed.change_feed_iterable import ChangeFeedIterable +from ._change_feed.change_feed_state import ChangeFeedState +from ._constants import _Constants as Constants +from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy +from ._range_partition_resolver import RangePartitionResolver +from ._request_object import RequestObject +from ._retry_utility import ConnectionRetryPolicy +from ._routing import routing_map_provider, routing_range +from .documents import ConnectionPolicy, DatabaseAccount from .partition_key import ( _Undefined, _Empty, @@ -67,9 +72,6 @@ _return_undefined_or_empty_partition_key, NonePartitionKeyValue ) -from ._auth_policy import CosmosBearerTokenCredentialPolicy -from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy -from ._range_partition_resolver import RangePartitionResolver PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long @@ -1191,11 +1193,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str return ItemPaged( self, - None, options, fetch_function=fetch_fn, collection_link=collection_link, - page_iterator_class=query_iterable.QueryIterable + page_iterator_class=ChangeFeedIterable ) def _ReadPartitionKeyRanges( @@ -3023,6 +3024,11 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: options, partition_key_range_id ) + + change_feed_state = options.pop("changeFeedState", None) + if change_feed_state and isinstance(change_feed_state, ChangeFeedState): + change_feed_state.populate_request_headers(self._routing_map_provider, headers) + result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers if response_hook: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index 0d61fbbbe1d7..f3269af47271 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -22,6 +22,9 @@ """Internal class for partition key range implementation in the Azure Cosmos database service. """ +import base64 +import binascii +import json class PartitionKeyRange(object): @@ -81,6 +84,76 @@ def ParseFromDict(cls, range_as_dict): ) return self + def to_dict(self): + return { + self.MinPath: self.min, + self.MaxPath: self.max, + self.IsMinInclusivePath: self.isMinInclusive, + self.IsMaxInclusivePath: self.isMaxInclusive + } + + def to_normalized_range(self): + if self.isMinInclusive and not self.isMaxInclusive: + return self + + normalized_min = self.min + normalized_max = self.max + + if not self.isMinInclusive: + normalized_min = self.add_to_effective_partition_key(self.min, -1) + + if self.isMaxInclusive: + normalized_max = self.add_to_effective_partition_key(self.max, 1) + + return Range(normalized_min, normalized_max, True, False) + + def add_to_effective_partition_key(self, effective_partition_key: str, value: int): + if value != 1 and value != -1: + raise ValueError("Invalid value - only 1 or -1 is allowed") + + byte_array = self.hex_binary_to_byte_array(effective_partition_key) + if value == 1: + for i in range(len(byte_array) -1, -1, -1): + if byte_array[i] < 255: + byte_array[i] += 1 + break + else: + byte_array[i] = 0 + else: + for i in range(len(byte_array) - 1, -1, -1): + if byte_array[i] != 0: + byte_array[i] -= 1 + break + else: + byte_array[i] = 255 + + return binascii.hexlify(byte_array).decode() + + def hex_binary_to_byte_array(self, hex_binary_string: str): + if hex_binary_string is None: + raise ValueError("hex_binary_string is missing") + if len(hex_binary_string) % 2 != 0: + raise ValueError("hex_binary_string must not have an odd number of characters") + + return bytearray.fromhex(hex_binary_string) + + @classmethod + def from_base64_encoded_json_string(cls, data: str): + try: + feed_range_json_string = base64.b64decode(data, validate=True).decode('utf-8') + feed_range_json = json.loads(feed_range_json_string) + return cls.ParseFromDict(feed_range_json) + except Exception: + raise ValueError(f"Invalid feed_range json string {data}") + + def to_base64_encoded_string(self): + data_json = json.dumps(self.to_dict()) + json_bytes = data_json.encode('utf-8') + # Encode the bytes to a Base64 string + base64_bytes = base64.b64encode(json_bytes) + # Convert the Base64 bytes to a string + return base64_bytes.decode('utf-8') + def isSingleValue(self): return self.isMinInclusive and self.isMaxInclusive and self.min == self.max diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 1b3d0370e6ef..1c03b8a054c5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -69,3 +69,17 @@ def get_index_metrics_info(delimited_string: Optional[str]) -> Dict[str, Any]: return result except (json.JSONDecodeError, ValueError): return {} + + +def is_base64_encoded(data: str) -> bool: + if data is None: + return False + try: + base64.b64decode(data, validate=True).decode('utf-8') + return True + except (json.JSONDecodeError, ValueError): + return False + + +def is_key_exists_and_not_none(data: dict[str, Any], key: str) -> bool: + return key in data and data[key] is not None diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 385d7f7af236..a8559839aad7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -21,16 +21,18 @@ """Create, read, update and delete items in the Azure Cosmos DB SQL API service. """ -from datetime import datetime, timezone -from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast -from typing_extensions import Literal +import warnings +from datetime import datetime +from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload from azure.core import MatchConditions from azure.core.async_paging import AsyncItemPaged from azure.core.tracing.decorator import distributed_trace from azure.core.tracing.decorator_async import distributed_trace_async # type: ignore +from typing_extensions import Literal from ._cosmos_client_connection_async import CosmosClientConnection +from ._scripts import ScriptsProxy from .._base import ( build_options as _build_options, validate_cache_staleness_value, @@ -39,13 +41,16 @@ GenerateGuidId, _set_properties_cache ) +from .._change_feed.aio.change_feed_state import ChangeFeedState +from .._routing import routing_range +from .._routing.routing_range import Range +from .._utils import is_key_exists_and_not_none, is_base64_encoded from ..offer import ThroughputProperties -from ._scripts import ScriptsProxy from ..partition_key import ( NonePartitionKeyValue, _return_undefined_or_empty_partition_key, _Empty, - _Undefined + _Undefined, PartitionKey ) __all__ = ("ContainerProxy",) @@ -132,6 +137,26 @@ async def _set_partition_key( return _return_undefined_or_empty_partition_key(await self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) + async def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: + container_properties = await self._get_properties() + partition_key_definition = container_properties.get("partitionKey") + partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) + + is_prefix_partition_key = await self.__is_prefix_partition_key(partition_key_value) + + return partition_key._get_epk_range_for_partition_key(partition_key_value, is_prefix_partition_key) + + async def __is_prefix_partition_key(self, partition_key: PartitionKeyType) -> bool: + + properties = await self._get_properties() + pk_properties = properties.get("partitionKey") + partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) + if partition_key_definition.kind != "MultiHash": + return False + if isinstance(partition_key, list) and len(partition_key_definition['paths']) == len(partition_key): + return False + return True + @distributed_trace_async async def read( self, @@ -480,62 +505,196 @@ def query_items( response_hook(self.client_connection.last_response_headers, items) return items - @distributed_trace - def query_items_change_feed( - self, - *, - partition_key_range_id: Optional[str] = None, - is_start_from_beginning: bool = False, - start_time: Optional[datetime] = None, - continuation: Optional[str] = None, - max_item_count: Optional[int] = None, - partition_key: Optional[PartitionKeyType] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + @overload + async def query_items_change_feed( + self, + *, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: Optional[PartitionKeyType] = None, + # -> would RU usage be more efficient, bug to backend team? deprecate it or using FeedRange to convert? + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword bool is_start_from_beginning: Get whether change feed should start from - beginning (true) or from current (false). By default, it's start from current (false). - :keyword ~datetime.datetime start_time: Specifies a point of time to start change feed. Provided value will be - converted to UTC. This value will be ignored if `is_start_from_beginning` is set to true. - :keyword str partition_key_range_id: ChangeFeed requests can be executed against specific partition key - ranges. This is used to process the change feed in parallel across multiple consumers. - :keyword str continuation: e_tag value to be used as continuation for reading change feed. - :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword partition_key: partition key at which ChangeFeed requests are targeted. - :paramtype partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] - :keyword response_hook: A callable invoked with the response metadata. - :paramtype response_hook: Callable[[Dict[str, str], AsyncItemPaged[Dict[str, Any]]], None] - :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :param int max_item_count: Max number of items to be returned in the enumeration operation. + :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :param PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) + :param 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. :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ - response_hook = kwargs.pop('response_hook', None) - if priority is not None: - kwargs['priority'] = priority + ... + + @overload + async def query_items_change_feed( + self, + *, + feed_range: Optional[str] = None, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> AsyncItemPaged[Dict[str, Any]]: + """Get a sorted list of items that were changed, in the order in which they were modified. + + :param str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. + :param int max_item_count: Max number of items to be returned in the enumeration operation. + :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :param 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. + :returns: An AsyncItemPaged of items (dicts). + :rtype: AsyncItemPaged[Dict[str, Any]] + """ + ... + + @overload + async def query_items_change_feed( + self, + *, + continuation: Optional[str] = None, + max_item_count: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> AsyncItemPaged[Dict[str, Any]]: + """Get a sorted list of items that were changed, in the order in which they were modified. + + :param str continuation: The continuation token retrieved from previous response. + :param int max_item_count: Max number of items to be returned in the enumeration operation. + :param 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. + :returns: An AsyncItemPaged of items (dicts). + :rtype: AsyncItemPaged[Dict[str, Any]] + """ + ... + + @distributed_trace + async def query_items_change_feed( + self, + *args: Any, + **kwargs: Any + ) -> AsyncItemPaged[Dict[str, Any]]: + + if is_key_exists_and_not_none(kwargs, "priority"): + kwargs['priority'] = kwargs['priority'] feed_options = _build_options(kwargs) - feed_options["isStartFromBeginning"] = is_start_from_beginning - if start_time is not None and is_start_from_beginning is False: - feed_options["startTime"] = start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') - if partition_key_range_id is not None: - feed_options["partitionKeyRangeId"] = partition_key_range_id - if partition_key is not None: - feed_options["partitionKey"] = self._set_partition_key(partition_key) - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if continuation is not None: - feed_options["continuation"] = continuation + + change_feed_state_context = {} + # Back compatibility with deprecation warnings for partition_key_range_id + if (args and args[0] is not None) or is_key_exists_and_not_none(kwargs, "partition_key_range_id"): + warnings.warn( + "partition_key_range_id is deprecated. Please pass in feed_range instead.", + DeprecationWarning + ) + + try: + change_feed_state_context["partitionKeyRangeId"] = kwargs.pop('partition_key_range_id') + except KeyError: + change_feed_state_context['partitionKeyRangeId'] = args[0] + + # Back compatibility with deprecation warnings for is_start_from_beginning + if (len(args) >= 2 and args[1] is not None) or is_key_exists_and_not_none(kwargs, "is_start_from_beginning"): + warnings.warn( + "is_start_from_beginning is deprecated. Please pass in start_time instead.", + DeprecationWarning + ) + + try: + is_start_from_beginning = kwargs.pop('is_start_from_beginning') + except KeyError: + is_start_from_beginning = args[1] + + if is_start_from_beginning: + change_feed_state_context["startTime"] = "Beginning" + + # parse start_time + if is_key_exists_and_not_none(kwargs, "start_time"): + if change_feed_state_context.get("startTime") is not None: + raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") + + start_time = kwargs.pop('start_time') + if not isinstance(start_time, (datetime, str)): + raise TypeError( + "'start_time' must be either a datetime object, or either the values 'now' or 'beginning'.") + change_feed_state_context["startTime"] = start_time + + # parse continuation token + if len(args) >= 3 and args[2] is not None or is_key_exists_and_not_none(feed_options, "continuation"): + try: + continuation = feed_options.pop('continuation') + except KeyError: + continuation = args[2] + + # there are two types of continuation token we support currently: + # v1 version: the continuation token would just be the _etag, + # which is being returned when customer is using partition_key_range_id, + # which is under deprecation and does not support split/merge + # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + if is_base64_encoded(continuation): + change_feed_state_context["continuationFeedRange"] = continuation + else: + change_feed_state_context["continuationPkRangeId"] = continuation + + if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): + try: + feed_options["maxItemCount"] = kwargs.pop('max_item_count') + except KeyError: + feed_options["maxItemCount"] = args[3] + + if is_key_exists_and_not_none(kwargs, "partition_key"): + partition_key = kwargs.pop("partition_key") + change_feed_state_context["partitionKey"] = await self._set_partition_key(partition_key) + change_feed_state_context["partitionKeyFeedRange"] = await self._get_epk_range_for_partition_key(partition_key) + + if is_key_exists_and_not_none(kwargs, "feed_range"): + change_feed_state_context["feedRange"] = kwargs.pop('feed_range') + + # validate exclusive or in-compatible parameters + if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + # if continuation token is in v1 format, throw exception if feed_range is set + if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + raise ValueError("feed_range and continuation are incompatible") + elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + # if continuation token is in v2 format, since the token itself contains the full change feed state + # so we will ignore other parameters if they passed in + if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId"): + raise ValueError("partition_key_range_id and continuation are incompatible") + else: + # validation when no continuation is passed + exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] + count = sum(1 for key in exclusive_keys if + key in change_feed_state_context and change_feed_state_context[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") + + container_properties = await self._get_properties() + container_rid = container_properties.get("_rid") + change_feed_state = ChangeFeedState.from_json(self.container_link, container_rid, change_feed_state_context) + feed_options["changeFeedState"] = change_feed_state + feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + + response_hook = kwargs.pop('response_hook', None) if 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"] result = self.client_connection.QueryItemsChangeFeed( self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) + if response_hook: response_hook(self.client_connection.last_response_headers, result) return result @@ -1098,3 +1257,17 @@ async def execute_item_batch( return await self.client_connection.Batch( collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) + + async def read_feed_ranges( + self, + **kwargs: Any + ) -> List[str]: + partition_key_ranges =\ + await self.client_connection._routing_map_provider.get_overlapping_ranges( + self.container_link, + # default to full range + [Range("", "FF", True, False)]) + + return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] + + 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 72ea03668909..7cccac695769 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 @@ -50,6 +50,8 @@ from .. import _base as base from .._base import _set_properties_cache from .. import documents +from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable +from .._change_feed.aio.change_feed_state import ChangeFeedState from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants @@ -2310,11 +2312,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di return AsyncItemPaged( self, - None, options, fetch_function=fetch_fn, collection_link=collection_link, - page_iterator_class=query_iterable.QueryIterable + page_iterator_class=ChangeFeedIterable ) def QueryOffers( @@ -2812,6 +2813,11 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed ) headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) + + change_feed_state = options.pop("changeFeedState", None) + if change_feed_state and isinstance(change_feed_state, ChangeFeedState): + await change_feed_state.populate_request_headers(self._routing_map_provider, headers) + result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) if response_hook: response_hook(self.last_response_headers, result) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 9cf697ee7f5a..32fd818075f8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -21,16 +21,15 @@ """Create, read, update and delete items in the Azure Cosmos DB SQL API service. """ -from datetime import datetime, timezone import warnings -from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast -from typing_extensions import Literal +from datetime import datetime +from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload from azure.core import MatchConditions -from azure.core.tracing.decorator import distributed_trace from azure.core.paging import ItemPaged +from azure.core.tracing.decorator import distributed_trace +from typing_extensions import Literal -from ._cosmos_client_connection import CosmosClientConnection from ._base import ( build_options, validate_cache_staleness_value, @@ -39,8 +38,12 @@ GenerateGuidId, _set_properties_cache ) +from ._change_feed.change_feed_state import ChangeFeedState +from ._cosmos_client_connection import CosmosClientConnection +from ._routing import routing_range +from ._routing.routing_range import Range +from ._utils import is_key_exists_and_not_none, is_base64_encoded from .offer import Offer, ThroughputProperties -from .scripts import ScriptsProxy from .partition_key import ( NonePartitionKeyValue, PartitionKey, @@ -48,6 +51,7 @@ _Undefined, _return_undefined_or_empty_partition_key ) +from .scripts import ScriptsProxy __all__ = ("ContainerProxy",) @@ -132,6 +136,13 @@ def _set_partition_key( def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: return self.client_connection._container_properties_cache + def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: + container_properties = self._get_properties() + partition_key_definition = container_properties.get("partitionKey") + partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) + + return partition_key._get_epk_range_for_partition_key(partition_key_value, self.__is_prefix_partitionkey(partition_key_value)) + @distributed_trace def read( # pylint:disable=docstring-missing-param self, @@ -309,56 +320,186 @@ def read_all_items( # pylint:disable=docstring-missing-param response_hook(self.client_connection.last_response_headers, items) return items - @distributed_trace + @overload def query_items_change_feed( - self, - partition_key_range_id: Optional[str] = None, - is_start_from_beginning: bool = False, - continuation: Optional[str] = None, - max_item_count: Optional[int] = None, - *, - start_time: Optional[datetime] = None, - partition_key: Optional[PartitionKeyType] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + *, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: Optional[PartitionKeyType] = None, + # -> would RU usage be more efficient, bug to backend team? deprecate it or using FeedRange to convert? + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. - :param str partition_key_range_id: ChangeFeed requests can be executed against specific partition key ranges. - This is used to process the change feed in parallel across multiple consumers. - :param bool is_start_from_beginning: Get whether change feed should start from - beginning (true) or from current (false). By default, it's start from current (false). - :param max_item_count: Max number of items to be returned in the enumeration operation. - :param str continuation: e_tag value to be used as continuation for reading change feed. :param int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword ~datetime.datetime start_time: Specifies a point of time to start change feed. Provided value will be - converted to UTC. This value will be ignored if `is_start_from_beginning` is set to true. - :keyword partition_key: partition key at which ChangeFeed requests are targeted. - :paramtype partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] - :keyword Callable response_hook: A callable invoked with the response metadata. - :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :param PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) + :param 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. :returns: An Iterable of items (dicts). - :rtype: Iterable[dict[str, Any]] + :rtype: Iterable[Dict[str, Any]] """ - if priority is not None: - kwargs['priority'] = priority + ... + + @overload + def query_items_change_feed( + self, + *, + feed_range: Optional[str] = None, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> ItemPaged[Dict[str, Any]]: + """Get a sorted list of items that were changed, in the order in which they were modified. + + :param str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. + :param int max_item_count: Max number of items to be returned in the enumeration operation. + :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :param 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. + :returns: An Iterable of items (dicts). + :rtype: Iterable[Dict[str, Any]] + """ + ... + + @overload + def query_items_change_feed( + self, + *, + continuation: Optional[str] = None, + max_item_count: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> ItemPaged[Dict[str, Any]]: + """Get a sorted list of items that were changed, in the order in which they were modified. + + :param str continuation: The continuation token retrieved from previous response. + :param int max_item_count: Max number of items to be returned in the enumeration operation. + :param 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. + :returns: An Iterable of items (dicts). + :rtype: Iterable[Dict[str, Any]] + """ + ... + + @distributed_trace + def query_items_change_feed( + self, + *args: Any, + **kwargs: Any + ) -> ItemPaged[Dict[str, Any]]: + + if is_key_exists_and_not_none(kwargs, "priority"): + kwargs['priority'] = kwargs['priority'] feed_options = build_options(kwargs) - response_hook = kwargs.pop('response_hook', None) - if partition_key_range_id is not None: - feed_options["partitionKeyRangeId"] = partition_key_range_id - if partition_key is not None: - feed_options["partitionKey"] = self._set_partition_key(partition_key) - if is_start_from_beginning is not None: - feed_options["isStartFromBeginning"] = is_start_from_beginning - if start_time is not None and is_start_from_beginning is False: - feed_options["startTime"] = start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') - if max_item_count is not None: - feed_options["maxItemCount"] = max_item_count - if continuation is not None: - feed_options["continuation"] = continuation + change_feed_state_context = {} + # Back compatibility with deprecation warnings for partition_key_range_id + if (args and args[0] is not None) or is_key_exists_and_not_none(kwargs, "partition_key_range_id"): + warnings.warn( + "partition_key_range_id is deprecated. Please pass in feed_range instead.", + DeprecationWarning + ) + + try: + change_feed_state_context["partitionKeyRangeId"] = kwargs.pop('partition_key_range_id') + except KeyError: + change_feed_state_context['partitionKeyRangeId'] = args[0] + + # Back compatibility with deprecation warnings for is_start_from_beginning + if (len(args) >= 2 and args[1] is not None) or is_key_exists_and_not_none(kwargs, "is_start_from_beginning"): + warnings.warn( + "is_start_from_beginning is deprecated. Please pass in start_time instead.", + DeprecationWarning + ) + + try: + is_start_from_beginning = kwargs.pop('is_start_from_beginning') + except KeyError: + is_start_from_beginning = args[1] + + if is_start_from_beginning: + change_feed_state_context["startTime"] = "Beginning" + + # parse start_time + if is_key_exists_and_not_none(kwargs, "start_time"): + if change_feed_state_context.get("startTime") is not None: + raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") + + start_time = kwargs.pop('start_time') + if not isinstance(start_time, (datetime, str)): + raise TypeError( + "'start_time' must be either a datetime object, or either the values 'now' or 'beginning'.") + change_feed_state_context["startTime"] = start_time + + # parse continuation token + if len(args) >= 3 and args[2] is not None or is_key_exists_and_not_none(feed_options, "continuation"): + try: + continuation = feed_options.pop('continuation') + except KeyError: + continuation = args[2] + + # there are two types of continuation token we support currently: + # v1 version: the continuation token would just be the _etag, + # which is being returned when customer is using partition_key_range_id, + # which is under deprecation and does not support split/merge + # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + if is_base64_encoded(continuation): + change_feed_state_context["continuationFeedRange"] = continuation + else: + change_feed_state_context["continuationPkRangeId"] = continuation + + if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): + try: + feed_options["maxItemCount"] = kwargs.pop('max_item_count') + except KeyError: + feed_options["maxItemCount"] = args[3] + + if is_key_exists_and_not_none(kwargs, "partition_key"): + partition_key = kwargs.pop("partition_key") + change_feed_state_context["partitionKey"] = self._set_partition_key(partition_key) + change_feed_state_context["partitionKeyFeedRange"] = self._get_epk_range_for_partition_key(partition_key) + + if is_key_exists_and_not_none(kwargs, "feed_range"): + change_feed_state_context["feedRange"] = kwargs.pop('feed_range') + + # validate exclusive or in-compatible parameters + if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + # if continuation token is in v1 format, throw exception if feed_range is set + if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + raise ValueError("feed_range and continuation are incompatible") + elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + # if continuation token is in v2 format, since the token itself contains the full change feed state + # so we will ignore other parameters if they passed in + if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId"): + raise ValueError("partition_key_range_id and continuation are incompatible") + else: + # validation when no continuation is passed + exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] + count = sum(1 for key in exclusive_keys if key in change_feed_state_context and change_feed_state_context[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") + + container_properties = self._get_properties() + container_rid = container_properties.get("_rid") + change_feed_state = ChangeFeedState.from_json(self.container_link, container_rid, change_feed_state_context) + feed_options["changeFeedState"] = change_feed_state + + response_hook = kwargs.pop('response_hook', None) if hasattr(response_hook, "clear"): response_hook.clear() if self.container_link in self.__get_client_container_caches(): @@ -1162,3 +1303,15 @@ def delete_all_items_by_partition_key( self.client_connection.DeleteAllItemsByPartitionKey( collection_link=self.container_link, options=request_options, **kwargs) + + def read_feed_ranges( + self, + **kwargs: Any + ) -> List[str]: + partition_key_ranges =\ + self.client_connection._routing_map_provider.get_overlapping_ranges( + self.container_link, + # default to full range + [Range("", "FF", True, False)]) + + return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py index 5092fd0de7cf..768890dacfa6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py @@ -28,7 +28,7 @@ ResourceNotFoundError ) from . import http_constants - +from .http_constants import StatusCodes, SubStatusCodes class CosmosHttpResponseError(HttpResponseError): """An HTTP request to the Azure Cosmos database service has failed.""" @@ -136,6 +136,19 @@ def __init__(self, **kwargs): super(CosmosClientTimeoutError, self).__init__(message, **kwargs) +class CosmosFeedRangeGoneError(CosmosHttpResponseError): + """An HTTP error response with status code 404.""" + def __init__(self, message=None, response=None, **kwargs): + """ + :param int sub_status_code: HTTP response sub code. + """ + self.status_code = StatusCodes.GONE + self.sub_status = SubStatusCodes.PARTITION_KEY_RANGE_GONE + self.http_error_message = message + formatted_message = "Status code: %d Sub-status: %d\n%s" % (self.status_code, self.sub_status, str(message)) + super(CosmosHttpResponseError, self).__init__(message=formatted_message, response=response, **kwargs) + + def _partition_range_is_gone(e): if (e.status_code == http_constants.StatusCodes.GONE and e.sub_status == http_constants.SubStatusCodes.PARTITION_KEY_RANGE_GONE): @@ -151,3 +164,7 @@ def _container_recreate_exception(e) -> bool: is_throughput_not_found = e.sub_status == http_constants.SubStatusCodes.THROUGHPUT_OFFER_NOT_FOUND return (is_bad_request and is_collection_rid_mismatch) or (is_not_found and is_throughput_not_found) + + +def _is_partition_split_or_merge(e): + return e.status_code == StatusCodes.GONE and e.status_code == SubStatusCodes.COMPLETING_SPLIT \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 22fcb19dae06..9f0a5cde29a2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -173,6 +173,20 @@ 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( + self, + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]], + is_prefix_pk_value: bool = False + ) -> _Range: + if is_prefix_pk_value: + return self._get_epk_range_for_prefix_partition_key(pk_value) + + # else return point range + effective_partition_key_string = self._get_effective_partition_key_string(pk_value) + partition_key_range = _Range(effective_partition_key_string, effective_partition_key_string, True, True) + + return partition_key_range.to_normalized_range() + def _get_effective_partition_key_for_hash_partitioning(self) -> str: # We shouldn't be supporting V1 return "" diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py new file mode 100644 index 000000000000..a1d34262cae7 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -0,0 +1,295 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import time +import unittest +import uuid +from datetime import datetime, timedelta, timezone +from time import sleep + +import pytest +from _pytest.outcomes import fail + +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +import test_config +from azure.cosmos import DatabaseProxy +from azure.cosmos.partition_key import PartitionKey + + +@pytest.fixture(scope="class") +def setup(): + if (TestChangeFeed.masterKey == '[YOUR_KEY_HERE]' or + TestChangeFeed.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = cosmos_client.CosmosClient(test_config.TestConfig.host, test_config.TestConfig.masterKey), + return { + "created_db": test_client[0].get_database_client(TestChangeFeed.TEST_DATABASE_ID) + } + +@pytest.mark.cosmosEmulator +@pytest.mark.unittest +@pytest.mark.usefixtures("setup") +class TestChangeFeed: + """Test to ensure escaping of non-ascii characters from partition key""" + + created_db: DatabaseProxy = None + client: cosmos_client.CosmosClient = None + config = test_config.TestConfig + host = config.host + masterKey = config.masterKey + connectionPolicy = config.connectionPolicy + TEST_DATABASE_ID = config.TEST_DATABASE_ID + + def test_get_feed_ranges(self, setup): + created_collection = setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), + PartitionKey(path="/pk")) + result = created_collection.read_feed_ranges() + assert len(result) == 1 + + @pytest.mark.parametrize("change_feed_filter_param", ["partitionKey", "partitionKeyRangeId", "feedRange"]) + def test_query_change_feed_with_different_filter(self, change_feed_filter_param, setup): + created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk")) + + # Read change feed without passing any options + query_iterable = created_collection.query_items_change_feed() + iter_list = list(query_iterable) + assert len(iter_list) == 0 + + if change_feed_filter_param == "partitionKey": + filter_param = {"partition_key": "pk"} + elif change_feed_filter_param == "partitionKeyRangeId": + filter_param = {"partition_key_range_id": "0"} + elif change_feed_filter_param == "feedRange": + feed_ranges = created_collection.read_feed_ranges() + assert len(feed_ranges) == 1 + filter_param = {"feed_range": feed_ranges[0]} + else: + filter_param = None + + # Read change feed from current should return an empty list + query_iterable = created_collection.query_items_change_feed(filter_param) + iter_list = list(query_iterable) + assert len(iter_list) == 0 + assert 'etag' in created_collection.client_connection.last_response_headers + assert created_collection.client_connection.last_response_headers['etag'] !='' + + # Read change feed from beginning should return an empty list + query_iterable = created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + iter_list = list(query_iterable) + assert len(iter_list) == 0 + assert 'etag' in created_collection.client_connection.last_response_headers + continuation1 = created_collection.client_connection.last_response_headers['etag'] + assert continuation1 != '' + + # Create a document. Read change feed should return be able to read that document + document_definition = {'pk': 'pk', 'id': 'doc1'} + created_collection.create_item(body=document_definition) + query_iterable = created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + iter_list = list(query_iterable) + assert len(iter_list) == 1 + assert iter_list[0]['id'] == 'doc1' + assert 'etag' in created_collection.client_connection.last_response_headers + continuation2 = created_collection.client_connection.last_response_headers['etag'] + assert continuation2 != '' + assert continuation2 != continuation1 + + # Create two new documents. Verify that change feed contains the 2 new documents + # with page size 1 and page size 100 + document_definition = {'pk': 'pk', 'id': 'doc2'} + created_collection.create_item(body=document_definition) + document_definition = {'pk': 'pk', 'id': 'doc3'} + created_collection.create_item(body=document_definition) + + for pageSize in [1, 100]: + # verify iterator + query_iterable = created_collection.query_items_change_feed( + continuation=continuation2, + max_item_count=pageSize, + **filter_param + ) + it = query_iterable.__iter__() + expected_ids = 'doc2.doc3.' + actual_ids = '' + for item in it: + actual_ids += item['id'] + '.' + assert actual_ids == expected_ids + + # verify by_page + # the options is not copied, therefore it need to be restored + query_iterable = created_collection.query_items_change_feed( + continuation=continuation2, + max_item_count=pageSize, + **filter_param + ) + count = 0 + expected_count = 2 + all_fetched_res = [] + for page in query_iterable.by_page(): + fetched_res = list(page) + assert len(fetched_res) == min(pageSize, expected_count - count) + count += len(fetched_res) + all_fetched_res.extend(fetched_res) + + actual_ids = '' + for item in all_fetched_res: + actual_ids += item['id'] + '.' + assert actual_ids == expected_ids + + # verify reading change feed from the beginning + query_iterable = created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + expected_ids = ['doc1', 'doc2', 'doc3'] + it = query_iterable.__iter__() + for i in range(0, len(expected_ids)): + doc = next(it) + assert doc['id'] == expected_ids[i] + assert 'etag' in created_collection.client_connection.last_response_headers + continuation3 = created_collection.client_connection.last_response_headers['etag'] + + # verify reading empty change feed + query_iterable = created_collection.query_items_change_feed( + continuation=continuation3, + is_start_from_beginning=True, + **filter_param + ) + iter_list = list(query_iterable) + assert len(iter_list) == 0 + setup["created_db"].delete_container(created_collection.id) + + def test_query_change_feed_with_start_time(self, setup): + created_collection = setup["created_db"].create_container_if_not_exists("query_change_feed_start_time_test", + PartitionKey(path="/pk")) + batchSize = 50 + + def round_time(): + utc_now = datetime.now(timezone.utc) + return utc_now - timedelta(microseconds=utc_now.microsecond) + def create_random_items(container, batch_size): + for _ in range(batch_size): + # Generate a Random partition key + partition_key = 'pk' + str(uuid.uuid4()) + + # Generate a random item + item = { + 'id': 'item' + str(uuid.uuid4()), + 'partitionKey': partition_key, + 'content': 'This is some random content', + } + + try: + # Create the item in the container + container.upsert_item(item) + except exceptions.CosmosHttpResponseError as e: + fail(e) + + # Create first batch of random items + create_random_items(created_collection, batchSize) + + # wait for 1 second and record the time, then wait another second + sleep(1) + start_time = round_time() + not_utc_time = datetime.now() + sleep(1) + + # now create another batch of items + create_random_items(created_collection, batchSize) + + # now query change feed based on start time + change_feed_iter = list(created_collection.query_items_change_feed(start_time=start_time)) + totalCount = len(change_feed_iter) + + # now check if the number of items that were changed match the batch size + assert totalCount == batchSize + + # negative test: pass in a valid time in the future + future_time = start_time + timedelta(hours=1) + change_feed_iter = list(created_collection.query_items_change_feed(start_time=future_time)) + totalCount = len(change_feed_iter) + # A future time should return 0 + assert totalCount == 0 + + # test a date that is not utc, will be converted to utc by sdk + change_feed_iter = list(created_collection.query_items_change_feed(start_time=not_utc_time)) + totalCount = len(change_feed_iter) + # Should equal batch size + assert totalCount == batchSize + + # test an invalid value, Attribute error will be raised for passing non datetime object + invalid_time = "Invalid value" + try: + list(created_collection.query_items_change_feed(start_time=invalid_time)) + fail("Cannot format date on a non datetime object.") + except ValueError as e: #TODO: previously it is throwing AttributeError, now has changed into ValueError, is it breaking change? + assert "Invalid start_time 'Invalid value'" == e.args[0] + + setup["created_db"].delete_container(created_collection.id) + + def test_query_change_feed_with_split(self, setup): + created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + + # initial change feed query returns empty result + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") + iter_list = list(query_iterable) + assert len(iter_list) == 0 + continuation = created_collection.client_connection.last_response_headers['etag'] + assert continuation != '' + + # create one doc and make sure change feed query can return the document + document_definition = {'pk': 'pk', 'id': 'doc1'} + created_collection.create_item(body=document_definition) + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + iter_list = list(query_iterable) + assert len(iter_list) == 1 + continuation = created_collection.client_connection.last_response_headers['etag'] + + print("Triggering a split in test_query_change_feed_with_split") + created_collection.replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = created_collection.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in test_query_change_feed_with_split has completed") + print("creating few more documents") + new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc2', 'doc3', 'doc4'] + for document in new_documents: + created_collection.create_item(body=document) + + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + it = query_iterable.__iter__() + actual_ids = [] + for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + setup["created_db"].delete_container(created_collection.id) + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py new file mode 100644 index 000000000000..c3246768a796 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -0,0 +1,322 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import time +import unittest +import uuid +from asyncio import sleep +from datetime import datetime, timedelta, timezone + +import pytest +import pytest_asyncio +from _pytest.outcomes import fail + +import azure.cosmos.exceptions as exceptions +import test_config +from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy +from azure.cosmos.partition_key import PartitionKey + + +@pytest_asyncio.fixture() +async def setup(): + config = test_config.TestConfig() + if config.masterKey == '[YOUR_KEY_HERE]' or config.host == '[YOUR_ENDPOINT_HERE]': + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = CosmosClient(config.host, config.masterKey) + created_db = await test_client.create_database_if_not_exists(config.TEST_DATABASE_ID) + created_db_data = { + "created_db": created_db + } + + yield created_db_data + await test_client.delete_database(config.TEST_DATABASE_ID) + await test_client.close() + +@pytest.mark.cosmosEmulator +@pytest.mark.asyncio +@pytest.mark.usefixtures("setup") +class TestChangeFeedAsync: + """Test to ensure escaping of non-ascii characters from partition key""" + + created_db: DatabaseProxy = None + created_container: ContainerProxy = None + client: CosmosClient = None + config = test_config.TestConfig + TEST_CONTAINER_ID = config.TEST_MULTI_PARTITION_CONTAINER_ID + TEST_DATABASE_ID = config.TEST_DATABASE_ID + host = config.host + masterKey = config.masterKey + connectionPolicy = config.connectionPolicy + + async def test_get_feed_ranges(self, setup): + created_collection = await setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), + PartitionKey(path="/pk")) + result = await created_collection.read_feed_ranges() + assert len(result) == 1 + + @pytest.mark.parametrize("change_feed_filter_param", ["partitionKey", "partitionKeyRangeId", "feedRange"]) + async def test_query_change_feed_with_different_filter_async(self, change_feed_filter_param, setup): + + created_collection = await setup["created_db"].create_container( + "change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk")) + + if change_feed_filter_param == "partitionKey": + filter_param = {"partition_key": "pk"} + elif change_feed_filter_param == "partitionKeyRangeId": + filter_param = {"partition_key_range_id": "0"} + elif change_feed_filter_param == "feedRange": + feed_ranges = await created_collection.read_feed_ranges() + assert len(feed_ranges) == 1 + filter_param = {"feed_range": feed_ranges[0]} + else: + filter_param = None + + # Read change feed without passing any options + query_iterable = await created_collection.query_items_change_feed() + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + + # Read change feed from current should return an empty list + query_iterable = await created_collection.query_items_change_feed(filter_param) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + if 'Etag' in created_collection.client_connection.last_response_headers: + assert created_collection.client_connection.last_response_headers['Etag'] != '' + elif 'etag' in created_collection.client_connection.last_response_headers: + assert created_collection.client_connection.last_response_headers['etag'] != '' + else: + fail("No Etag or etag found in last response headers") + + # Read change feed from beginning should return an empty list + query_iterable = await created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + if 'Etag' in created_collection.client_connection.last_response_headers: + continuation1 = created_collection.client_connection.last_response_headers['Etag'] + elif 'etag' in created_collection.client_connection.last_response_headers: + continuation1 = created_collection.client_connection.last_response_headers['etag'] + else: + fail("No Etag or etag found in last response headers") + assert continuation1 != '' + + # Create a document. Read change feed should return be able to read that document + document_definition = {'pk': 'pk', 'id': 'doc1'} + await created_collection.create_item(body=document_definition) + query_iterable = await created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 1 + assert iter_list[0]['id'] == 'doc1' + if 'Etag' in created_collection.client_connection.last_response_headers: + continuation2 = created_collection.client_connection.last_response_headers['Etag'] + elif 'etag' in created_collection.client_connection.last_response_headers: + continuation2 = created_collection.client_connection.last_response_headers['etag'] + else: + fail("No Etag or etag found in last response headers") + assert continuation2 != '' + assert continuation2 != continuation1 + + # Create two new documents. Verify that change feed contains the 2 new documents + # with page size 1 and page size 100 + document_definition = {'pk': 'pk', 'id': 'doc2'} + await created_collection.create_item(body=document_definition) + document_definition = {'pk': 'pk', 'id': 'doc3'} + await created_collection.create_item(body=document_definition) + + for pageSize in [2, 100]: + # verify iterator + query_iterable = await created_collection.query_items_change_feed( + continuation=continuation2, + max_item_count=pageSize, + **filter_param) + it = query_iterable.__aiter__() + expected_ids = 'doc2.doc3.' + actual_ids = '' + async for item in it: + actual_ids += item['id'] + '.' + assert actual_ids == expected_ids + + # verify by_page + # the options is not copied, therefore it need to be restored + query_iterable = await created_collection.query_items_change_feed( + continuation=continuation2, + max_item_count=pageSize, + **filter_param + ) + count = 0 + expected_count = 2 + all_fetched_res = [] + pages = query_iterable.by_page() + async for items in await pages.__anext__(): + count += 1 + all_fetched_res.append(items) + assert count == expected_count + + actual_ids = '' + for item in all_fetched_res: + actual_ids += item['id'] + '.' + assert actual_ids == expected_ids + + # verify reading change feed from the beginning + query_iterable = await created_collection.query_items_change_feed( + is_start_from_beginning=True, + **filter_param + ) + expected_ids = ['doc1', 'doc2', 'doc3'] + it = query_iterable.__aiter__() + for i in range(0, len(expected_ids)): + doc = await it.__anext__() + assert doc['id'] == expected_ids[i] + if 'Etag' in created_collection.client_connection.last_response_headers: + continuation3 = created_collection.client_connection.last_response_headers['Etag'] + elif 'etag' in created_collection.client_connection.last_response_headers: + continuation3 = created_collection.client_connection.last_response_headers['etag'] + else: + fail("No Etag or etag found in last response headers") + + # verify reading empty change feed + query_iterable = await created_collection.query_items_change_feed( + continuation=continuation3, + is_start_from_beginning=True, + **filter_param + ) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + + await setup["created_db"].delete_container(created_collection.id) + + @pytest.mark.asyncio + async def test_query_change_feed_with_start_time(self, setup): + created_collection = await setup["created_db"].create_container_if_not_exists("query_change_feed_start_time_test", + PartitionKey(path="/pk")) + batchSize = 50 + + def round_time(): + utc_now = datetime.now(timezone.utc) + return utc_now - timedelta(microseconds=utc_now.microsecond) + + async def create_random_items(container, batch_size): + for _ in range(batch_size): + # Generate a Random partition key + partition_key = 'pk' + str(uuid.uuid4()) + + # Generate a random item + item = { + 'id': 'item' + str(uuid.uuid4()), + 'partitionKey': partition_key, + 'content': 'This is some random content', + } + + try: + # Create the item in the container + await container.upsert_item(item) + except exceptions.CosmosHttpResponseError as e: + pytest.fail(e) + + # Create first batch of random items + await create_random_items(created_collection, batchSize) + + # wait for 1 second and record the time, then wait another second + await sleep(1) + start_time = round_time() + not_utc_time = datetime.now() + await sleep(1) + + # now create another batch of items + await create_random_items(created_collection, batchSize) + + # now query change feed based on start time + change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=start_time)] + totalCount = len(change_feed_iter) + + # now check if the number of items that were changed match the batch size + assert totalCount == batchSize + + # negative test: pass in a valid time in the future + future_time = start_time + timedelta(hours=1) + change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=future_time)] + totalCount = len(change_feed_iter) + # A future time should return 0 + assert totalCount == 0 + + # test a date that is not utc, will be converted to utc by sdk + change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=not_utc_time)] + totalCount = len(change_feed_iter) + # Should equal batch size + assert totalCount == batchSize + + # test an invalid value, Attribute error will be raised for passing non datetime object + invalid_time = "Invalid value" + try: + change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=invalid_time)] + fail("Cannot format date on a non datetime object.") + except ValueError as e: + assert ("Invalid start_time 'Invalid value'" == e.args[0]) + + await setup["created_db"].delete_container(created_collection.id) + + async def test_query_change_feed_with_split_async(self, setup): + created_collection = await setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + + # initial change feed query returns empty result + query_iterable = await created_collection.query_items_change_feed(start_time="Beginning") + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + continuation = created_collection.client_connection.last_response_headers['etag'] + assert continuation != '' + + # create one doc and make sure change feed query can return the document + document_definition = {'pk': 'pk', 'id': 'doc1'} + await created_collection.create_item(body=document_definition) + query_iterable = await created_collection.query_items_change_feed(continuation=continuation) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 1 + continuation = created_collection.client_connection.last_response_headers['etag'] + + print("Triggering a split in test_query_change_feed_with_split") + await created_collection.replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = await created_collection.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in test_query_change_feed_with_split has completed") + print("creating few more documents") + new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc2', 'doc3', 'doc4'] + for document in new_documents: + await created_collection.create_item(body=document) + + query_iterable = await created_collection.query_items_change_feed(continuation=continuation) + it = query_iterable.__aiter__() + actual_ids = [] + async for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + setup["created_db"].delete_container(created_collection.id) + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_query.py b/sdk/cosmos/azure-cosmos/test/test_query.py index 73249e562c14..3cc8f57a6d21 100644 --- a/sdk/cosmos/azure-cosmos/test/test_query.py +++ b/sdk/cosmos/azure-cosmos/test/test_query.py @@ -4,8 +4,7 @@ import os import unittest import uuid -from datetime import datetime, timedelta, timezone -from time import sleep + import pytest import azure.cosmos._retry_utility as retry_utility @@ -61,293 +60,6 @@ def test_first_and_last_slashes_trimmed_for_query_string(self): self.assertEqual(iter_list[0]['id'], doc_id) self.created_db.delete_container(created_collection.id) - def test_query_change_feed_with_pk(self): - created_collection = self.created_db.create_container("change_feed_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk")) - # The test targets partition #3 - partition_key = "pk" - - # Read change feed without passing any options - query_iterable = created_collection.query_items_change_feed() - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - - # Read change feed from current should return an empty list - query_iterable = created_collection.query_items_change_feed(partition_key=partition_key) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - self.assertNotEqual(created_collection.client_connection.last_response_headers['etag'], '') - - # Read change feed from beginning should return an empty list - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation1 = created_collection.client_connection.last_response_headers['etag'] - self.assertNotEqual(continuation1, '') - - # Create a document. Read change feed should return be able to read that document - document_definition = {'pk': 'pk', 'id': 'doc1'} - created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 1) - self.assertEqual(iter_list[0]['id'], 'doc1') - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation2 = created_collection.client_connection.last_response_headers['etag'] - self.assertNotEqual(continuation2, '') - self.assertNotEqual(continuation2, continuation1) - - # Create two new documents. Verify that change feed contains the 2 new documents - # with page size 1 and page size 100 - document_definition = {'pk': 'pk', 'id': 'doc2'} - created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk', 'id': 'doc3'} - created_collection.create_item(body=document_definition) - - for pageSize in [1, 100]: - # verify iterator - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - partition_key=partition_key - ) - it = query_iterable.__iter__() - expected_ids = 'doc2.doc3.' - actual_ids = '' - for item in it: - actual_ids += item['id'] + '.' - self.assertEqual(actual_ids, expected_ids) - - # verify by_page - # the options is not copied, therefore it need to be restored - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - partition_key=partition_key - ) - count = 0 - expected_count = 2 - all_fetched_res = [] - for page in query_iterable.by_page(): - fetched_res = list(page) - self.assertEqual(len(fetched_res), min(pageSize, expected_count - count)) - count += len(fetched_res) - all_fetched_res.extend(fetched_res) - - actual_ids = '' - for item in all_fetched_res: - actual_ids += item['id'] + '.' - self.assertEqual(actual_ids, expected_ids) - - # verify reading change feed from the beginning - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - expected_ids = ['doc1', 'doc2', 'doc3'] - it = query_iterable.__iter__() - for i in range(0, len(expected_ids)): - doc = next(it) - self.assertEqual(doc['id'], expected_ids[i]) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation3 = created_collection.client_connection.last_response_headers['etag'] - - # verify reading empty change feed - query_iterable = created_collection.query_items_change_feed( - continuation=continuation3, - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.created_db.delete_container(created_collection.id) - - # TODO: partition key range id 0 is relative to the way collection is created - @pytest.mark.skip - def test_query_change_feed_with_pk_range_id(self): - created_collection = self.created_db.create_container("change_feed_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk")) - # The test targets partition #3 - partition_key_range_id = 0 - partitionParam = {"partition_key_range_id": partition_key_range_id} - - # Read change feed without passing any options - query_iterable = created_collection.query_items_change_feed() - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - - # Read change feed from current should return an empty list - query_iterable = created_collection.query_items_change_feed(**partitionParam) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - self.assertNotEqual(created_collection.client_connection.last_response_headers['etag'], '') - - # Read change feed from beginning should return an empty list - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partitionParam - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation1 = created_collection.client_connection.last_response_headers['etag'] - self.assertNotEqual(continuation1, '') - - # Create a document. Read change feed should return be able to read that document - document_definition = {'pk': 'pk', 'id': 'doc1'} - created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partitionParam - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 1) - self.assertEqual(iter_list[0]['id'], 'doc1') - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation2 = created_collection.client_connection.last_response_headers['etag'] - self.assertNotEqual(continuation2, '') - self.assertNotEqual(continuation2, continuation1) - - # Create two new documents. Verify that change feed contains the 2 new documents - # with page size 1 and page size 100 - document_definition = {'pk': 'pk', 'id': 'doc2'} - created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk', 'id': 'doc3'} - created_collection.create_item(body=document_definition) - - for pageSize in [1, 100]: - # verify iterator - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - **partitionParam - ) - it = query_iterable.__iter__() - expected_ids = 'doc2.doc3.' - actual_ids = '' - for item in it: - actual_ids += item['id'] + '.' - self.assertEqual(actual_ids, expected_ids) - - # verify by_page - # the options is not copied, therefore it need to be restored - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - **partitionParam - ) - count = 0 - expected_count = 2 - all_fetched_res = [] - for page in query_iterable.by_page(): - fetched_res = list(page) - self.assertEqual(len(fetched_res), min(pageSize, expected_count - count)) - count += len(fetched_res) - all_fetched_res.extend(fetched_res) - - actual_ids = '' - for item in all_fetched_res: - actual_ids += item['id'] + '.' - self.assertEqual(actual_ids, expected_ids) - - # verify reading change feed from the beginning - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partitionParam - ) - expected_ids = ['doc1', 'doc2', 'doc3'] - it = query_iterable.__iter__() - for i in range(0, len(expected_ids)): - doc = next(it) - self.assertEqual(doc['id'], expected_ids[i]) - self.assertTrue('etag' in created_collection.client_connection.last_response_headers) - continuation3 = created_collection.client_connection.last_response_headers['etag'] - - # verify reading empty change feed - query_iterable = created_collection.query_items_change_feed( - continuation=continuation3, - is_start_from_beginning=True, - **partitionParam - ) - iter_list = list(query_iterable) - self.assertEqual(len(iter_list), 0) - self.created_db.delete_container(created_collection.id) - - def test_query_change_feed_with_start_time(self): - created_collection = self.created_db.create_container_if_not_exists("query_change_feed_start_time_test", - PartitionKey(path="/pk")) - batchSize = 50 - - def round_time(): - utc_now = datetime.now(timezone.utc) - return utc_now - timedelta(microseconds=utc_now.microsecond) - def create_random_items(container, batch_size): - for _ in range(batch_size): - # Generate a Random partition key - partition_key = 'pk' + str(uuid.uuid4()) - - # Generate a random item - item = { - 'id': 'item' + str(uuid.uuid4()), - 'partitionKey': partition_key, - 'content': 'This is some random content', - } - - try: - # Create the item in the container - container.upsert_item(item) - except exceptions.CosmosHttpResponseError as e: - self.fail(e) - - # Create first batch of random items - create_random_items(created_collection, batchSize) - - # wait for 1 second and record the time, then wait another second - sleep(1) - start_time = round_time() - not_utc_time = datetime.now() - sleep(1) - - # now create another batch of items - create_random_items(created_collection, batchSize) - - # now query change feed based on start time - change_feed_iter = list(created_collection.query_items_change_feed(start_time=start_time)) - totalCount = len(change_feed_iter) - - # now check if the number of items that were changed match the batch size - self.assertEqual(totalCount, batchSize) - - # negative test: pass in a valid time in the future - future_time = start_time + timedelta(hours=1) - change_feed_iter = list(created_collection.query_items_change_feed(start_time=future_time)) - totalCount = len(change_feed_iter) - # A future time should return 0 - self.assertEqual(totalCount, 0) - - # test a date that is not utc, will be converted to utc by sdk - change_feed_iter = list(created_collection.query_items_change_feed(start_time=not_utc_time)) - totalCount = len(change_feed_iter) - # Should equal batch size - self.assertEqual(totalCount, batchSize) - - # test an invalid value, Attribute error will be raised for passing non datetime object - invalid_time = "Invalid value" - try: - change_feed_iter = list(created_collection.query_items_change_feed(start_time=invalid_time)) - self.fail("Cannot format date on a non datetime object.") - except AttributeError as e: - self.assertTrue("'str' object has no attribute 'astimezone'" == e.args[0]) - def test_populate_query_metrics(self): created_collection = self.created_db.create_container("query_metrics_test", PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/test/test_query_async.py b/sdk/cosmos/azure-cosmos/test/test_query_async.py index 51018126462d..718c544193a3 100644 --- a/sdk/cosmos/azure-cosmos/test/test_query_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_query_async.py @@ -4,8 +4,7 @@ import os import unittest import uuid -from asyncio import sleep, gather -from datetime import datetime, timedelta, timezone +from asyncio import gather import pytest @@ -14,10 +13,10 @@ import test_config from azure.cosmos import http_constants from azure.cosmos._execution_context.query_execution_info import _PartitionedQueryExecutionInfo +from azure.cosmos._retry_options import RetryOptions from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy from azure.cosmos.documents import _DistinctType from azure.cosmos.partition_key import PartitionKey -from azure.cosmos._retry_options import RetryOptions @pytest.mark.cosmosEmulator @@ -69,329 +68,6 @@ async def test_first_and_last_slashes_trimmed_for_query_string_async(self): await self.created_db.delete_container(created_collection.id) - async def test_query_change_feed_with_pk_async(self): - created_collection = await self.created_db.create_container( - "change_feed_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk")) - # The test targets partition #3 - partition_key = "pk" - - # Read change feed without passing any options - query_iterable = created_collection.query_items_change_feed() - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - - # Read change feed from current should return an empty list - query_iterable = created_collection.query_items_change_feed(partition_key=partition_key) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - if 'Etag' in created_collection.client_connection.last_response_headers: - assert created_collection.client_connection.last_response_headers['Etag'] != '' - elif 'etag' in created_collection.client_connection.last_response_headers: - assert created_collection.client_connection.last_response_headers['etag'] != '' - else: - self.fail("No Etag or etag found in last response headers") - - # Read change feed from beginning should return an empty list - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation1 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation1 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - assert continuation1 != '' - - # Create a document. Read change feed should return be able to read that document - document_definition = {'pk': 'pk', 'id': 'doc1'} - await created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 1 - assert iter_list[0]['id'] == 'doc1' - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation2 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation2 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - assert continuation2 != '' - assert continuation2 != continuation1 - - # Create two new documents. Verify that change feed contains the 2 new documents - # with page size 1 and page size 100 - document_definition = {'pk': 'pk', 'id': 'doc2'} - await created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk', 'id': 'doc3'} - await created_collection.create_item(body=document_definition) - - for pageSize in [2, 100]: - # verify iterator - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - partition_key=partition_key) - it = query_iterable.__aiter__() - expected_ids = 'doc2.doc3.' - actual_ids = '' - async for item in it: - actual_ids += item['id'] + '.' - assert actual_ids == expected_ids - - # verify by_page - # the options is not copied, therefore it need to be restored - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - partition_key=partition_key - ) - count = 0 - expected_count = 2 - all_fetched_res = [] - pages = query_iterable.by_page() - async for items in await pages.__anext__(): - count += 1 - all_fetched_res.append(items) - assert count == expected_count - - actual_ids = '' - for item in all_fetched_res: - actual_ids += item['id'] + '.' - assert actual_ids == expected_ids - - # verify reading change feed from the beginning - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - partition_key=partition_key - ) - expected_ids = ['doc1', 'doc2', 'doc3'] - it = query_iterable.__aiter__() - for i in range(0, len(expected_ids)): - doc = await it.__anext__() - assert doc['id'] == expected_ids[i] - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation3 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation3 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - - # verify reading empty change feed - query_iterable = created_collection.query_items_change_feed( - continuation=continuation3, - is_start_from_beginning=True, - partition_key=partition_key - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - - await self.created_db.delete_container(created_collection.id) - - # TODO: partition key range id 0 is relative to the way collection is created - @pytest.mark.skip - async def test_query_change_feed_with_pk_range_id_async(self): - created_collection = await self.created_db.create_container("cf_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk")) - # The test targets partition #3 - partition_key_range_id = 0 - partition_param = {"partition_key_range_id": partition_key_range_id} - - # Read change feed without passing any options - query_iterable = created_collection.query_items_change_feed() - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - - # Read change feed from current should return an empty list - query_iterable = created_collection.query_items_change_feed(**partition_param) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - if 'Etag' in created_collection.client_connection.last_response_headers: - assert created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - assert created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - - # Read change feed from beginning should return an empty list - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partition_param - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation1 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation1 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - assert continuation1 != '' - - # Create a document. Read change feed should return be able to read that document - document_definition = {'pk': 'pk', 'id': 'doc1'} - await created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partition_param - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 1 - assert iter_list[0]['id'] == 'doc1' - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation2 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation2 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - assert continuation2 != '' - assert continuation2 != continuation1 - - # Create two new documents. Verify that change feed contains the 2 new documents - # with page size 1 and page size 100 - document_definition = {'pk': 'pk', 'id': 'doc2'} - await created_collection.create_item(body=document_definition) - document_definition = {'pk': 'pk', 'id': 'doc3'} - await created_collection.create_item(body=document_definition) - - for pageSize in [2, 100]: - # verify iterator - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - **partition_param - ) - it = query_iterable.__aiter__() - expected_ids = 'doc2.doc3.' - actual_ids = '' - async for item in it: - actual_ids += item['id'] + '.' - assert actual_ids == expected_ids - - # verify by_page - # the options is not copied, therefore it need to be restored - query_iterable = created_collection.query_items_change_feed( - continuation=continuation2, - max_item_count=pageSize, - **partition_param - ) - count = 0 - expected_count = 2 - all_fetched_res = [] - pages = query_iterable.by_page() - async for items in await pages.__anext__(): - count += 1 - all_fetched_res.append(items) - assert count == expected_count - - actual_ids = '' - for item in all_fetched_res: - actual_ids += item['id'] + '.' - assert actual_ids == expected_ids - - # verify reading change feed from the beginning - query_iterable = created_collection.query_items_change_feed( - is_start_from_beginning=True, - **partition_param - ) - expected_ids = ['doc1', 'doc2', 'doc3'] - it = query_iterable.__aiter__() - for i in range(0, len(expected_ids)): - doc = await it.__anext__() - assert doc['id'] == expected_ids[i] - if 'Etag' in created_collection.client_connection.last_response_headers: - continuation3 = created_collection.client_connection.last_response_headers['Etag'] - elif 'etag' in created_collection.client_connection.last_response_headers: - continuation3 = created_collection.client_connection.last_response_headers['etag'] - else: - self.fail("No Etag or etag found in last response headers") - - # verify reading empty change feed - query_iterable = created_collection.query_items_change_feed( - continuation=continuation3, - is_start_from_beginning=True, - **partition_param - ) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - - @pytest.mark.asyncio - async def test_query_change_feed_with_start_time(self): - created_collection = await self.created_db.create_container_if_not_exists("query_change_feed_start_time_test", - PartitionKey(path="/pk")) - batchSize = 50 - - def round_time(): - utc_now = datetime.now(timezone.utc) - return utc_now - timedelta(microseconds=utc_now.microsecond) - - async def create_random_items(container, batch_size): - for _ in range(batch_size): - # Generate a Random partition key - partition_key = 'pk' + str(uuid.uuid4()) - - # Generate a random item - item = { - 'id': 'item' + str(uuid.uuid4()), - 'partitionKey': partition_key, - 'content': 'This is some random content', - } - - try: - # Create the item in the container - await container.upsert_item(item) - except exceptions.CosmosHttpResponseError as e: - pytest.fail(e) - - # Create first batch of random items - await create_random_items(created_collection, batchSize) - - # wait for 1 second and record the time, then wait another second - await sleep(1) - start_time = round_time() - not_utc_time = datetime.now() - await sleep(1) - - # now create another batch of items - await create_random_items(created_collection, batchSize) - - # now query change feed based on start time - change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=start_time)] - totalCount = len(change_feed_iter) - - # now check if the number of items that were changed match the batch size - assert totalCount == batchSize - - # negative test: pass in a valid time in the future - future_time = start_time + timedelta(hours=1) - change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=future_time)] - totalCount = len(change_feed_iter) - # A future time should return 0 - assert totalCount == 0 - - # test a date that is not utc, will be converted to utc by sdk - change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=not_utc_time)] - totalCount = len(change_feed_iter) - # Should equal batch size - assert totalCount == batchSize - - # test an invalid value, Attribute error will be raised for passing non datetime object - invalid_time = "Invalid value" - try: - change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=invalid_time)] - self.fail("Cannot format date on a non datetime object.") - except AttributeError as e: - assert ("'str' object has no attribute 'astimezone'" == e.args[0]) - - await self.created_db.delete_container(created_collection.id) - @pytest.mark.asyncio async def test_populate_query_metrics_async(self): created_collection = await self.created_db.create_container( From 7a1a1ebb7e06e0260c755cf5f4f305576489144e Mon Sep 17 00:00:00 2001 From: annie-mac Date: Sat, 17 Aug 2024 18:17:56 -0700 Subject: [PATCH 02/59] remove async keyword from changeFeed query in aio package --- .../_change_feed/aio/change_feed_iterable.py | 104 +++++++++++++----- .../_change_feed/aio/change_feed_state.py | 56 +++++++--- ...feed_range_composite_continuation_token.py | 42 ++++--- .../azure/cosmos/_change_feed/feed_range.py | 102 +++++++++++++++++ .../azure/cosmos/aio/_container.py | 71 ++---------- .../azure-cosmos/azure/cosmos/container.py | 1 - .../azure/cosmos/partition_key.py | 18 ++- .../test/test_change_feed_async.py | 30 ++--- 8 files changed, 281 insertions(+), 143 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 501f3a7e4150..16a431653e9c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,11 +21,13 @@ """Iterable change feed results in the Azure Cosmos database service. """ + from azure.core.async_paging import AsyncPageIterator +from azure.cosmos import PartitionKey from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.aio.change_feed_state import ChangeFeedStateV1, ChangeFeedState -from azure.cosmos._utils import is_base64_encoded +from azure.cosmos._utils import is_base64_encoded, is_key_exists_and_not_none class ChangeFeedIterable(AsyncPageIterator): @@ -57,40 +59,30 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link + self._change_feed_fetcher = None - change_feed_state = self._options.get("changeFeedState") - if not change_feed_state: - raise ValueError("Missing changeFeedState in feed options") + if not is_key_exists_and_not_none(self._options, "changeFeedStateContext"): + raise ValueError("Missing changeFeedStateContext in feed options") - if isinstance(change_feed_state, ChangeFeedStateV1): - if continuation_token: - if is_base64_encoded(continuation_token): - raise ValueError("Incompatible continuation token") - else: - change_feed_state.apply_server_response_continuation(continuation_token) + change_feed_state_context = self._options.pop("changeFeedStateContext") - self._change_feed_fetcher = ChangeFeedFetcherV1( - self._client, - self._collection_link, - self._options, - fetch_function - ) - else: - if continuation_token: - if not is_base64_encoded(continuation_token): - raise ValueError("Incompatible continuation token") + continuation = continuation_token if continuation_token is not None else change_feed_state_context.pop("continuation", None) - effective_change_feed_context = {"continuationFeedRange": continuation_token} - effective_change_feed_state = ChangeFeedState.from_json(change_feed_state.container_rid, effective_change_feed_context) - # replace with the effective change feed state - self._options["continuationFeedRange"] = effective_change_feed_state + # analysis and validate continuation token + # there are two types of continuation token we support currently: + # v1 version: the continuation token would just be the _etag, + # which is being returned when customer is using partition_key_range_id, + # which is under deprecation and does not support split/merge + # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + if continuation is not None: + if is_base64_encoded(continuation): + change_feed_state_context["continuationFeedRange"] = continuation + else: + change_feed_state_context["continuationPkRangeId"] = continuation + + self._validate_change_feed_state_context(change_feed_state_context) + self._options["changeFeedStateContext"] = change_feed_state_context - self._change_feed_fetcher = ChangeFeedFetcherV2( - self._client, - self._collection_link, - self._options, - fetch_function - ) super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) async def _unpack(self, block): @@ -112,7 +104,59 @@ async def _fetch_next(self, *args): # pylint: disable=unused-argument :return: List of results. :rtype: list """ + if self._change_feed_fetcher is None: + await self._initialize_change_feed_fetcher() + block = await self._change_feed_fetcher.fetch_next_block() if not block: raise StopAsyncIteration return block + + async def _initialize_change_feed_fetcher(self): + change_feed_state_context = self._options.pop("changeFeedStateContext") + conn_properties = await change_feed_state_context.pop("containerProperties") + if is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): + change_feed_state_context["partitionKey"] = await change_feed_state_context.pop("partitionKey") + + pk_properties = conn_properties.get("partitionKey") + partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) + + change_feed_state =\ + ChangeFeedState.from_json(self._collection_link, conn_properties["_rid"], partition_key_definition, change_feed_state_context) + self._options["changeFeedState"] = change_feed_state + + if isinstance(change_feed_state, ChangeFeedStateV1): + self._change_feed_fetcher = ChangeFeedFetcherV1( + self._client, + self._collection_link, + self._options, + self._fetch_function + ) + else: + self._change_feed_fetcher = ChangeFeedFetcherV2( + self._client, + self._collection_link, + self._options, + self._fetch_function + ) + + def _validate_change_feed_state_context(self, change_feed_state_context: dict[str, any]) -> None: + + if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + # if continuation token is in v1 format, throw exception if feed_range is set + if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + raise ValueError("feed_range and continuation are incompatible") + elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + # if continuation token is in v2 format, since the token itself contains the full change feed state + # so we will ignore other parameters (including incompatible parameters) if they passed in + pass + else: + # validation when no continuation is passed + exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] + count = sum(1 for key in exclusive_keys if + key in change_feed_state_context and change_feed_state_context[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") + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py index ae2e37568bd4..eede9bd4fe15 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py @@ -29,11 +29,12 @@ from abc import ABC, abstractmethod from typing import Optional, Union, List, Any -from azure.cosmos import http_constants +from azure.cosmos import http_constants, PartitionKey from azure.cosmos._change_feed.aio.change_feed_start_from import ChangeFeedStartFromETagAndFeedRange, \ ChangeFeedStartFromInternal from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken from azure.cosmos._change_feed.aio.feed_range_composite_continuation_token import FeedRangeCompositeContinuation +from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey, FeedRange from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range from azure.cosmos._utils import is_key_exists_and_not_none @@ -49,7 +50,10 @@ def populate_feed_options(self, feed_options: dict[str, any]) -> None: pass @abstractmethod - async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, request_headers: dict[str, any]) -> None: + async def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + request_headers: dict[str, any]) -> None: pass @abstractmethod @@ -57,7 +61,11 @@ def apply_server_response_continuation(self, continuation: str) -> None: pass @staticmethod - def from_json(container_link: str, container_rid: str, data: dict[str, Any]): + def from_json( + container_link: str, + container_rid: str, + partition_key_definition: PartitionKey, + data: dict[str, Any]): if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): return ChangeFeedStateV1.from_json(container_link, container_rid, data) else: @@ -69,11 +77,11 @@ def from_json(container_link: str, container_rid: str, data: dict[str, Any]): if version is None: raise ValueError("Invalid base64 encoded continuation string [Missing version]") elif version == "V2": - return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) + return ChangeFeedStateV2.from_continuation(container_link, container_rid, partition_key_definition, continuation_json) else: raise ValueError("Invalid base64 encoded continuation string [Invalid version]") # when there is no continuation token, by default construct ChangeFeedStateV2 - return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, partition_key_definition, data) class ChangeFeedStateV1(ChangeFeedState): """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag @@ -110,7 +118,10 @@ def from_json(cls, container_link: str, container_rid: str, data: dict[str, Any] data.get("continuationPkRangeId") ) - async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, Any]) -> None: + async def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + headers: dict[str, Any]) -> None: headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time @@ -140,7 +151,8 @@ def __init__( self, container_link: str, container_rid: str, - feed_range: Range, + partition_key_definition: PartitionKey, + feed_range: FeedRange, change_feed_start_from: ChangeFeedStartFromInternal, continuation: Optional[FeedRangeCompositeContinuation] = None): @@ -151,7 +163,9 @@ def __init__( self._continuation = continuation if self._continuation is None: composite_continuation_token_queue = collections.deque() - composite_continuation_token_queue.append(CompositeContinuationToken(self._feed_range, None)) + composite_continuation_token_queue.append(CompositeContinuationToken( + self._feed_range.get_normalized_range(partition_key_definition), + None)) self._continuation =\ FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) @@ -168,7 +182,10 @@ def to_dict(self) -> dict[str, Any]: self.continuation_property_name: self._continuation.to_dict() } - async def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, any]) -> None: + async def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + headers: dict[str, any]) -> None: headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time @@ -224,6 +241,7 @@ def from_continuation( cls, container_link: str, container_rid: str, + partition_key_definition: PartitionKey, continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) @@ -244,6 +262,7 @@ def from_continuation( return ChangeFeedStateV2( container_link=container_link, container_rid=container_rid, + partition_key_definition=partition_key_definition, feed_range=continuation.feed_range, change_feed_start_from=change_feed_start_from, continuation=continuation) @@ -253,26 +272,29 @@ def from_initial_state( cls, container_link: str, collection_rid: str, + partition_key_definition: PartitionKey, data: dict[str, Any]) -> 'ChangeFeedStateV2': if is_key_exists_and_not_none(data, "feedRange"): feed_range_str = base64.b64decode(data["feedRange"]).decode('utf-8') feed_range_json = json.loads(feed_range_str) - feed_range = Range.ParseFromDict(feed_range_json) - elif is_key_exists_and_not_none(data, "partitionKeyFeedRange"): - feed_range = data["partitionKeyFeedRange"] + feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) + elif is_key_exists_and_not_none(data, "partitionKey"): + feed_range = FeedRangePartitionKey(data["partitionKey"]) else: # default to full range - feed_range = Range( - "", - "FF", - True, - False) + feed_range = FeedRangeEpk( + Range( + "", + "FF", + True, + False)) change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(data.get("startTime")) return cls( container_link=container_link, container_rid=collection_rid, + partition_key_definition=partition_key_definition, feed_range=feed_range, change_feed_start_from=change_feed_start_from, continuation=None) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py index 6e1b8f974eea..d7bf97c0a903 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py @@ -27,20 +27,21 @@ from typing import Any from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range +from azure.cosmos._utils import is_key_exists_and_not_none class FeedRangeCompositeContinuation(object): - _version_property_name = "V" - _container_rid_property_name = "Rid" - _continuation_property_name = "Continuation" - _feed_range_property_name = "Range" + _version_property_name = "v" + _container_rid_property_name = "rid" + _continuation_property_name = "continuation" def __init__( self, container_rid: str, - feed_range: Range, + feed_range: FeedRange, continuation: collections.deque[CompositeContinuationToken]): if container_rid is None: raise ValueError("container_rid is missing") @@ -55,31 +56,34 @@ def __init__( def current_token(self): return self._current_token + def get_feed_range(self) -> FeedRange: + if isinstance(self._feed_range, FeedRangeEpk): + return FeedRangeEpk(self.current_token.feed_range) + else: + return self._feed_range + def to_dict(self) -> dict[str, Any]: - return { - self._version_property_name: "v1", #TODO: should this start from v2 + json_data = { + self._version_property_name: "v2", self._container_rid_property_name: self._container_rid, self._continuation_property_name: [childToken.to_dict() for childToken in self._continuation], - self._feed_range_property_name: self._feed_range.to_dict() } + json_data.update(self._feed_range.to_dict()) + return json_data + @classmethod def from_json(cls, data) -> 'FeedRangeCompositeContinuation': version = data.get(cls._version_property_name) if version is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._version_property_name}]") - if version != "v1": + if version != "v2": raise ValueError("Invalid feed range composite continuation token [Invalid version]") container_rid = data.get(cls._container_rid_property_name) if container_rid is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") - feed_range_data = data.get(cls._feed_range_property_name) - if feed_range_data is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._feed_range_property_name}]") - feed_range = Range.ParseFromDict(feed_range_data) - continuation_data = data.get(cls._continuation_property_name) if continuation_data is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") @@ -87,6 +91,14 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] + # parsing feed range + if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): + feed_range = FeedRangeEpk.from_json({ FeedRangeEpk.type_property_name: data[FeedRangeEpk.type_property_name] }) + elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): + feed_range = FeedRangePartitionKey.from_json({ FeedRangePartitionKey.type_property_name: data[FeedRangePartitionKey.type_property_name] }) + else: + raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]") + return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: @@ -130,5 +142,5 @@ def apply_not_modified_response(self) -> None: self._initial_no_result_range = self._current_token.feed_range @property - def feed_range(self) -> Range: + def feed_range(self) -> FeedRange: return self._feed_range diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py new file mode 100644 index 000000000000..a4b4b5dfedda --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -0,0 +1,102 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal class for feed range implementation in the Azure Cosmos +database service. +""" +import json +from abc import ABC, abstractmethod +from typing import Union, List + +from azure.cosmos import PartitionKey +from azure.cosmos._routing.routing_range import Range +from azure.cosmos._utils import is_key_exists_and_not_none +from azure.cosmos.partition_key import _Undefined, _Empty + + +class FeedRange(ABC): + + @abstractmethod + def get_normalized_range(self, partition_key_range_definition: PartitionKey) -> Range: + pass + + @abstractmethod + def to_dict(self) -> dict[str, any]: + pass + +class FeedRangePartitionKey(FeedRange): + type_property_name = "PK" + + def __init__( + self, + pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]): + if pk_value is None: + raise ValueError("PartitionKey cannot be None") + + self._pk_value = pk_value + + def get_normalized_range(self, partition_key_definition: PartitionKey) -> Range: + return partition_key_definition._get_epk_range_for_partition_key(self._pk_value).to_normalized_range() + + def to_dict(self) -> dict[str, any]: + if isinstance(self._pk_value, _Undefined): + return { self.type_property_name: [{}] } + elif isinstance(self._pk_value, _Empty): + return { self.type_property_name: [] } + else: + return { self.type_property_name: json.dumps(self._pk_value) } + + @classmethod + def from_json(cls, data: dict[str, any]) -> 'FeedRangePartitionKey': + if is_key_exists_and_not_none(data, cls.type_property_name): + pk_value = data.get(cls.type_property_name) + if isinstance(pk_value, list): + if not pk_value: + return cls(_Empty()) + if pk_value == [{}]: + return cls(_Undefined()) + + return cls(json.loads(data.get(cls.type_property_name))) + raise ValueError(f"Can not parse FeedRangePartitionKey from the json, there is no property {cls.type_property_name}") + +class FeedRangeEpk(FeedRange): + type_property_name = "Range" + + def __init__(self, feed_range: Range): + if feed_range is None: + raise ValueError("feed_range cannot be None") + + self._range = feed_range + + def get_normalized_range(self, partition_key_definition: PartitionKey) -> Range: + return self._range.to_normalized_range() + + def to_dict(self) -> dict[str, any]: + return { + self.type_property_name: self._range.to_dict() + } + + @classmethod + def from_json(cls, data: dict[str, any]) -> 'FeedRangeEpk': + if is_key_exists_and_not_none(data, cls.type_property_name): + feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) + return cls(feed_range) + raise ValueError(f"Can not parse FeedRangeEPK from the json, there is no property {cls.type_property_name}") \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index a8559839aad7..c68ddad7eb0d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -41,10 +41,9 @@ GenerateGuidId, _set_properties_cache ) -from .._change_feed.aio.change_feed_state import ChangeFeedState from .._routing import routing_range from .._routing.routing_range import Range -from .._utils import is_key_exists_and_not_none, is_base64_encoded +from .._utils import is_key_exists_and_not_none from ..offer import ThroughputProperties from ..partition_key import ( NonePartitionKeyValue, @@ -137,25 +136,12 @@ async def _set_partition_key( return _return_undefined_or_empty_partition_key(await self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) - async def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: - container_properties = await self._get_properties() - partition_key_definition = container_properties.get("partitionKey") - partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) - - is_prefix_partition_key = await self.__is_prefix_partition_key(partition_key_value) - - return partition_key._get_epk_range_for_partition_key(partition_key_value, is_prefix_partition_key) - async def __is_prefix_partition_key(self, partition_key: PartitionKeyType) -> bool: properties = await self._get_properties() pk_properties = properties.get("partitionKey") partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) - if partition_key_definition.kind != "MultiHash": - return False - if isinstance(partition_key, list) and len(partition_key_definition['paths']) == len(partition_key): - return False - return True + return partition_key_definition._is_prefix_partition_key(partition_key) @distributed_trace_async async def read( @@ -506,13 +492,12 @@ def query_items( return items @overload - async def query_items_change_feed( + def query_items_change_feed( self, *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, partition_key: Optional[PartitionKeyType] = None, - # -> would RU usage be more efficient, bug to backend team? deprecate it or using FeedRange to convert? priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: @@ -561,7 +546,7 @@ async def query_items_change_feed( ... @overload - async def query_items_change_feed( + def query_items_change_feed( self, *, continuation: Optional[str] = None, @@ -582,7 +567,7 @@ async def query_items_change_feed( ... @distributed_trace - async def query_items_change_feed( + def query_items_change_feed( self, *args: Any, **kwargs: Any @@ -637,16 +622,7 @@ async def query_items_change_feed( continuation = feed_options.pop('continuation') except KeyError: continuation = args[2] - - # there are two types of continuation token we support currently: - # v1 version: the continuation token would just be the _etag, - # which is being returned when customer is using partition_key_range_id, - # which is under deprecation and does not support split/merge - # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state - if is_base64_encoded(continuation): - change_feed_state_context["continuationFeedRange"] = continuation - else: - change_feed_state_context["continuationPkRangeId"] = continuation + change_feed_state_context["continuation"] = continuation if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): try: @@ -655,42 +631,21 @@ async def query_items_change_feed( feed_options["maxItemCount"] = args[3] if is_key_exists_and_not_none(kwargs, "partition_key"): - partition_key = kwargs.pop("partition_key") - change_feed_state_context["partitionKey"] = await self._set_partition_key(partition_key) - change_feed_state_context["partitionKeyFeedRange"] = await self._get_epk_range_for_partition_key(partition_key) + change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop("partition_key")) if is_key_exists_and_not_none(kwargs, "feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') - # validate exclusive or in-compatible parameters - if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): - # if continuation token is in v1 format, throw exception if feed_range is set - if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): - raise ValueError("feed_range and continuation are incompatible") - elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): - # if continuation token is in v2 format, since the token itself contains the full change feed state - # so we will ignore other parameters if they passed in - if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId"): - raise ValueError("partition_key_range_id and continuation are incompatible") - else: - # validation when no continuation is passed - exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] - count = sum(1 for key in exclusive_keys if - key in change_feed_state_context and change_feed_state_context[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") - - container_properties = await self._get_properties() - container_rid = container_properties.get("_rid") - change_feed_state = ChangeFeedState.from_json(self.container_link, container_rid, change_feed_state_context) - feed_options["changeFeedState"] = change_feed_state - feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + change_feed_state_context["containerProperties"] = self._get_properties() + feed_options["changeFeedStateContext"] = change_feed_state_context response_hook = kwargs.pop('response_hook', None) if 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"] + result = self.client_connection.QueryItemsChangeFeed( self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) @@ -1269,5 +1224,3 @@ async def read_feed_ranges( [Range("", "FF", True, False)]) return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] - - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 32fd818075f8..d4f1d5480241 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -327,7 +327,6 @@ def query_items_change_feed( max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, partition_key: Optional[PartitionKeyType] = None, - # -> would RU usage be more efficient, bug to backend team? deprecate it or using FeedRange to convert? priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 9f0a5cde29a2..c89e1d9ac771 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -175,17 +175,14 @@ def _get_epk_range_for_prefix_partition_key( def _get_epk_range_for_partition_key( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]], - is_prefix_pk_value: bool = False + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> _Range: - if is_prefix_pk_value: + if self._is_prefix_partition_key(pk_value): return self._get_epk_range_for_prefix_partition_key(pk_value) # else return point range effective_partition_key_string = self._get_effective_partition_key_string(pk_value) - partition_key_range = _Range(effective_partition_key_string, effective_partition_key_string, True, True) - - return partition_key_range.to_normalized_range() + return _Range(effective_partition_key_string, effective_partition_key_string, True, True) def _get_effective_partition_key_for_hash_partitioning(self) -> str: # We shouldn't be supporting V1 @@ -279,6 +276,15 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( return ''.join(sb).upper() + def _is_prefix_partition_key( + self, + partition_key: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> bool: + if self.kind!= "MultiHash": + return False + if isinstance(partition_key, list) and len(self.path) == len(partition_key): + return False + return True + def _return_undefined_or_empty_partition_key(is_system_key: bool) -> Union[_Empty, _Undefined]: if is_system_key: diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index c3246768a796..b4165af0601c 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -76,12 +76,12 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f filter_param = None # Read change feed without passing any options - query_iterable = await created_collection.query_items_change_feed() + query_iterable = created_collection.query_items_change_feed() iter_list = [item async for item in query_iterable] assert len(iter_list) == 0 # Read change feed from current should return an empty list - query_iterable = await created_collection.query_items_change_feed(filter_param) + query_iterable = created_collection.query_items_change_feed(filter_param) iter_list = [item async for item in query_iterable] assert len(iter_list) == 0 if 'Etag' in created_collection.client_connection.last_response_headers: @@ -92,7 +92,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f fail("No Etag or etag found in last response headers") # Read change feed from beginning should return an empty list - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( is_start_from_beginning=True, **filter_param ) @@ -109,7 +109,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f # Create a document. Read change feed should return be able to read that document document_definition = {'pk': 'pk', 'id': 'doc1'} await created_collection.create_item(body=document_definition) - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( is_start_from_beginning=True, **filter_param ) @@ -134,7 +134,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f for pageSize in [2, 100]: # verify iterator - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( continuation=continuation2, max_item_count=pageSize, **filter_param) @@ -147,7 +147,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f # verify by_page # the options is not copied, therefore it need to be restored - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( continuation=continuation2, max_item_count=pageSize, **filter_param @@ -167,7 +167,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f assert actual_ids == expected_ids # verify reading change feed from the beginning - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( is_start_from_beginning=True, **filter_param ) @@ -184,7 +184,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f fail("No Etag or etag found in last response headers") # verify reading empty change feed - query_iterable = await created_collection.query_items_change_feed( + query_iterable = created_collection.query_items_change_feed( continuation=continuation3, is_start_from_beginning=True, **filter_param @@ -235,7 +235,7 @@ async def create_random_items(container, batch_size): await create_random_items(created_collection, batchSize) # now query change feed based on start time - change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=start_time)] + change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=start_time)] totalCount = len(change_feed_iter) # now check if the number of items that were changed match the batch size @@ -243,13 +243,13 @@ async def create_random_items(container, batch_size): # negative test: pass in a valid time in the future future_time = start_time + timedelta(hours=1) - change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=future_time)] + change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=future_time)] totalCount = len(change_feed_iter) # A future time should return 0 assert totalCount == 0 # test a date that is not utc, will be converted to utc by sdk - change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=not_utc_time)] + change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=not_utc_time)] totalCount = len(change_feed_iter) # Should equal batch size assert totalCount == batchSize @@ -257,7 +257,7 @@ async def create_random_items(container, batch_size): # test an invalid value, Attribute error will be raised for passing non datetime object invalid_time = "Invalid value" try: - change_feed_iter = [i async for i in await created_collection.query_items_change_feed(start_time=invalid_time)] + change_feed_iter = [i async for i in created_collection.query_items_change_feed(start_time=invalid_time)] fail("Cannot format date on a non datetime object.") except ValueError as e: assert ("Invalid start_time 'Invalid value'" == e.args[0]) @@ -270,7 +270,7 @@ async def test_query_change_feed_with_split_async(self, setup): offer_throughput=400) # initial change feed query returns empty result - query_iterable = await created_collection.query_items_change_feed(start_time="Beginning") + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") iter_list = [item async for item in query_iterable] assert len(iter_list) == 0 continuation = created_collection.client_connection.last_response_headers['etag'] @@ -279,7 +279,7 @@ async def test_query_change_feed_with_split_async(self, setup): # create one doc and make sure change feed query can return the document document_definition = {'pk': 'pk', 'id': 'doc1'} await created_collection.create_item(body=document_definition) - query_iterable = await created_collection.query_items_change_feed(continuation=continuation) + query_iterable = created_collection.query_items_change_feed(continuation=continuation) iter_list = [item async for item in query_iterable] assert len(iter_list) == 1 continuation = created_collection.client_connection.last_response_headers['etag'] @@ -309,7 +309,7 @@ async def test_query_change_feed_with_split_async(self, setup): for document in new_documents: await created_collection.create_item(body=document) - query_iterable = await created_collection.query_items_change_feed(continuation=continuation) + query_iterable = created_collection.query_items_change_feed(continuation=continuation) it = query_iterable.__aiter__() actual_ids = [] async for item in it: From b6c53fb71465a01512ecf01c22c247cae97e55a9 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Sat, 17 Aug 2024 22:50:20 -0700 Subject: [PATCH 03/59] refactor --- .../_change_feed/aio/change_feed_iterable.py | 13 ++- .../_change_feed/aio/change_feed_state.py | 26 +++-- ...feed_range_composite_continuation_token.py | 13 +-- .../_change_feed/change_feed_fetcher.py | 1 - .../_change_feed/change_feed_iterable.py | 100 ++++++++++++------ .../cosmos/_change_feed/change_feed_state.py | 46 +++++--- .../azure/cosmos/_change_feed/feed_range.py | 24 +++-- ...feed_range_composite_continuation_token.py | 38 ++++--- .../azure/cosmos/aio/_container.py | 2 +- .../azure-cosmos/azure/cosmos/container.py | 55 ++-------- 10 files changed, 170 insertions(+), 148 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 16a431653e9c..c792d1357550 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -114,15 +114,16 @@ async def _fetch_next(self, *args): # pylint: disable=unused-argument async def _initialize_change_feed_fetcher(self): change_feed_state_context = self._options.pop("changeFeedStateContext") - conn_properties = await change_feed_state_context.pop("containerProperties") + conn_properties = await self._options.pop("containerProperties") if is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): change_feed_state_context["partitionKey"] = await change_feed_state_context.pop("partitionKey") - - pk_properties = conn_properties.get("partitionKey") - partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) + pk_properties = conn_properties.get("partitionKey") + partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) + change_feed_state_context["partitionKeyFeedRange"] =\ + partition_key_definition._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) change_feed_state =\ - ChangeFeedState.from_json(self._collection_link, conn_properties["_rid"], partition_key_definition, change_feed_state_context) + ChangeFeedState.from_json(self._collection_link, conn_properties["_rid"], change_feed_state_context) self._options["changeFeedState"] = change_feed_state if isinstance(change_feed_state, ChangeFeedStateV1): @@ -158,5 +159,3 @@ def _validate_change_feed_state_context(self, change_feed_state_context: dict[st if count > 1: raise ValueError( "partition_key_range_id, partition_key, feed_range are exclusive parameters, please only set one of them") - - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py index eede9bd4fe15..ceb83166bdab 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py @@ -29,7 +29,7 @@ from abc import ABC, abstractmethod from typing import Optional, Union, List, Any -from azure.cosmos import http_constants, PartitionKey +from azure.cosmos import http_constants from azure.cosmos._change_feed.aio.change_feed_start_from import ChangeFeedStartFromETagAndFeedRange, \ ChangeFeedStartFromInternal from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken @@ -64,7 +64,6 @@ def apply_server_response_continuation(self, continuation: str) -> None: def from_json( container_link: str, container_rid: str, - partition_key_definition: PartitionKey, data: dict[str, Any]): if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): return ChangeFeedStateV1.from_json(container_link, container_rid, data) @@ -77,11 +76,11 @@ def from_json( if version is None: raise ValueError("Invalid base64 encoded continuation string [Missing version]") elif version == "V2": - return ChangeFeedStateV2.from_continuation(container_link, container_rid, partition_key_definition, continuation_json) + return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) else: raise ValueError("Invalid base64 encoded continuation string [Invalid version]") # when there is no continuation token, by default construct ChangeFeedStateV2 - return ChangeFeedStateV2.from_initial_state(container_link, container_rid, partition_key_definition, data) + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) class ChangeFeedStateV1(ChangeFeedState): """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag @@ -151,7 +150,6 @@ def __init__( self, container_link: str, container_rid: str, - partition_key_definition: PartitionKey, feed_range: FeedRange, change_feed_start_from: ChangeFeedStartFromInternal, continuation: Optional[FeedRangeCompositeContinuation] = None): @@ -163,9 +161,10 @@ def __init__( self._continuation = continuation if self._continuation is None: composite_continuation_token_queue = collections.deque() - composite_continuation_token_queue.append(CompositeContinuationToken( - self._feed_range.get_normalized_range(partition_key_definition), - None)) + composite_continuation_token_queue.append( + CompositeContinuationToken( + self._feed_range.get_normalized_range(), + None)) self._continuation =\ FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) @@ -241,7 +240,6 @@ def from_continuation( cls, container_link: str, container_rid: str, - partition_key_definition: PartitionKey, continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) @@ -262,8 +260,7 @@ def from_continuation( return ChangeFeedStateV2( container_link=container_link, container_rid=container_rid, - partition_key_definition=partition_key_definition, - feed_range=continuation.feed_range, + feed_range=continuation._feed_range, change_feed_start_from=change_feed_start_from, continuation=continuation) @@ -272,7 +269,6 @@ def from_initial_state( cls, container_link: str, collection_rid: str, - partition_key_definition: PartitionKey, data: dict[str, Any]) -> 'ChangeFeedStateV2': if is_key_exists_and_not_none(data, "feedRange"): @@ -280,7 +276,10 @@ def from_initial_state( feed_range_json = json.loads(feed_range_str) feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) elif is_key_exists_and_not_none(data, "partitionKey"): - feed_range = FeedRangePartitionKey(data["partitionKey"]) + if is_key_exists_and_not_none(data, "partitionKeyFeedRange"): + feed_range = FeedRangePartitionKey(data["partitionKey"], data["partitionKeyFeedRange"]) + else: + raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") else: # default to full range feed_range = FeedRangeEpk( @@ -294,7 +293,6 @@ def from_initial_state( return cls( container_link=container_link, container_rid=collection_rid, - partition_key_definition=partition_key_definition, feed_range=feed_range, change_feed_start_from=change_feed_start_from, continuation=None) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py index d7bf97c0a903..32122145009c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py @@ -56,12 +56,6 @@ def __init__( def current_token(self): return self._current_token - def get_feed_range(self) -> FeedRange: - if isinstance(self._feed_range, FeedRangeEpk): - return FeedRangeEpk(self.current_token.feed_range) - else: - return self._feed_range - def to_dict(self) -> dict[str, Any]: json_data = { self._version_property_name: "v2", @@ -93,16 +87,17 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': # parsing feed range if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): - feed_range = FeedRangeEpk.from_json({ FeedRangeEpk.type_property_name: data[FeedRangeEpk.type_property_name] }) + feed_range = FeedRangeEpk.from_json(data) elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): - feed_range = FeedRangePartitionKey.from_json({ FeedRangePartitionKey.type_property_name: data[FeedRangePartitionKey.type_property_name] }) + feed_range =\ + FeedRangePartitionKey.from_json(data, continuation[0].feed_range) else: raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]") return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: - overlapping_ranges = await routing_provider.get_overlapping_ranges(collection_link, self._current_token.feed_range) + overlapping_ranges = await routing_provider.get_overlapping_ranges(collection_link, [self._current_token.feed_range]) if len(overlapping_ranges) == 1: # merge,reusing the existing the feedRange and continuationToken diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index fd8ac2787a8f..6edaf8e73fd9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -115,7 +115,6 @@ def __init__( self._change_feed_state = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV2): raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") - self._change_feed_state.__class__ = ChangeFeedStateV2 self._resource_link = resource_link self._fetch_function = fetch_function diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 676036180d29..39e94a30c4c0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -26,7 +26,7 @@ from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedState -from azure.cosmos._utils import is_base64_encoded +from azure.cosmos._utils import is_base64_encoded, is_key_exists_and_not_none class ChangeFeedIterable(PageIterator): @@ -57,40 +57,29 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link + self._change_feed_fetcher = None - change_feed_state = self._options.get("changeFeedState") - if not change_feed_state: - raise ValueError("Missing changeFeedState in feed options") + if not is_key_exists_and_not_none(self._options, "changeFeedStateContext"): + raise ValueError("Missing changeFeedStateContext in feed options") - if isinstance(change_feed_state, ChangeFeedStateV1): - if continuation_token: - if is_base64_encoded(continuation_token): - raise ValueError("Incompatible continuation token") - else: - change_feed_state.apply_server_response_continuation(continuation_token) + change_feed_state_context = self._options.pop("changeFeedStateContext") + continuation = continuation_token if continuation_token is not None else change_feed_state_context.pop("continuation", None) - self._change_feed_fetcher = ChangeFeedFetcherV1( - self._client, - self._collection_link, - self._options, - fetch_function - ) - else: - if continuation_token: - if not is_base64_encoded(continuation_token): - raise ValueError("Incompatible continuation token") + # analysis and validate continuation token + # there are two types of continuation token we support currently: + # v1 version: the continuation token would just be the _etag, + # which is being returned when customer is using partition_key_range_id, + # which is under deprecation and does not support split/merge + # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + if continuation is not None: + if is_base64_encoded(continuation): + change_feed_state_context["continuationFeedRange"] = continuation + else: + change_feed_state_context["continuationPkRangeId"] = continuation - effective_change_feed_context = {"continuationFeedRange": continuation_token} - effective_change_feed_state = ChangeFeedState.from_json(change_feed_state.container_rid, effective_change_feed_context) - # replace with the effective change feed state - self._options["continuationFeedRange"] = effective_change_feed_state + self._validate_change_feed_state_context(change_feed_state_context) + self._options["changeFeedStateContext"] = change_feed_state_context - self._change_feed_fetcher = ChangeFeedFetcherV2( - self._client, - self._collection_link, - self._options, - fetch_function - ) super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) def _unpack(self, block): @@ -112,7 +101,58 @@ def _fetch_next(self, *args): # pylint: disable=unused-argument :return: List of results. :rtype: list """ + + if self._change_feed_fetcher is None: + self._initialize_change_feed_fetcher() + block = self._change_feed_fetcher.fetch_next_block() if not block: raise StopIteration return block + + def _initialize_change_feed_fetcher(self): + change_feed_state_context = self._options.pop("changeFeedStateContext") + change_feed_state = \ + ChangeFeedState.from_json( + self._collection_link, + self._options.get("containerRID"), + change_feed_state_context) + + self._options["changeFeedState"] = change_feed_state + + if isinstance(change_feed_state, ChangeFeedStateV1): + self._change_feed_fetcher = ChangeFeedFetcherV1( + self._client, + self._collection_link, + self._options, + self._fetch_function + ) + else: + self._change_feed_fetcher = ChangeFeedFetcherV2( + self._client, + self._collection_link, + self._options, + self._fetch_function + ) + + def _validate_change_feed_state_context(self, change_feed_state_context: dict[str, any]) -> None: + + if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + # if continuation token is in v1 format, throw exception if feed_range is set + if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + raise ValueError("feed_range and continuation are incompatible") + elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + # if continuation token is in v2 format, since the token itself contains the full change feed state + # so we will ignore other parameters (including incompatible parameters) if they passed in + pass + else: + # validation when no continuation is passed + exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] + count = sum(1 for key in exclusive_keys if + key in change_feed_state_context and change_feed_state_context[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") + + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 8c61c306b94e..210563ab8411 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -33,6 +33,7 @@ from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromInternal, \ ChangeFeedStartFromETagAndFeedRange from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._change_feed.feed_range_composite_continuation_token import FeedRangeCompositeContinuation from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range @@ -57,13 +58,18 @@ def apply_server_response_continuation(self, continuation: str) -> None: pass @staticmethod - def from_json(container_link: str, container_rid: str, data: dict[str, Any]): - if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): - return ChangeFeedStateV1.from_json(container_link, container_rid, data) + def from_json( + container_link: str, + container_rid: str, + change_feed_state_context: dict[str, Any]): + + if (is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId") + or is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId")): + return ChangeFeedStateV1.from_json(container_link, container_rid, change_feed_state_context) else: - if is_key_exists_and_not_none(data, "continuationFeedRange"): + if is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): # get changeFeedState from continuation - continuation_json_str = base64.b64decode(data["continuationFeedRange"]).decode('utf-8') + continuation_json_str = base64.b64decode(change_feed_state_context["continuationFeedRange"]).decode('utf-8') continuation_json = json.loads(continuation_json_str) version = continuation_json.get(ChangeFeedState.version_property_name) if version is None: @@ -73,7 +79,7 @@ def from_json(container_link: str, container_rid: str, data: dict[str, Any]): else: raise ValueError("Invalid base64 encoded continuation string [Invalid version]") # when there is no continuation token, by default construct ChangeFeedStateV2 - return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, change_feed_state_context) class ChangeFeedStateV1(ChangeFeedState): """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag @@ -140,7 +146,7 @@ def __init__( self, container_link: str, container_rid: str, - feed_range: Range, + feed_range: FeedRange, change_feed_start_from: ChangeFeedStartFromInternal, continuation: Optional[FeedRangeCompositeContinuation] = None): @@ -151,7 +157,10 @@ def __init__( self._continuation = continuation if self._continuation is None: composite_continuation_token_queue = collections.deque() - composite_continuation_token_queue.append(CompositeContinuationToken(self._feed_range, None)) + composite_continuation_token_queue.append( + CompositeContinuationToken( + self._feed_range.get_normalized_range(), + None)) self._continuation =\ FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) @@ -253,23 +262,28 @@ def from_initial_state( cls, container_link: str, collection_rid: str, - data: dict[str, Any]) -> 'ChangeFeedStateV2': + change_feed_state_context: dict[str, Any]) -> 'ChangeFeedStateV2': - if is_key_exists_and_not_none(data, "feedRange"): - feed_range_str = base64.b64decode(data["feedRange"]).decode('utf-8') + if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') feed_range_json = json.loads(feed_range_str) - feed_range = Range.ParseFromDict(feed_range_json) - elif is_key_exists_and_not_none(data, "partitionKeyFeedRange"): - feed_range = data["partitionKeyFeedRange"] + feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) + elif is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): + if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyFeedRange"): + feed_range = FeedRangePartitionKey(change_feed_state_context["partitionKey"], change_feed_state_context["partitionKeyFeedRange"]) + else: + raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") else: # default to full range - feed_range = Range( + feed_range = FeedRangeEpk( + Range( "", "FF", True, False) + ) - change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(data.get("startTime")) + change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(change_feed_state_context.get("startTime")) return cls( container_link=container_link, container_rid=collection_rid, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index a4b4b5dfedda..1b6ef79e4176 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -26,7 +26,6 @@ from abc import ABC, abstractmethod from typing import Union, List -from azure.cosmos import PartitionKey from azure.cosmos._routing.routing_range import Range from azure.cosmos._utils import is_key_exists_and_not_none from azure.cosmos.partition_key import _Undefined, _Empty @@ -35,7 +34,7 @@ class FeedRange(ABC): @abstractmethod - def get_normalized_range(self, partition_key_range_definition: PartitionKey) -> Range: + def get_normalized_range(self) -> Range: pass @abstractmethod @@ -47,14 +46,19 @@ class FeedRangePartitionKey(FeedRange): def __init__( self, - pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]): + pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], + feed_range: Range): + if pk_value is None: raise ValueError("PartitionKey cannot be None") + if feed_range is None: + raise ValueError("Feed range cannot be None") self._pk_value = pk_value + self._feed_range = feed_range - def get_normalized_range(self, partition_key_definition: PartitionKey) -> Range: - return partition_key_definition._get_epk_range_for_partition_key(self._pk_value).to_normalized_range() + def get_normalized_range(self) -> Range: + return self._feed_range.to_normalized_range() def to_dict(self) -> dict[str, any]: if isinstance(self._pk_value, _Undefined): @@ -65,16 +69,16 @@ def to_dict(self) -> dict[str, any]: return { self.type_property_name: json.dumps(self._pk_value) } @classmethod - def from_json(cls, data: dict[str, any]) -> 'FeedRangePartitionKey': + def from_json(cls, data: dict[str, any], feed_range: Range) -> 'FeedRangePartitionKey': if is_key_exists_and_not_none(data, cls.type_property_name): pk_value = data.get(cls.type_property_name) if isinstance(pk_value, list): if not pk_value: - return cls(_Empty()) + return cls(_Empty(), feed_range) if pk_value == [{}]: - return cls(_Undefined()) + return cls(_Undefined(), feed_range) - return cls(json.loads(data.get(cls.type_property_name))) + return cls(json.loads(data.get(cls.type_property_name)), feed_range) raise ValueError(f"Can not parse FeedRangePartitionKey from the json, there is no property {cls.type_property_name}") class FeedRangeEpk(FeedRange): @@ -86,7 +90,7 @@ def __init__(self, feed_range: Range): self._range = feed_range - def get_normalized_range(self, partition_key_definition: PartitionKey) -> Range: + def get_normalized_range(self) -> Range: return self._range.to_normalized_range() def to_dict(self) -> dict[str, any]: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index 2461436924aa..26abd66ba132 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -27,20 +27,21 @@ from typing import Any from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken +from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range +from azure.cosmos._utils import is_key_exists_and_not_none class FeedRangeCompositeContinuation(object): - _version_property_name = "V" - _container_rid_property_name = "Rid" - _continuation_property_name = "Continuation" - _feed_range_property_name = "Range" + _version_property_name = "v" + _container_rid_property_name = "rid" + _continuation_property_name = "continuation" def __init__( self, container_rid: str, - feed_range: Range, + feed_range: FeedRange, continuation: collections.deque[CompositeContinuationToken]): if container_rid is None: raise ValueError("container_rid is missing") @@ -56,30 +57,26 @@ def current_token(self): return self._current_token def to_dict(self) -> dict[str, Any]: - return { - self._version_property_name: "v1", #TODO: should this start from v2 + json_data = { + self._version_property_name: "v2", self._container_rid_property_name: self._container_rid, self._continuation_property_name: [childToken.to_dict() for childToken in self._continuation], - self._feed_range_property_name: self._feed_range.to_dict() } + json_data.update(self._feed_range.to_dict()) + return json_data @classmethod def from_json(cls, data) -> 'FeedRangeCompositeContinuation': version = data.get(cls._version_property_name) if version is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._version_property_name}]") - if version != "v1": + if version != "v2": raise ValueError("Invalid feed range composite continuation token [Invalid version]") container_rid = data.get(cls._container_rid_property_name) if container_rid is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") - feed_range_data = data.get(cls._feed_range_property_name) - if feed_range_data is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._feed_range_property_name}]") - feed_range = Range.ParseFromDict(feed_range_data) - continuation_data = data.get(cls._continuation_property_name) if continuation_data is None: raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") @@ -87,10 +84,18 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] + # parsing feed range + if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): + feed_range = FeedRangeEpk.from_json(data) + elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): + feed_range = FeedRangePartitionKey.from_json(data, continuation[0].feed_range) + else: + raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]") + return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: - overlapping_ranges = routing_provider.get_overlapping_ranges(collection_link, self._current_token.feed_range) + overlapping_ranges = routing_provider.get_overlapping_ranges(collection_link, [self._current_token.feed_range]) if len(overlapping_ranges) == 1: # merge,reusing the existing the feedRange and continuationToken @@ -130,5 +135,6 @@ def apply_not_modified_response(self) -> None: self._initial_no_result_range = self._current_token.feed_range @property - def feed_range(self) -> Range: + def feed_range(self) -> FeedRange: return self._feed_range + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index c68ddad7eb0d..5e8da4f54cf4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -636,7 +636,7 @@ def query_items_change_feed( if is_key_exists_and_not_none(kwargs, "feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') - change_feed_state_context["containerProperties"] = self._get_properties() + feed_options["containerProperties"] = self._get_properties() feed_options["changeFeedStateContext"] = change_feed_state_context response_hook = kwargs.pop('response_hook', None) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index d4f1d5480241..e2fa8ad4071e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -38,11 +38,10 @@ GenerateGuidId, _set_properties_cache ) -from ._change_feed.change_feed_state import ChangeFeedState from ._cosmos_client_connection import CosmosClientConnection from ._routing import routing_range from ._routing.routing_range import Range -from ._utils import is_key_exists_and_not_none, is_base64_encoded +from ._utils import is_key_exists_and_not_none from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -133,15 +132,15 @@ def _set_partition_key( 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) - def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: - return self.client_connection._container_properties_cache - def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: container_properties = self._get_properties() partition_key_definition = container_properties.get("partitionKey") partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) - return partition_key._get_epk_range_for_partition_key(partition_key_value, self.__is_prefix_partitionkey(partition_key_value)) + return partition_key._get_epk_range_for_partition_key(partition_key_value) + + def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: + return self.client_connection._container_properties_cache @distributed_trace def read( # pylint:disable=docstring-missing-param @@ -451,16 +450,7 @@ def query_items_change_feed( continuation = feed_options.pop('continuation') except KeyError: continuation = args[2] - - # there are two types of continuation token we support currently: - # v1 version: the continuation token would just be the _etag, - # which is being returned when customer is using partition_key_range_id, - # which is under deprecation and does not support split/merge - # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state - if is_base64_encoded(continuation): - change_feed_state_context["continuationFeedRange"] = continuation - else: - change_feed_state_context["continuationPkRangeId"] = continuation + change_feed_state_context["continuation"] = continuation if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): try: @@ -469,40 +459,21 @@ def query_items_change_feed( feed_options["maxItemCount"] = args[3] if is_key_exists_and_not_none(kwargs, "partition_key"): - partition_key = kwargs.pop("partition_key") + partition_key = kwargs.pop('partition_key') change_feed_state_context["partitionKey"] = self._set_partition_key(partition_key) change_feed_state_context["partitionKeyFeedRange"] = self._get_epk_range_for_partition_key(partition_key) if is_key_exists_and_not_none(kwargs, "feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') - # validate exclusive or in-compatible parameters - if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): - # if continuation token is in v1 format, throw exception if feed_range is set - if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): - raise ValueError("feed_range and continuation are incompatible") - elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): - # if continuation token is in v2 format, since the token itself contains the full change feed state - # so we will ignore other parameters if they passed in - if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId"): - raise ValueError("partition_key_range_id and continuation are incompatible") - else: - # validation when no continuation is passed - exclusive_keys = ["partitionKeyRangeId", "partitionKey", "feedRange"] - count = sum(1 for key in exclusive_keys if key in change_feed_state_context and change_feed_state_context[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") - container_properties = self._get_properties() - container_rid = container_properties.get("_rid") - change_feed_state = ChangeFeedState.from_json(self.container_link, container_rid, change_feed_state_context) - feed_options["changeFeedState"] = change_feed_state + feed_options["changeFeedStateContext"] = change_feed_state_context + feed_options["containerRID"] = container_properties["_rid"] response_hook = kwargs.pop('response_hook', None) if 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"] + result = self.client_connection.QueryItemsChangeFeed( self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) @@ -639,11 +610,7 @@ def __is_prefix_partitionkey( properties = self._get_properties() pk_properties = properties["partitionKey"] partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) - if partition_key_definition.kind != "MultiHash": - return False - if isinstance(partition_key, list) and len(partition_key_definition['paths']) == len(partition_key): - return False - return True + return partition_key_definition._is_prefix_partition_key(partition_key) @distributed_trace def replace_item( # pylint:disable=docstring-missing-param From 5f16b143fd341b8ef82e5ea49f515c846f32f52e Mon Sep 17 00:00:00 2001 From: annie-mac Date: Sun, 18 Aug 2024 11:02:49 -0700 Subject: [PATCH 04/59] refactor --- .../_change_feed/aio/change_feed_fetcher.py | 62 ++-- .../_change_feed/aio/change_feed_iterable.py | 15 +- .../aio/change_feed_start_from.py | 189 ----------- .../_change_feed/aio/change_feed_state.py | 299 ------------------ .../aio/composite_continuation_token.py | 70 ---- ...feed_range_composite_continuation_token.py | 141 --------- .../_change_feed/change_feed_fetcher.py | 66 ++-- .../_change_feed/change_feed_iterable.py | 14 +- .../_change_feed/change_feed_start_from.py | 46 +-- .../cosmos/_change_feed/change_feed_state.py | 238 ++++++++++---- .../azure/cosmos/_change_feed/feed_range.py | 38 ++- ...feed_range_composite_continuation_token.py | 61 +++- .../azure/cosmos/_cosmos_client_connection.py | 6 +- .../azure/cosmos/_routing/routing_range.py | 12 +- .../azure-cosmos/azure/cosmos/_utils.py | 6 +- .../azure/cosmos/aio/_container.py | 46 ++- .../aio/_cosmos_client_connection_async.py | 6 +- .../azure-cosmos/azure/cosmos/container.py | 54 ++-- .../azure-cosmos/azure/cosmos/exceptions.py | 2 +- .../azure/cosmos/partition_key.py | 2 +- 20 files changed, 436 insertions(+), 937 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index 83ca3025ee07..4d85a891ac3f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -26,12 +26,15 @@ import copy import json from abc import ABC, abstractmethod +from typing import Dict, Any, List from azure.cosmos import http_constants, exceptions -from azure.cosmos._change_feed.aio.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 +from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromPointInTime +from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 from azure.cosmos.aio import _retry_utility_async from azure.cosmos.exceptions import CosmosHttpResponseError +# pylint: disable=protected-access class ChangeFeedFetcher(ABC): @@ -49,7 +52,7 @@ def __init__( self, client, resource_link: str, - feed_options: dict[str, any], + feed_options: Dict[str, Any], fetch_function): self._client = client @@ -57,8 +60,8 @@ def __init__( self._change_feed_state = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV1): - raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version {type(self._change_feed_state)}") - self._change_feed_state.__class__ = ChangeFeedStateV1 + raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version" + f" {type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function @@ -74,24 +77,27 @@ async def callback(): return await _retry_utility_async.ExecuteAsync(self._client, self._client._global_endpoint_manager, callback) - async def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: new_options = copy.deepcopy(self._feed_options) new_options["changeFeedState"] = self._change_feed_state self._change_feed_state.populate_feed_options(new_options) - is_s_time_first_fetch = True + is_s_time_first_fetch = self._change_feed_state._continuation is None while True: (fetched_items, response_headers) = await fetch_function(new_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. - # For start time however we get no initial results, so we need to pass continuation token? Is this true? self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) if fetched_items: break - elif is_s_time_first_fetch: + + # When processing from point in time, there will be no initial results being returned, + # so we will retry with the new continuation token again + if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + and is_s_time_first_fetch): is_s_time_first_fetch = False else: break @@ -106,16 +112,15 @@ def __init__( self, client, resource_link: str, - feed_options: dict[str, any], + feed_options: Dict[str, Any], fetch_function): self._client = client self._feed_options = feed_options - self._change_feed_state = self._feed_options.pop("changeFeedState") + self._change_feed_state: ChangeFeedStateV2 = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV2): raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") - self._change_feed_state.__class__ = ChangeFeedStateV2 self._resource_link = resource_link self._fetch_function = fetch_function @@ -131,17 +136,22 @@ async def callback(): return await self.fetch_change_feed_items(self._fetch_function) try: - return await _retry_utility_async.ExecuteAsync(self._client, self._client._global_endpoint_manager, callback) + return await _retry_utility_async.ExecuteAsync( + self._client, + self._client._global_endpoint_manager, + callback) except CosmosHttpResponseError as e: if exceptions._partition_range_is_gone(e) or exceptions._is_partition_split_or_merge(e): # refresh change feed state - await self._change_feed_state.handle_feed_range_gone(self._client._routing_map_provider, self._resource_link) + await self._change_feed_state.handle_feed_range_gone_async( + self._client._routing_map_provider, + self._resource_link) else: raise e return await self.fetch_next_block() - async def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: new_options = copy.deepcopy(self._feed_options) new_options["changeFeedState"] = self._change_feed_state @@ -154,19 +164,33 @@ async def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. - # For start time however we get no initial results, so we need to pass continuation token? Is this true? if fetched_items: self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) response_headers[continuation_key] = self._get_base64_encoded_continuation() break - else: + + # when there is no items being returned, we will decide to retry based on: + # 1. When processing from point in time, there will be no initial results being returned, + # so we will retry with the new continuation token + # 2. if the feed range of the changeFeedState span multiple physical partitions + # then we will read from the next feed range until we have looped through all physical partitions self._change_feed_state.apply_not_modified_response() self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) - response_headers[continuation_key] = self._get_base64_encoded_continuation() - should_retry = self._change_feed_state.should_retry_on_not_modified_response() or is_s_time_first_fetch - is_s_time_first_fetch = False + + #TODO: can this part logic be simplified + if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + and is_s_time_first_fetch): + response_headers[continuation_key] = self._get_base64_encoded_continuation() + is_s_time_first_fetch = False + should_retry = True + else: + self._change_feed_state._continuation._move_to_next_token() + response_headers[continuation_key] = self._get_base64_encoded_continuation() + should_retry = self._change_feed_state.should_retry_on_not_modified_response() + is_s_time_first_fetch = False + if not should_retry: break diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index c792d1357550..83c12f59157c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,14 +21,16 @@ """Iterable change feed results in the Azure Cosmos database service. """ +from typing import Dict, Any from azure.core.async_paging import AsyncPageIterator from azure.cosmos import PartitionKey from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 -from azure.cosmos._change_feed.aio.change_feed_state import ChangeFeedStateV1, ChangeFeedState +from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateV1 from azure.cosmos._utils import is_base64_encoded, is_key_exists_and_not_none +# pylint: disable=protected-access class ChangeFeedIterable(AsyncPageIterator): """Represents an iterable object of the change feed results. @@ -66,14 +68,16 @@ def __init__( change_feed_state_context = self._options.pop("changeFeedStateContext") - continuation = continuation_token if continuation_token is not None else change_feed_state_context.pop("continuation", None) + continuation = continuation_token if continuation_token is not None\ + else change_feed_state_context.pop("continuation", None) # analysis and validate continuation token # there are two types of continuation token we support currently: # v1 version: the continuation token would just be the _etag, # which is being returned when customer is using partition_key_range_id, # which is under deprecation and does not support split/merge - # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + # v2 version: the continuation token will be base64 encoded composition token + # which includes full change feed state if continuation is not None: if is_base64_encoded(continuation): change_feed_state_context["continuationFeedRange"] = continuation @@ -141,7 +145,7 @@ async def _initialize_change_feed_fetcher(self): self._fetch_function ) - def _validate_change_feed_state_context(self, change_feed_state_context: dict[str, any]) -> None: + def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): # if continuation token is in v1 format, throw exception if feed_range is set @@ -158,4 +162,5 @@ def _validate_change_feed_state_context(self, change_feed_state_context: dict[st key in change_feed_state_context and change_feed_state_context[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") + "partition_key_range_id, partition_key, feed_range are exclusive parameters," + " please only set one of them") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py deleted file mode 100644 index 99aeeb6eb914..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_start_from.py +++ /dev/null @@ -1,189 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Internal class for change feed start from implementation in the Azure Cosmos database service. -""" - -from abc import ABC, abstractmethod -from datetime import datetime, timezone -from enum import Enum -from typing import Optional, Union, Literal, Any - -from azure.cosmos import http_constants -from azure.cosmos._routing.routing_range import Range - -class ChangeFeedStartFromType(Enum): - BEGINNING = "Beginning" - NOW = "Now" - LEASE = "Lease" - POINT_IN_TIME = "PointInTime" - -class ChangeFeedStartFromInternal(ABC): - """Abstract class for change feed start from implementation in the Azure Cosmos database service. - """ - - _type_property_name = "Type" - - @abstractmethod - def to_dict(self) -> dict[str, Any]: - pass - - @staticmethod - def from_start_time(start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]]) -> 'ChangeFeedStartFromInternal': - if start_time is None: - return ChangeFeedStartFromNow() - elif isinstance(start_time, datetime): - return ChangeFeedStartFromPointInTime(start_time) - elif start_time.lower() == ChangeFeedStartFromType.NOW.value.lower(): - return ChangeFeedStartFromNow() - elif start_time.lower() == ChangeFeedStartFromType.BEGINNING.value.lower(): - return ChangeFeedStartFromBeginning() - else: - raise ValueError(f"Invalid start_time '{start_time}'") - - @staticmethod - def from_json(data: dict[str, any]) -> 'ChangeFeedStartFromInternal': - change_feed_start_from_type = data.get(ChangeFeedStartFromInternal._type_property_name) - if change_feed_start_from_type is None: - raise ValueError(f"Invalid start from json [Missing {ChangeFeedStartFromInternal._type_property_name}]") - - if change_feed_start_from_type == ChangeFeedStartFromType.BEGINNING.value: - return ChangeFeedStartFromBeginning.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.LEASE.value: - return ChangeFeedStartFromETagAndFeedRange.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.NOW.value: - return ChangeFeedStartFromNow.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.POINT_IN_TIME.value: - return ChangeFeedStartFromPointInTime.from_json(data) - else: - raise ValueError(f"Can not process changeFeedStartFrom for type {change_feed_start_from_type}") - - @abstractmethod - def populate_request_headers(self, request_headers) -> None: - pass - - -class ChangeFeedStartFromBeginning(ChangeFeedStartFromInternal): - """Class for change feed start from beginning implementation in the Azure Cosmos database service. - """ - - def to_dict(self) -> dict[str, Any]: - return { - self._type_property_name: ChangeFeedStartFromType.BEGINNING.value - } - - def populate_request_headers(self, request_headers) -> None: - pass # there is no headers need to be set for start from beginning - - @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromBeginning': - return ChangeFeedStartFromBeginning() - - -class ChangeFeedStartFromETagAndFeedRange(ChangeFeedStartFromInternal): - """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. - """ - - _etag_property_name = "Etag" - _feed_range_property_name = "FeedRange" - - def __init__(self, etag, feed_range): - if feed_range is None: - raise ValueError("feed_range is missing") - - self._etag = etag - self._feed_range = feed_range - - def to_dict(self) -> dict[str, Any]: - return { - self._type_property_name: ChangeFeedStartFromType.LEASE.value, - self._etag_property_name: self._etag, - self._feed_range_property_name: self._feed_range.to_dict() - } - - @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromETagAndFeedRange': - etag = data.get(cls._etag_property_name) - if etag is None: - raise ValueError(f"Invalid change feed start from [Missing {cls._etag_property_name}]") - - feed_range_data = data.get(cls._feed_range_property_name) - if feed_range_data is None: - raise ValueError(f"Invalid change feed start from [Missing {cls._feed_range_property_name}]") - feed_range = Range.ParseFromDict(feed_range_data) - return cls(etag, feed_range) - - def populate_request_headers(self, request_headers) -> None: - # change feed uses etag as the continuationToken - if self._etag: - request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._etag - - -class ChangeFeedStartFromNow(ChangeFeedStartFromInternal): - """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. - """ - - def to_dict(self) -> dict[str, Any]: - return { - self._type_property_name: ChangeFeedStartFromType.NOW.value - } - - def populate_request_headers(self, request_headers) -> None: - request_headers[http_constants.HttpHeaders.IfNoneMatch] = "*" - - @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromNow': - return ChangeFeedStartFromNow() - - -class ChangeFeedStartFromPointInTime(ChangeFeedStartFromInternal): - """Class for change feed start from point in time implementation in the Azure Cosmos database service. - """ - - _point_in_time_ms_property_name = "PointInTimeMs" - - def __init__(self, start_time: datetime): - if start_time is None: - raise ValueError("start_time is missing") - - self._start_time = start_time - - def to_dict(self) -> dict[str, Any]: - return { - self._type_property_name: ChangeFeedStartFromType.POINT_IN_TIME.value, - self._point_in_time_ms_property_name: - int(self._start_time.astimezone(timezone.utc).timestamp() * 1000) - } - - def populate_request_headers(self, request_headers) -> None: - request_headers[http_constants.HttpHeaders.IfModified_since] =\ - self._start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') - - @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': - point_in_time_ms = data.get(cls._point_in_time_ms_property_name) - if point_in_time_ms is None: - raise ValueError(f"Invalid change feed start from {cls._point_in_time_ms_property_name} ") - - point_in_time = datetime.fromtimestamp(point_in_time_ms).astimezone(timezone.utc) - return ChangeFeedStartFromPointInTime(point_in_time) - - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py deleted file mode 100644 index ceb83166bdab..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_state.py +++ /dev/null @@ -1,299 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Internal class for change feed state implementation in the Azure Cosmos -database service. -""" - -import base64 -import collections -import json -from abc import ABC, abstractmethod -from typing import Optional, Union, List, Any - -from azure.cosmos import http_constants -from azure.cosmos._change_feed.aio.change_feed_start_from import ChangeFeedStartFromETagAndFeedRange, \ - ChangeFeedStartFromInternal -from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken -from azure.cosmos._change_feed.aio.feed_range_composite_continuation_token import FeedRangeCompositeContinuation -from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey, FeedRange -from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider -from azure.cosmos._routing.routing_range import Range -from azure.cosmos._utils import is_key_exists_and_not_none -from azure.cosmos.exceptions import CosmosFeedRangeGoneError -from azure.cosmos.partition_key import _Empty, _Undefined - - -class ChangeFeedState(ABC): - version_property_name = "v" - - @abstractmethod - def populate_feed_options(self, feed_options: dict[str, any]) -> None: - pass - - @abstractmethod - async def populate_request_headers( - self, - routing_provider: SmartRoutingMapProvider, - request_headers: dict[str, any]) -> None: - pass - - @abstractmethod - def apply_server_response_continuation(self, continuation: str) -> None: - pass - - @staticmethod - def from_json( - container_link: str, - container_rid: str, - data: dict[str, Any]): - if is_key_exists_and_not_none(data, "partitionKeyRangeId") or is_key_exists_and_not_none(data, "continuationPkRangeId"): - return ChangeFeedStateV1.from_json(container_link, container_rid, data) - else: - if is_key_exists_and_not_none(data, "continuationFeedRange"): - # get changeFeedState from continuation - continuation_json_str = base64.b64decode(data["continuationFeedRange"]).decode('utf-8') - continuation_json = json.loads(continuation_json_str) - version = continuation_json.get(ChangeFeedState.version_property_name) - if version is None: - raise ValueError("Invalid base64 encoded continuation string [Missing version]") - elif version == "V2": - return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) - else: - raise ValueError("Invalid base64 encoded continuation string [Invalid version]") - # when there is no continuation token, by default construct ChangeFeedStateV2 - return ChangeFeedStateV2.from_initial_state(container_link, container_rid, data) - -class ChangeFeedStateV1(ChangeFeedState): - """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag - """ - - def __init__( - self, - container_link: str, - container_rid: str, - change_feed_start_from: ChangeFeedStartFromInternal, - partition_key_range_id: Optional[str] = None, - partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, - continuation: Optional[str] = None): - - self._container_link = container_link - self._container_rid = container_rid - self._change_feed_start_from = change_feed_start_from - self._partition_key_range_id = partition_key_range_id - self._partition_key = partition_key - self._continuation = continuation - - @property - def container_rid(self): - return self._container_rid - - @classmethod - def from_json(cls, container_link: str, container_rid: str, data: dict[str, Any]) -> 'ChangeFeedStateV1': - return cls( - container_link, - container_rid, - ChangeFeedStartFromInternal.from_start_time(data.get("startTime")), - data.get("partitionKeyRangeId"), - data.get("partitionKey"), - data.get("continuationPkRangeId") - ) - - async def populate_request_headers( - self, - routing_provider: SmartRoutingMapProvider, - headers: dict[str, Any]) -> None: - headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - - # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. - self._change_feed_start_from.populate_request_headers(headers) - if self._continuation: - headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation - - def populate_feed_options(self, feed_options: dict[str, any]) -> None: - if self._partition_key_range_id is not None: - feed_options["partitionKeyRangeId"] = self._partition_key_range_id - if self._partition_key is not None: - feed_options["partitionKey"] = self._partition_key - - def apply_server_response_continuation(self, continuation: str) -> None: - self._continuation = continuation - -class ChangeFeedStateV2(ChangeFeedState): - container_rid_property_name = "containerRid" - change_feed_mode_property_name = "mode" - change_feed_start_from_property_name = "startFrom" - continuation_property_name = "continuation" - - # TODO: adding change feed mode - def __init__( - self, - container_link: str, - container_rid: str, - feed_range: FeedRange, - change_feed_start_from: ChangeFeedStartFromInternal, - continuation: Optional[FeedRangeCompositeContinuation] = None): - - self._container_link = container_link - self._container_rid = container_rid - self._feed_range = feed_range - self._change_feed_start_from = change_feed_start_from - self._continuation = continuation - if self._continuation is None: - composite_continuation_token_queue = collections.deque() - composite_continuation_token_queue.append( - CompositeContinuationToken( - self._feed_range.get_normalized_range(), - None)) - self._continuation =\ - FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) - - @property - def container_rid(self) -> str : - return self._container_rid - - def to_dict(self) -> dict[str, Any]: - return { - self.version_property_name: "V2", - self.container_rid_property_name: self._container_rid, - self.change_feed_mode_property_name: "Incremental", - self.change_feed_start_from_property_name: self._change_feed_start_from.to_dict(), - self.continuation_property_name: self._continuation.to_dict() - } - - async def populate_request_headers( - self, - routing_provider: SmartRoutingMapProvider, - headers: dict[str, any]) -> None: - headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - - # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. - self._change_feed_start_from.populate_request_headers(headers) - - if self._continuation.current_token is not None and self._continuation.current_token.token is not None: - change_feed_start_from_feed_range_and_etag =\ - ChangeFeedStartFromETagAndFeedRange(self._continuation.current_token.token, self._continuation.current_token.feed_range) - change_feed_start_from_feed_range_and_etag.populate_request_headers(headers) - - # based on the feed range to find the overlapping partition key range id - over_lapping_ranges =\ - await routing_provider.get_overlapping_ranges( - self._container_link, - [self._continuation.current_token.feed_range]) - - if len(over_lapping_ranges) > 1: - raise CosmosFeedRangeGoneError(message= - f"Range {self._continuation.current_token.feed_range}" - f" spans {len(over_lapping_ranges)}" - f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") - else: - overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) - if overlapping_feed_range == self._continuation.current_token.feed_range: - # exactly mapping to one physical partition, only need to set the partitionKeyRangeId - headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] - else: - # the current token feed range spans less than single physical partition - # for this case, need to set both the partition key range id and epk filter headers - headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] - headers[http_constants.HttpHeaders.StartEpkString] = self._continuation.current_token.feed_range.min - headers[http_constants.HttpHeaders.EndEpkString] = self._continuation.current_token.feed_range.max - - def populate_feed_options(self, feed_options: dict[str, any]) -> None: - pass - - async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, resource_link: str) -> None: - await self._continuation.handle_feed_range_gone(routing_provider, resource_link) - - def apply_server_response_continuation(self, continuation: str) -> None: - self._continuation.apply_server_response_continuation(continuation) - - def should_retry_on_not_modified_response(self): - self._continuation.should_retry_on_not_modified_response() - - def apply_not_modified_response(self) -> None: - self._continuation.apply_not_modified_response() - - @classmethod - def from_continuation( - cls, - container_link: str, - container_rid: str, - continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': - - container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) - if container_rid_from_continuation is None: - raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.container_rid_property_name}]") - elif container_rid_from_continuation != container_rid: - raise ValueError("Invalid continuation: [Mismatch collection rid]") - - change_feed_start_from_data = continuation_json.get(ChangeFeedStateV2.change_feed_start_from_property_name) - if change_feed_start_from_data is None: - raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.change_feed_start_from_property_name}]") - change_feed_start_from = ChangeFeedStartFromInternal.from_json(change_feed_start_from_data) - - continuation_data = continuation_json.get(ChangeFeedStateV2.continuation_property_name) - if continuation_data is None: - raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.continuation_property_name}]") - continuation = FeedRangeCompositeContinuation.from_json(continuation_data) - return ChangeFeedStateV2( - container_link=container_link, - container_rid=container_rid, - feed_range=continuation._feed_range, - change_feed_start_from=change_feed_start_from, - continuation=continuation) - - @classmethod - def from_initial_state( - cls, - container_link: str, - collection_rid: str, - data: dict[str, Any]) -> 'ChangeFeedStateV2': - - if is_key_exists_and_not_none(data, "feedRange"): - feed_range_str = base64.b64decode(data["feedRange"]).decode('utf-8') - feed_range_json = json.loads(feed_range_str) - feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) - elif is_key_exists_and_not_none(data, "partitionKey"): - if is_key_exists_and_not_none(data, "partitionKeyFeedRange"): - feed_range = FeedRangePartitionKey(data["partitionKey"], data["partitionKeyFeedRange"]) - else: - raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") - else: - # default to full range - feed_range = FeedRangeEpk( - Range( - "", - "FF", - True, - False)) - - change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(data.get("startTime")) - return cls( - container_link=container_link, - container_rid=collection_rid, - feed_range=feed_range, - change_feed_start_from=change_feed_start_from, - continuation=None) - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py deleted file mode 100644 index 6d779fed1037..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/composite_continuation_token.py +++ /dev/null @@ -1,70 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Internal class for change feed composite continuation token in the Azure Cosmos -database service. -""" -from azure.cosmos._routing.routing_range import Range - - -class CompositeContinuationToken(object): - _token_property_name = "token" - _feed_range_property_name = "range" - - def __init__(self, feed_range: Range, token): - if range is None: - raise ValueError("range is missing") - - self._token = token - self._feed_range = feed_range - - def to_dict(self): - return { - self._token_property_name: self._token, - self._feed_range_property_name: self._feed_range.to_dict() - } - - @property - def feed_range(self): - return self._feed_range - - @property - def token(self): - return self._token - - def update_token(self, etag): - self._token = etag - - @classmethod - def from_json(cls, data): - token = data.get(cls._token_property_name) - if token is None: - raise ValueError(f"Invalid composite token [Missing {cls._token_property_name}]") - - feed_range_data = data.get(cls._feed_range_property_name) - if feed_range_data is None: - raise ValueError(f"Invalid composite token [Missing {cls._feed_range_property_name}]") - - feed_range = Range.ParseFromDict(feed_range_data) - return cls(feed_range=feed_range, token=token) - - def __repr__(self): - return f"CompositeContinuationToken(token={self.token}, range={self._feed_range})" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py deleted file mode 100644 index 32122145009c..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/feed_range_composite_continuation_token.py +++ /dev/null @@ -1,141 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Internal class for change feed continuation token by feed range in the Azure Cosmos -database service. -""" -import collections -from collections import deque -from typing import Any - -from azure.cosmos._change_feed.aio.composite_continuation_token import CompositeContinuationToken -from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey -from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider -from azure.cosmos._routing.routing_range import Range -from azure.cosmos._utils import is_key_exists_and_not_none - - -class FeedRangeCompositeContinuation(object): - _version_property_name = "v" - _container_rid_property_name = "rid" - _continuation_property_name = "continuation" - - def __init__( - self, - container_rid: str, - feed_range: FeedRange, - continuation: collections.deque[CompositeContinuationToken]): - if container_rid is None: - raise ValueError("container_rid is missing") - - self._container_rid = container_rid - self._feed_range = feed_range - self._continuation = continuation - self._current_token = self._continuation[0] - self._initial_no_result_range = None - - @property - def current_token(self): - return self._current_token - - def to_dict(self) -> dict[str, Any]: - json_data = { - self._version_property_name: "v2", - self._container_rid_property_name: self._container_rid, - self._continuation_property_name: [childToken.to_dict() for childToken in self._continuation], - } - - json_data.update(self._feed_range.to_dict()) - return json_data - - @classmethod - def from_json(cls, data) -> 'FeedRangeCompositeContinuation': - version = data.get(cls._version_property_name) - if version is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._version_property_name}]") - if version != "v2": - raise ValueError("Invalid feed range composite continuation token [Invalid version]") - - container_rid = data.get(cls._container_rid_property_name) - if container_rid is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") - - continuation_data = data.get(cls._continuation_property_name) - if continuation_data is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") - if not isinstance(continuation_data, list) or len(continuation_data) == 0: - raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") - continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] - - # parsing feed range - if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): - feed_range = FeedRangeEpk.from_json(data) - elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): - feed_range =\ - FeedRangePartitionKey.from_json(data, continuation[0].feed_range) - else: - raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]") - - return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) - - async def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: - overlapping_ranges = await routing_provider.get_overlapping_ranges(collection_link, [self._current_token.feed_range]) - - if len(overlapping_ranges) == 1: - # merge,reusing the existing the feedRange and continuationToken - pass - else: - # split, remove the parent range and then add new child ranges. - # For each new child range, using the continuation token from the parent - self._continuation.popleft() - for child_range in overlapping_ranges: - self._continuation.append(CompositeContinuationToken(Range.PartitionKeyRangeToRange(child_range), self._current_token.token)) - - self._current_token = self._continuation[0] - - def should_retry_on_not_modified_response(self) -> bool: - # when getting 304(Not Modified) response from one sub feed range, we will try to fetch for the next sub feed range - # we will repeat the above logic until we have looped through all sub feed ranges - - # TODO: validate the response headers, can we get the status code - if len(self._continuation) > 1: - return self._current_token.feed_range != self._initial_no_result_range - - else: - return False - - def _move_to_next_token(self) -> None: - first_composition_token = self._continuation.popleft() - # add the composition token to the end of the list - self._continuation.append(first_composition_token) - self._current_token = self._continuation[0] - - def apply_server_response_continuation(self, etag) -> None: - self._current_token.update_token(etag) - self._move_to_next_token() - - def apply_not_modified_response(self) -> None: - if self._initial_no_result_range is None: - self._initial_no_result_range = self._current_token.feed_range - - @property - def feed_range(self) -> FeedRange: - return self._feed_range diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index 6edaf8e73fd9..abad86dfd119 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -26,11 +26,14 @@ import copy import json from abc import ABC, abstractmethod +from typing import Dict, Any, List from azure.cosmos import _retry_utility, http_constants, exceptions +from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromPointInTime from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 from azure.cosmos.exceptions import CosmosHttpResponseError +# pylint: disable=protected-access class ChangeFeedFetcher(ABC): @@ -48,16 +51,16 @@ def __init__( self, client, resource_link: str, - feed_options: dict[str, any], + feed_options: Dict[str, Any], fetch_function): self._client = client self._feed_options = feed_options - self._change_feed_state = self._feed_options.pop("changeFeedState") + self._change_feed_state: ChangeFeedStateV1 = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV1): - raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version {type(self._change_feed_state)}") - self._change_feed_state.__class__ = ChangeFeedStateV1 + raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version" + f" {type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function @@ -73,28 +76,30 @@ def callback(): return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback) - def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: new_options = copy.deepcopy(self._feed_options) new_options["changeFeedState"] = self._change_feed_state self._change_feed_state.populate_feed_options(new_options) - is_s_time_first_fetch = True + is_s_time_first_fetch = self._change_feed_state._continuation is None while True: (fetched_items, response_headers) = fetch_function(new_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. - # For start time however we get no initial results, so we need to pass continuation token? Is this true? self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) if fetched_items: break - elif is_s_time_first_fetch: + + # When processing from point in time, there will be no initial results being returned, + # so we will retry with the new continuation token again + if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + and is_s_time_first_fetch): is_s_time_first_fetch = False else: break - return fetched_items @@ -106,15 +111,16 @@ def __init__( self, client, resource_link: str, - feed_options: dict[str, any], + feed_options: Dict[str, Any], fetch_function): self._client = client self._feed_options = feed_options - self._change_feed_state = self._feed_options.pop("changeFeedState") + self._change_feed_state: ChangeFeedStateV2 = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV2): - raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") + raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version " + f"{type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function @@ -140,34 +146,47 @@ def callback(): return self.fetch_next_block() - def fetch_change_feed_items(self, fetch_function) -> list[dict[str, any]]: + def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: new_options = copy.deepcopy(self._feed_options) new_options["changeFeedState"] = self._change_feed_state self._change_feed_state.populate_feed_options(new_options) - is_s_time_first_fetch = True + is_s_time_first_fetch = self._change_feed_state._continuation.current_token.token is None while True: (fetched_items, response_headers) = fetch_function(new_options) continuation_key = http_constants.HttpHeaders.ETag - # In change feed queries, the continuation token is always populated. The hasNext() test is whether - # there is any items in the response or not. - # For start time however we get no initial results, so we need to pass continuation token? Is this true? + # In change feed queries, the continuation token is always populated. if fetched_items: self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) + self._change_feed_state._continuation._move_to_next_token() response_headers[continuation_key] = self._get_base64_encoded_continuation() break + + # when there is no items being returned, we will decide to retry based on: + # 1. When processing from point in time, there will be no initial results being returned, + # so we will retry with the new continuation token + # 2. if the feed range of the changeFeedState span multiple physical partitions + # then we will read from the next feed range until we have looped through all physical partitions + self._change_feed_state.apply_not_modified_response() + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + + if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + and is_s_time_first_fetch): + response_headers[continuation_key] = self._get_base64_encoded_continuation() + is_s_time_first_fetch = False + should_retry = True else: - self._change_feed_state.apply_not_modified_response() - self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + self._change_feed_state._continuation._move_to_next_token() response_headers[continuation_key] = self._get_base64_encoded_continuation() - should_retry = self._change_feed_state.should_retry_on_not_modified_response() or is_s_time_first_fetch + should_retry = self._change_feed_state.should_retry_on_not_modified_response() is_s_time_first_fetch = False - if not should_retry: - break + + if not should_retry: + break return fetched_items @@ -178,3 +197,4 @@ def _get_base64_encoded_continuation(self) -> str: base64_bytes = base64.b64encode(json_bytes) # Convert the Base64 bytes to a string return base64_bytes.decode('utf-8') + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 39e94a30c4c0..55f98374252a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -21,6 +21,7 @@ """Iterable change feed results in the Azure Cosmos database service. """ +from typing import Dict, Any from azure.core.paging import PageIterator @@ -63,14 +64,16 @@ def __init__( raise ValueError("Missing changeFeedStateContext in feed options") change_feed_state_context = self._options.pop("changeFeedStateContext") - continuation = continuation_token if continuation_token is not None else change_feed_state_context.pop("continuation", None) + continuation = continuation_token if continuation_token is not None\ + else change_feed_state_context.pop("continuation", None) # analysis and validate continuation token # there are two types of continuation token we support currently: # v1 version: the continuation token would just be the _etag, # which is being returned when customer is using partition_key_range_id, # which is under deprecation and does not support split/merge - # v2 version: the continuation token will be base64 encoded composition token which includes full change feed state + # v2 version: the continuation token will be base64 encoded composition token + # which includes full change feed state if continuation is not None: if is_base64_encoded(continuation): change_feed_state_context["continuationFeedRange"] = continuation @@ -135,7 +138,7 @@ def _initialize_change_feed_fetcher(self): self._fetch_function ) - def _validate_change_feed_state_context(self, change_feed_state_context: dict[str, any]) -> None: + def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): # if continuation token is in v1 format, throw exception if feed_range is set @@ -152,7 +155,6 @@ def _validate_change_feed_state_context(self, change_feed_state_context: dict[st key in change_feed_state_context and change_feed_state_context[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") - - + "partition_key_range_id, partition_key, feed_range are exclusive parameters," + " please only set one of them") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py index 76a4d6b56803..632f87715819 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py @@ -25,7 +25,7 @@ from abc import ABC, abstractmethod from datetime import datetime, timezone from enum import Enum -from typing import Optional, Union, Literal, Any +from typing import Optional, Union, Literal, Any, Dict from azure.cosmos import http_constants from azure.cosmos._routing.routing_range import Range @@ -43,38 +43,39 @@ class ChangeFeedStartFromInternal(ABC): type_property_name = "Type" @abstractmethod - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: pass @staticmethod - def from_start_time(start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]]) -> 'ChangeFeedStartFromInternal': + def from_start_time( + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]]) -> 'ChangeFeedStartFromInternal': if start_time is None: return ChangeFeedStartFromNow() - elif isinstance(start_time, datetime): + if isinstance(start_time, datetime): return ChangeFeedStartFromPointInTime(start_time) - elif start_time.lower() == ChangeFeedStartFromType.NOW.value.lower(): + if start_time.lower() == ChangeFeedStartFromType.NOW.value.lower(): return ChangeFeedStartFromNow() - elif start_time.lower() == ChangeFeedStartFromType.BEGINNING.value.lower(): + if start_time.lower() == ChangeFeedStartFromType.BEGINNING.value.lower(): return ChangeFeedStartFromBeginning() - else: - raise ValueError(f"Invalid start_time '{start_time}'") + + raise ValueError(f"Invalid start_time '{start_time}'") @staticmethod - def from_json(data: dict[str, any]) -> 'ChangeFeedStartFromInternal': + def from_json(data: Dict[str, Any]) -> 'ChangeFeedStartFromInternal': change_feed_start_from_type = data.get(ChangeFeedStartFromInternal.type_property_name) if change_feed_start_from_type is None: raise ValueError(f"Invalid start from json [Missing {ChangeFeedStartFromInternal.type_property_name}]") if change_feed_start_from_type == ChangeFeedStartFromType.BEGINNING.value: return ChangeFeedStartFromBeginning.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.LEASE.value: + if change_feed_start_from_type == ChangeFeedStartFromType.LEASE.value: return ChangeFeedStartFromETagAndFeedRange.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.NOW.value: + if change_feed_start_from_type == ChangeFeedStartFromType.NOW.value: return ChangeFeedStartFromNow.from_json(data) - elif change_feed_start_from_type == ChangeFeedStartFromType.POINT_IN_TIME.value: + if change_feed_start_from_type == ChangeFeedStartFromType.POINT_IN_TIME.value: return ChangeFeedStartFromPointInTime.from_json(data) - else: - raise ValueError(f"Can not process changeFeedStartFrom for type {change_feed_start_from_type}") + + raise ValueError(f"Can not process changeFeedStartFrom for type {change_feed_start_from_type}") @abstractmethod def populate_request_headers(self, request_headers) -> None: @@ -85,7 +86,7 @@ class ChangeFeedStartFromBeginning(ChangeFeedStartFromInternal): """Class for change feed start from beginning implementation in the Azure Cosmos database service. """ - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.BEGINNING.value } @@ -94,7 +95,7 @@ def populate_request_headers(self, request_headers) -> None: pass # there is no headers need to be set for start from beginning @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromBeginning': + def from_json(cls, data: Dict[str, Any]) -> 'ChangeFeedStartFromBeginning': return ChangeFeedStartFromBeginning() @@ -112,7 +113,7 @@ def __init__(self, etag, feed_range): self._etag = etag self._feed_range = feed_range - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.LEASE.value, self._etag_property_name: self._etag, @@ -120,7 +121,7 @@ def to_dict(self) -> dict[str, Any]: } @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromETagAndFeedRange': + def from_json(cls, data: Dict[str, Any]) -> 'ChangeFeedStartFromETagAndFeedRange': etag = data.get(cls._etag_property_name) if etag is None: raise ValueError(f"Invalid change feed start from [Missing {cls._etag_property_name}]") @@ -141,7 +142,7 @@ class ChangeFeedStartFromNow(ChangeFeedStartFromInternal): """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. """ - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.NOW.value } @@ -150,7 +151,7 @@ def populate_request_headers(self, request_headers) -> None: request_headers[http_constants.HttpHeaders.IfNoneMatch] = "*" @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromNow': + def from_json(cls, data: Dict[str, Any]) -> 'ChangeFeedStartFromNow': return ChangeFeedStartFromNow() @@ -166,7 +167,7 @@ def __init__(self, start_time: datetime): self._start_time = start_time - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.POINT_IN_TIME.value, self._point_in_time_ms_property_name: @@ -178,7 +179,7 @@ def populate_request_headers(self, request_headers) -> None: self._start_time.astimezone(timezone.utc).strftime('%a, %d %b %Y %H:%M:%S GMT') @classmethod - def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': + def from_json(cls, data: Dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': point_in_time_ms = data.get(cls._point_in_time_ms_property_name) if point_in_time_ms is None: raise ValueError(f"Invalid change feed start from {cls._point_in_time_ms_property_name} ") @@ -187,3 +188,4 @@ def from_json(cls, data: dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': return ChangeFeedStartFromPointInTime(point_in_time) + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 210563ab8411..2d2f4d0d6ae2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -27,7 +27,7 @@ import collections import json from abc import ABC, abstractmethod -from typing import Optional, Union, List, Any +from typing import Optional, Union, List, Any, Dict, Deque from azure.cosmos import http_constants from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromInternal, \ @@ -35,6 +35,7 @@ from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._change_feed.feed_range_composite_continuation_token import FeedRangeCompositeContinuation +from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range from azure.cosmos._utils import is_key_exists_and_not_none @@ -46,11 +47,21 @@ class ChangeFeedState(ABC): version_property_name = "v" @abstractmethod - def populate_feed_options(self, feed_options: dict[str, any]) -> None: + def populate_feed_options(self, feed_options: Dict[str, Any]) -> None: pass @abstractmethod - def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, request_headers: dict[str, any]) -> None: + def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: + pass + + @abstractmethod + async def populate_request_headers_async( + self, + async_routing_provider: AsyncSmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: pass @abstractmethod @@ -61,28 +72,32 @@ def apply_server_response_continuation(self, continuation: str) -> None: def from_json( container_link: str, container_rid: str, - change_feed_state_context: dict[str, Any]): + change_feed_state_context: Dict[str, Any]): if (is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId") or is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId")): return ChangeFeedStateV1.from_json(container_link, container_rid, change_feed_state_context) - else: - if is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): - # get changeFeedState from continuation - continuation_json_str = base64.b64decode(change_feed_state_context["continuationFeedRange"]).decode('utf-8') - continuation_json = json.loads(continuation_json_str) - version = continuation_json.get(ChangeFeedState.version_property_name) - if version is None: - raise ValueError("Invalid base64 encoded continuation string [Missing version]") - elif version == "V2": - return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) - else: - raise ValueError("Invalid base64 encoded continuation string [Invalid version]") - # when there is no continuation token, by default construct ChangeFeedStateV2 - return ChangeFeedStateV2.from_initial_state(container_link, container_rid, change_feed_state_context) + + if is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + # get changeFeedState from continuation + continuation_json_str = base64.b64decode(change_feed_state_context["continuationFeedRange"]).decode( + 'utf-8') + continuation_json = json.loads(continuation_json_str) + version = continuation_json.get(ChangeFeedState.version_property_name) + if version is None: + raise ValueError("Invalid base64 encoded continuation string [Missing version]") + + if version == "V2": + return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) + + raise ValueError("Invalid base64 encoded continuation string [Invalid version]") + + # when there is no continuation token, by default construct ChangeFeedStateV2 + return ChangeFeedStateV2.from_initial_state(container_link, container_rid, change_feed_state_context) class ChangeFeedStateV1(ChangeFeedState): - """Change feed state v1 implementation. This is used when partition key range id is used or the continuation is just simple _etag + """Change feed state v1 implementation. + This is used when partition key range id is used or the continuation is just simple _etag """ def __init__( @@ -92,7 +107,7 @@ def __init__( change_feed_start_from: ChangeFeedStartFromInternal, partition_key_range_id: Optional[str] = None, partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, - continuation: Optional[str] = None): + continuation: Optional[str] = None): # pylint: disable=line-too-long self._container_link = container_link self._container_rid = container_rid @@ -106,27 +121,48 @@ def container_rid(self): return self._container_rid @classmethod - def from_json(cls, container_link: str, container_rid: str, data: dict[str, Any]) -> 'ChangeFeedStateV1': + def from_json( + cls, + container_link: str, + container_rid: str, + change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV1': return cls( container_link, container_rid, - ChangeFeedStartFromInternal.from_start_time(data.get("startTime")), - data.get("partitionKeyRangeId"), - data.get("partitionKey"), - data.get("continuationPkRangeId") + ChangeFeedStartFromInternal.from_start_time(change_feed_state_context.get("startTime")), + change_feed_state_context.get("partitionKeyRangeId"), + change_feed_state_context.get("partitionKey"), + change_feed_state_context.get("continuationPkRangeId") ) - def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, Any]) -> None: - headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: + request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. - self._change_feed_start_from.populate_request_headers(headers) + self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation: - headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation + request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation + + async def populate_request_headers_async( + self, + routing_provider: AsyncSmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: - def populate_feed_options(self, feed_options: dict[str, any]) -> None: + request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(request_headers) + if self._continuation: + request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation + + def populate_feed_options(self, feed_options: Dict[str, Any]) -> None: if self._partition_key_range_id is not None: feed_options["partitionKeyRangeId"] = self._partition_key_range_id if self._partition_key is not None: @@ -148,7 +184,7 @@ def __init__( container_rid: str, feed_range: FeedRange, change_feed_start_from: ChangeFeedStartFromInternal, - continuation: Optional[FeedRangeCompositeContinuation] = None): + continuation: Optional[FeedRangeCompositeContinuation]): self._container_link = container_link self._container_rid = container_rid @@ -156,39 +192,49 @@ def __init__( self._change_feed_start_from = change_feed_start_from self._continuation = continuation if self._continuation is None: - composite_continuation_token_queue = collections.deque() + composite_continuation_token_queue: Deque = collections.deque() composite_continuation_token_queue.append( CompositeContinuationToken( self._feed_range.get_normalized_range(), None)) self._continuation =\ - FeedRangeCompositeContinuation(self._container_rid, self._feed_range, composite_continuation_token_queue) + FeedRangeCompositeContinuation( + self._container_rid, + self._feed_range, + composite_continuation_token_queue) @property def container_rid(self) -> str : return self._container_rid - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: return { self.version_property_name: "V2", self.container_rid_property_name: self._container_rid, self.change_feed_mode_property_name: "Incremental", self.change_feed_start_from_property_name: self._change_feed_start_from.to_dict(), - self.continuation_property_name: self._continuation.to_dict() + self.continuation_property_name: self._continuation.to_dict() if self._continuation is not None else None } - def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, headers: dict[str, any]) -> None: - headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue + def populate_request_headers( + self, + routing_provider: SmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: + request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. - self._change_feed_start_from.populate_request_headers(headers) + # of the documents may not be sequential. + # So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, + # we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation.current_token is not None and self._continuation.current_token.token is not None: change_feed_start_from_feed_range_and_etag =\ - ChangeFeedStartFromETagAndFeedRange(self._continuation.current_token.token, self._continuation.current_token.feed_range) - change_feed_start_from_feed_range_and_etag.populate_request_headers(headers) + ChangeFeedStartFromETagAndFeedRange( + self._continuation.current_token.token, + self._continuation.current_token.feed_range) + change_feed_start_from_feed_range_and_etag.populate_request_headers(request_headers) # based on the feed range to find the overlapping partition key range id over_lapping_ranges =\ @@ -197,28 +243,87 @@ def populate_request_headers(self, routing_provider: SmartRoutingMapProvider, he [self._continuation.current_token.feed_range]) if len(over_lapping_ranges) > 1: - raise CosmosFeedRangeGoneError(message= - f"Range {self._continuation.current_token.feed_range}" - f" spans {len(over_lapping_ranges)}" - f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + raise CosmosFeedRangeGoneError( + message= + f"Range {self._continuation.current_token.feed_range}" + f" spans {len(over_lapping_ranges)}" + f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + + overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) + if overlapping_feed_range == self._continuation.current_token.feed_range: + # exactly mapping to one physical partition, only need to set the partitionKeyRangeId + request_headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] else: - overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) - if overlapping_feed_range == self._continuation.current_token.feed_range: - # exactly mapping to one physical partition, only need to set the partitionKeyRangeId - headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] - else: - # the current token feed range spans less than single physical partition - # for this case, need to set both the partition key range id and epk filter headers - headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] - headers[http_constants.HttpHeaders.StartEpkString] = self._continuation.current_token.feed_range.min - headers[http_constants.HttpHeaders.EndEpkString] = self._continuation.current_token.feed_range.max + # the current token feed range spans less than single physical partition + # for this case, need to set both the partition key range id and epk filter headers + request_headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + request_headers[ + http_constants.HttpHeaders.StartEpkString] = self._continuation.current_token.feed_range.min + request_headers[ + http_constants.HttpHeaders.EndEpkString] = self._continuation.current_token.feed_range.max + + async def populate_request_headers_async( + self, + async_routing_provider: AsyncSmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: + request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - def populate_feed_options(self, feed_options: dict[str, any]) -> None: + # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time + # of the documents may not be sequential. + # So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, + # we will need to always pass the start time in the header. + self._change_feed_start_from.populate_request_headers(request_headers) + + if self._continuation.current_token is not None and self._continuation.current_token.token is not None: + change_feed_start_from_feed_range_and_etag = \ + ChangeFeedStartFromETagAndFeedRange( + self._continuation.current_token.token, + self._continuation.current_token.feed_range) + change_feed_start_from_feed_range_and_etag.populate_request_headers(request_headers) + + # based on the feed range to find the overlapping partition key range id + over_lapping_ranges = \ + await async_routing_provider.get_overlapping_ranges( + self._container_link, + [self._continuation.current_token.feed_range]) + + if len(over_lapping_ranges) > 1: + raise CosmosFeedRangeGoneError( + message= + f"Range {self._continuation.current_token.feed_range}" + f" spans {len(over_lapping_ranges)}" + f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + + overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) + if overlapping_feed_range == self._continuation.current_token.feed_range: + # exactly mapping to one physical partition, only need to set the partitionKeyRangeId + request_headers[http_constants.HttpHeaders.PartitionKeyRangeID] = over_lapping_ranges[0]["id"] + else: + # the current token feed range spans less than single physical partition + # for this case, need to set both the partition key range id and epk filter headers + request_headers[http_constants.HttpHeaders.PartitionKeyRangeID] = \ + over_lapping_ranges[0]["id"] + request_headers[http_constants.HttpHeaders.StartEpkString] = \ + self._continuation.current_token.feed_range.min + request_headers[http_constants.HttpHeaders.EndEpkString] = \ + self._continuation.current_token.feed_range.max + + def populate_feed_options(self, feed_options: Dict[str, Any]) -> None: pass - def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, resource_link: str) -> None: + def handle_feed_range_gone( + self, + routing_provider: SmartRoutingMapProvider, + resource_link: str) -> None: self._continuation.handle_feed_range_gone(routing_provider, resource_link) + async def handle_feed_range_gone_async( + self, + routing_provider: AsyncSmartRoutingMapProvider, + resource_link: str) -> None: + await self._continuation.handle_feed_range_gone_async(routing_provider, resource_link) + def apply_server_response_continuation(self, continuation: str) -> None: self._continuation.apply_server_response_continuation(continuation) @@ -233,17 +338,18 @@ def from_continuation( cls, container_link: str, container_rid: str, - continuation_json: dict[str, Any]) -> 'ChangeFeedStateV2': + continuation_json: Dict[str, Any]) -> 'ChangeFeedStateV2': container_rid_from_continuation = continuation_json.get(ChangeFeedStateV2.container_rid_property_name) if container_rid_from_continuation is None: raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.container_rid_property_name}]") - elif container_rid_from_continuation != container_rid: + if container_rid_from_continuation != container_rid: raise ValueError("Invalid continuation: [Mismatch collection rid]") change_feed_start_from_data = continuation_json.get(ChangeFeedStateV2.change_feed_start_from_property_name) if change_feed_start_from_data is None: - raise ValueError(f"Invalid continuation: [Missing {ChangeFeedStateV2.change_feed_start_from_property_name}]") + raise ValueError(f"Invalid continuation:" + f" [Missing {ChangeFeedStateV2.change_feed_start_from_property_name}]") change_feed_start_from = ChangeFeedStartFromInternal.from_json(change_feed_start_from_data) continuation_data = continuation_json.get(ChangeFeedStateV2.continuation_property_name) @@ -270,7 +376,10 @@ def from_initial_state( feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) elif is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyFeedRange"): - feed_range = FeedRangePartitionKey(change_feed_state_context["partitionKey"], change_feed_state_context["partitionKeyFeedRange"]) + feed_range =\ + FeedRangePartitionKey( + change_feed_state_context["partitionKey"], + change_feed_state_context["partitionKeyFeedRange"]) else: raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") else: @@ -283,7 +392,8 @@ def from_initial_state( False) ) - change_feed_start_from = ChangeFeedStartFromInternal.from_start_time(change_feed_state_context.get("startTime")) + change_feed_start_from = ( + ChangeFeedStartFromInternal.from_start_time(change_feed_state_context.get("startTime"))) return cls( container_link=container_link, container_rid=collection_rid, @@ -291,3 +401,5 @@ def from_initial_state( change_feed_start_from=change_feed_start_from, continuation=None) + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index 1b6ef79e4176..3b9707371fb8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -22,9 +22,8 @@ """Internal class for feed range implementation in the Azure Cosmos database service. """ -import json from abc import ABC, abstractmethod -from typing import Union, List +from typing import Union, List, Dict, Any from azure.cosmos._routing.routing_range import Range from azure.cosmos._utils import is_key_exists_and_not_none @@ -38,7 +37,7 @@ def get_normalized_range(self) -> Range: pass @abstractmethod - def to_dict(self) -> dict[str, any]: + def to_dict(self) -> Dict[str, Any]: pass class FeedRangePartitionKey(FeedRange): @@ -47,7 +46,7 @@ class FeedRangePartitionKey(FeedRange): def __init__( self, pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], - feed_range: Range): + feed_range: Range): # pylint: disable=line-too-long if pk_value is None: raise ValueError("PartitionKey cannot be None") @@ -60,26 +59,31 @@ def __init__( def get_normalized_range(self) -> Range: return self._feed_range.to_normalized_range() - def to_dict(self) -> dict[str, any]: + def to_dict(self) -> Dict[str, Any]: if isinstance(self._pk_value, _Undefined): return { self.type_property_name: [{}] } - elif isinstance(self._pk_value, _Empty): + if isinstance(self._pk_value, _Empty): return { self.type_property_name: [] } - else: - return { self.type_property_name: json.dumps(self._pk_value) } + if isinstance(self._pk_value, list): + return { self.type_property_name: [item for item in self._pk_value] } + + return { self.type_property_name: self._pk_value } @classmethod - def from_json(cls, data: dict[str, any], feed_range: Range) -> 'FeedRangePartitionKey': + def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartitionKey': if is_key_exists_and_not_none(data, cls.type_property_name): pk_value = data.get(cls.type_property_name) - if isinstance(pk_value, list): - if not pk_value: - return cls(_Empty(), feed_range) - if pk_value == [{}]: + if not pk_value: + return cls(_Empty(), feed_range) + if pk_value is [{}]: return cls(_Undefined(), feed_range) + if isinstance(pk_value, list): + return cls([item for item in pk_value], feed_range) + return cls(data[cls.type_property_name], feed_range) + + raise ValueError(f"Can not parse FeedRangePartitionKey from the json," + f" there is no property {cls.type_property_name}") - return cls(json.loads(data.get(cls.type_property_name)), feed_range) - raise ValueError(f"Can not parse FeedRangePartitionKey from the json, there is no property {cls.type_property_name}") class FeedRangeEpk(FeedRange): type_property_name = "Range" @@ -93,13 +97,13 @@ def __init__(self, feed_range: Range): def get_normalized_range(self) -> Range: return self._range.to_normalized_range() - def to_dict(self) -> dict[str, any]: + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: self._range.to_dict() } @classmethod - def from_json(cls, data: dict[str, any]) -> 'FeedRangeEpk': + def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeEpk': if is_key_exists_and_not_none(data, cls.type_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) return cls(feed_range) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index 26abd66ba132..fb67b6b4603c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -22,13 +22,13 @@ """Internal class for change feed continuation token by feed range in the Azure Cosmos database service. """ -import collections from collections import deque -from typing import Any +from typing import Any, Deque, Dict from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider +from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range from azure.cosmos._utils import is_key_exists_and_not_none @@ -42,7 +42,7 @@ def __init__( self, container_rid: str, feed_range: FeedRange, - continuation: collections.deque[CompositeContinuationToken]): + continuation: Deque[CompositeContinuationToken]): if container_rid is None: raise ValueError("container_rid is missing") @@ -56,7 +56,7 @@ def __init__( def current_token(self): return self._current_token - def to_dict(self) -> dict[str, Any]: + def to_dict(self) -> Dict[str, Any]: json_data = { self._version_property_name: "v2", self._container_rid_property_name: self._container_rid, @@ -75,14 +75,18 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': container_rid = data.get(cls._container_rid_property_name) if container_rid is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._container_rid_property_name}]") + raise ValueError(f"Invalid feed range composite continuation token " + f"[Missing {cls._container_rid_property_name}]") continuation_data = data.get(cls._continuation_property_name) if continuation_data is None: - raise ValueError(f"Invalid feed range composite continuation token [Missing {cls._continuation_property_name}]") + raise ValueError(f"Invalid feed range composite continuation token " + f"[Missing {cls._continuation_property_name}]") if not isinstance(continuation_data, list) or len(continuation_data) == 0: - raise ValueError(f"Invalid feed range composite continuation token [The {cls._continuation_property_name} must be non-empty array]") - continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) for child_range_continuation_token in continuation_data] + raise ValueError(f"Invalid feed range composite continuation token " + f"[The {cls._continuation_property_name} must be non-empty array]") + continuation = [CompositeContinuationToken.from_json(child_range_continuation_token) + for child_range_continuation_token in continuation_data] # parsing feed range if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): @@ -94,7 +98,10 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': return cls(container_rid=container_rid, feed_range=feed_range, continuation=deque(continuation)) - def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, collection_link: str) -> None: + def handle_feed_range_gone( + self, + routing_provider: SmartRoutingMapProvider, + collection_link: str) -> None: overlapping_ranges = routing_provider.get_overlapping_ranges(collection_link, [self._current_token.feed_range]) if len(overlapping_ranges) == 1: @@ -105,20 +112,47 @@ def handle_feed_range_gone(self, routing_provider: SmartRoutingMapProvider, coll # For each new child range, using the continuation token from the parent self._continuation.popleft() for child_range in overlapping_ranges: - self._continuation.append(CompositeContinuationToken(Range.PartitionKeyRangeToRange(child_range), self._current_token.token)) + self._continuation.append( + CompositeContinuationToken( + Range.PartitionKeyRangeToRange(child_range), + self._current_token.token)) + + self._current_token = self._continuation[0] + + async def handle_feed_range_gone_async( + self, + routing_provider: AsyncSmartRoutingMapProvider, + collection_link: str) -> None: + overlapping_ranges = \ + await routing_provider.get_overlapping_ranges( + collection_link, + [self._current_token.feed_range]) + + if len(overlapping_ranges) == 1: + # merge,reusing the existing the feedRange and continuationToken + pass + else: + # split, remove the parent range and then add new child ranges. + # For each new child range, using the continuation token from the parent + self._continuation.popleft() + for child_range in overlapping_ranges: + self._continuation.append( + CompositeContinuationToken( + Range.PartitionKeyRangeToRange(child_range), + self._current_token.token)) self._current_token = self._continuation[0] def should_retry_on_not_modified_response(self) -> bool: - # when getting 304(Not Modified) response from one sub feed range, we will try to fetch for the next sub feed range + # when getting 304(Not Modified) response from one sub feed range, + # we will try to fetch for the next sub feed range # we will repeat the above logic until we have looped through all sub feed ranges # TODO: validate the response headers, can we get the status code if len(self._continuation) > 1: return self._current_token.feed_range != self._initial_no_result_range - else: - return False + return False def _move_to_next_token(self) -> None: first_composition_token = self._continuation.popleft() @@ -128,7 +162,6 @@ def _move_to_next_token(self) -> None: def apply_server_response_continuation(self, etag) -> None: self._current_token.update_token(etag) - self._move_to_next_token() def apply_not_modified_response(self) -> None: if self._initial_no_result_range is None: 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 a81ab438cbf2..87e72391cf0a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -26,6 +26,8 @@ import os import urllib.parse from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast, Type +from typing_extensions import TypedDict +from urllib3.util.retry import Retry from azure.core import PipelineClient from azure.core.credentials import TokenCredential @@ -42,8 +44,6 @@ ) from azure.core.pipeline.transport import HttpRequest, \ HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import -from typing_extensions import TypedDict -from urllib3.util.retry import Retry from . import _base as base from . import _global_endpoint_manager as global_endpoint_manager @@ -3025,7 +3025,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: partition_key_range_id ) - change_feed_state = options.pop("changeFeedState", None) + change_feed_state = options.get("changeFeedState", None) if change_feed_state and isinstance(change_feed_state, ChangeFeedState): change_feed_state.populate_request_headers(self._routing_map_provider, headers) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index f3269af47271..a2d789f20644 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -108,7 +108,7 @@ def to_normalized_range(self): return Range(normalized_min, normalized_max, True, False) def add_to_effective_partition_key(self, effective_partition_key: str, value: int): - if value != 1 and value != -1: + if value not in (-1, 1): raise ValueError("Invalid value - only 1 or -1 is allowed") byte_array = self.hex_binary_to_byte_array(effective_partition_key) @@ -117,15 +117,13 @@ def add_to_effective_partition_key(self, effective_partition_key: str, value: in if byte_array[i] < 255: byte_array[i] += 1 break - else: - byte_array[i] = 0 + byte_array[i] = 0 else: for i in range(len(byte_array) - 1, -1, -1): if byte_array[i] != 0: byte_array[i] -= 1 break - else: - byte_array[i] = 255 + byte_array[i] = 255 return binascii.hexlify(byte_array).decode() @@ -143,8 +141,8 @@ def from_base64_encoded_json_string(cls, data: str): feed_range_json_string = base64.b64decode(data, validate=True).decode('utf-8') feed_range_json = json.loads(feed_range_json_string) return cls.ParseFromDict(feed_range_json) - except Exception: - raise ValueError(f"Invalid feed_range json string {data}") + except Exception as exc: + raise ValueError(f"Invalid feed_range json string {data}") from exc def to_base64_encoded_string(self): data_json = json.dumps(self.to_dict()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 1c03b8a054c5..cf4b4977ed44 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -79,7 +79,7 @@ def is_base64_encoded(data: str) -> bool: return True except (json.JSONDecodeError, ValueError): return False - - -def is_key_exists_and_not_none(data: dict[str, Any], key: str) -> bool: + +def is_key_exists_and_not_none(data: Dict[str, Any], key: str) -> bool: return key in data and data[key] is not None + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 5e8da4f54cf4..e70bdf9ab546 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -24,12 +24,12 @@ import warnings from datetime import datetime from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload +from typing_extensions import Literal from azure.core import MatchConditions from azure.core.async_paging import AsyncItemPaged from azure.core.tracing.decorator import distributed_trace from azure.core.tracing.decorator_async import distributed_trace_async # type: ignore -from typing_extensions import Literal from ._cosmos_client_connection_async import CosmosClientConnection from ._scripts import ScriptsProxy @@ -49,13 +49,14 @@ NonePartitionKeyValue, _return_undefined_or_empty_partition_key, _Empty, - _Undefined, PartitionKey + _Undefined ) __all__ = ("ContainerProxy",) # pylint: disable=protected-access, too-many-lines # pylint: disable=missing-client-constructor-parameter-credential,missing-client-constructor-parameter-kwargs +# pylint: disable=too-many-public-methods PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long @@ -136,13 +137,6 @@ async def _set_partition_key( return _return_undefined_or_empty_partition_key(await self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) - async def __is_prefix_partition_key(self, partition_key: PartitionKeyType) -> bool: - - properties = await self._get_properties() - pk_properties = properties.get("partitionKey") - partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) - return partition_key_definition._is_prefix_partition_key(partition_key) - @distributed_trace_async async def read( self, @@ -500,17 +494,17 @@ def query_items_change_feed( partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: + ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. Now: Processing change feed from the current time, so only events for all future changes will be retrieved. ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current (NOW) - :param PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :keyword PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) + :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. :returns: An AsyncItemPaged of items (dicts). @@ -519,7 +513,7 @@ def query_items_change_feed( ... @overload - async def query_items_change_feed( + def query_items_change_feed( self, *, feed_range: Optional[str] = None, @@ -527,17 +521,17 @@ async def query_items_change_feed( start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: + ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + :keyword str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. Now: Processing change feed from the current time, so only events for all future changes will be retrieved. ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current (NOW) - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :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. :returns: An AsyncItemPaged of items (dicts). @@ -553,12 +547,12 @@ def query_items_change_feed( max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: + ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param str continuation: The continuation token retrieved from previous response. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :keyword str continuation: The continuation token retrieved from previous response. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :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. :returns: An AsyncItemPaged of items (dicts). @@ -571,7 +565,7 @@ def query_items_change_feed( self, *args: Any, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: + ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements if is_key_exists_and_not_none(kwargs, "priority"): kwargs['priority'] = kwargs['priority'] @@ -1216,7 +1210,7 @@ async def execute_item_batch( async def read_feed_ranges( self, **kwargs: Any - ) -> List[str]: + ) -> List[str]: # pylint: disable=unused-argument partition_key_ranges =\ await self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, 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 7cccac695769..b2a9aaff9ec1 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 @@ -51,7 +51,7 @@ from .._base import _set_properties_cache from .. import documents from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable -from .._change_feed.aio.change_feed_state import ChangeFeedState +from .._change_feed.change_feed_state import ChangeFeedState from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants @@ -2814,9 +2814,9 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: ) headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) - change_feed_state = options.pop("changeFeedState", None) + change_feed_state = options.get("changeFeedState", None) if change_feed_state and isinstance(change_feed_state, ChangeFeedState): - await change_feed_state.populate_request_headers(self._routing_map_provider, headers) + await change_feed_state.populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) if response_hook: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index e2fa8ad4071e..7a6e466bfc1d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -24,11 +24,11 @@ import warnings from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload +from typing_extensions import Literal from azure.core import MatchConditions from azure.core.paging import ItemPaged from azure.core.tracing.decorator import distributed_trace -from typing_extensions import Literal from ._base import ( build_options, @@ -328,17 +328,19 @@ def query_items_change_feed( partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: + ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) - :param PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: + The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :keyword PartitionKeyType partition_key: The partition key that is used to define the scope + (logical partition or a subset of a container) + :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. :returns: An Iterable of items (dicts). @@ -355,17 +357,19 @@ def query_items_change_feed( start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: + ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :keyword str feed_range: The feed range that is used to define the scope. + By default, the scope will be the entire container. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword Union[datetime, Literal["Now", "Beginning"]] + start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current (NOW) + :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. :returns: An Iterable of items (dicts). @@ -381,12 +385,12 @@ def query_items_change_feed( max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: + ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :param str continuation: The continuation token retrieved from previous response. - :param int max_item_count: Max number of items to be returned in the enumeration operation. - :param Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each + :keyword str continuation: The continuation token retrieved from previous response. + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :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. :returns: An Iterable of items (dicts). @@ -399,7 +403,7 @@ def query_items_change_feed( self, *args: Any, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: + ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements if is_key_exists_and_not_none(kwargs, "priority"): kwargs['priority'] = kwargs['priority'] @@ -1273,7 +1277,7 @@ def delete_all_items_by_partition_key( def read_feed_ranges( self, **kwargs: Any - ) -> List[str]: + ) -> List[str]: # pylint: disable=unused-argument partition_key_ranges =\ self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py index 768890dacfa6..6913979cf81d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py @@ -167,4 +167,4 @@ def _container_recreate_exception(e) -> bool: def _is_partition_split_or_merge(e): - return e.status_code == StatusCodes.GONE and e.status_code == SubStatusCodes.COMPLETING_SPLIT \ No newline at end of file + return e.status_code == StatusCodes.GONE and e.status_code == SubStatusCodes.COMPLETING_SPLIT diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index c89e1d9ac771..21aca775cbd5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -278,7 +278,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: + 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": return False if isinstance(partition_key, list) and len(self.path) == len(partition_key): From 36990ef428145a56660b0dc784f62250ef67145f Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 20 Aug 2024 09:47:59 -0700 Subject: [PATCH 05/59] fix pylint --- .../_change_feed/aio/change_feed_fetcher.py | 48 +++++++++---------- .../_change_feed/change_feed_fetcher.py | 1 - .../_change_feed/change_feed_iterable.py | 1 - .../_change_feed/change_feed_start_from.py | 3 -- .../cosmos/_change_feed/change_feed_state.py | 28 ++++++----- .../azure/cosmos/_change_feed/feed_range.py | 8 ++-- ...feed_range_composite_continuation_token.py | 1 - .../azure/cosmos/aio/_container.py | 16 +++++-- .../azure-cosmos/azure/cosmos/container.py | 29 ++++++----- 9 files changed, 72 insertions(+), 63 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index 4d85a891ac3f..ee926aa0e92c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -120,7 +120,8 @@ def __init__( self._change_feed_state: ChangeFeedStateV2 = self._feed_options.pop("changeFeedState") if not isinstance(self._change_feed_state, ChangeFeedStateV2): - raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version {type(self._change_feed_state)}") + raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version " + f"{type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function @@ -170,29 +171,29 @@ async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: response_headers[continuation_key] = self._get_base64_encoded_continuation() break - # when there is no items being returned, we will decide to retry based on: - # 1. When processing from point in time, there will be no initial results being returned, - # so we will retry with the new continuation token - # 2. if the feed range of the changeFeedState span multiple physical partitions - # then we will read from the next feed range until we have looped through all physical partitions - self._change_feed_state.apply_not_modified_response() - self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + # when there is no items being returned, we will decide to retry based on: + # 1. When processing from point in time, there will be no initial results being returned, + # so we will retry with the new continuation token + # 2. if the feed range of the changeFeedState span multiple physical partitions + # then we will read from the next feed range until we have looped through all physical partitions + self._change_feed_state.apply_not_modified_response() + self._change_feed_state.apply_server_response_continuation( + response_headers.get(continuation_key)) + + #TODO: can this part logic be simplified + if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + and is_s_time_first_fetch): + response_headers[continuation_key] = self._get_base64_encoded_continuation() + is_s_time_first_fetch = False + should_retry = True + else: + self._change_feed_state._continuation._move_to_next_token() + response_headers[continuation_key] = self._get_base64_encoded_continuation() + should_retry = self._change_feed_state.should_retry_on_not_modified_response() + is_s_time_first_fetch = False - #TODO: can this part logic be simplified - if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) - and is_s_time_first_fetch): - response_headers[continuation_key] = self._get_base64_encoded_continuation() - is_s_time_first_fetch = False - should_retry = True - else: - self._change_feed_state._continuation._move_to_next_token() - response_headers[continuation_key] = self._get_base64_encoded_continuation() - should_retry = self._change_feed_state.should_retry_on_not_modified_response() - is_s_time_first_fetch = False - - if not should_retry: - break + if not should_retry: + break return fetched_items @@ -203,4 +204,3 @@ def _get_base64_encoded_continuation(self) -> str: base64_bytes = base64.b64encode(json_bytes) # Convert the Base64 bytes to a string return base64_bytes.decode('utf-8') - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index abad86dfd119..92f0b2446f74 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -197,4 +197,3 @@ def _get_base64_encoded_continuation(self) -> str: base64_bytes = base64.b64encode(json_bytes) # Convert the Base64 bytes to a string return base64_bytes.decode('utf-8') - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 55f98374252a..4b03e33d0122 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -157,4 +157,3 @@ def _validate_change_feed_state_context(self, change_feed_state_context: Dict[st raise ValueError( "partition_key_range_id, partition_key, feed_range are exclusive parameters," " please only set one of them") - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py index 632f87715819..30d0ce787983 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py @@ -186,6 +186,3 @@ def from_json(cls, data: Dict[str, Any]) -> 'ChangeFeedStartFromPointInTime': point_in_time = datetime.fromtimestamp(point_in_time_ms).astimezone(timezone.utc) return ChangeFeedStartFromPointInTime(point_in_time) - - - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 2d2f4d0d6ae2..05d51daace59 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -106,8 +106,8 @@ def __init__( container_rid: str, change_feed_start_from: ChangeFeedStartFromInternal, partition_key_range_id: Optional[str] = None, - partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, - continuation: Optional[str] = None): # pylint: disable=line-too-long + partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, # pylint: disable=line-too-long + continuation: Optional[str] = None): self._container_link = container_link self._container_rid = container_rid @@ -142,22 +142,26 @@ def populate_request_headers( request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + # of the documents may not be sequential. So when reading the changeFeed by LSN, + # it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, + # we will need to always pass the start time in the header. self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation: request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation async def populate_request_headers_async( self, - routing_provider: AsyncSmartRoutingMapProvider, - request_headers: Dict[str, Any]) -> None: + async_routing_provider: AsyncSmartRoutingMapProvider, + request_headers: Dict[str, Any]) -> None: # pylint: disable=unused-argument request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, we will need to always pass the start time in the header. + # of the documents may not be sequential. + # So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. + # In order to guarantee we always get the documents after customer's point start time, + # we will need to always pass the start time in the header. self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation: request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation @@ -190,7 +194,6 @@ def __init__( self._container_rid = container_rid self._feed_range = feed_range self._change_feed_start_from = change_feed_start_from - self._continuation = continuation if self._continuation is None: composite_continuation_token_queue: Deque = collections.deque() composite_continuation_token_queue.append( @@ -202,6 +205,8 @@ def __init__( self._container_rid, self._feed_range, composite_continuation_token_queue) + else: + self._continuation = continuation @property def container_rid(self) -> str : @@ -368,7 +373,7 @@ def from_initial_state( cls, container_link: str, collection_rid: str, - change_feed_state_context: dict[str, Any]) -> 'ChangeFeedStateV2': + change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV2': if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') @@ -400,6 +405,3 @@ def from_initial_state( feed_range=feed_range, change_feed_start_from=change_feed_start_from, continuation=None) - - - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index 3b9707371fb8..856ccd6c5b48 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -65,7 +65,7 @@ def to_dict(self) -> Dict[str, Any]: if isinstance(self._pk_value, _Empty): return { self.type_property_name: [] } if isinstance(self._pk_value, list): - return { self.type_property_name: [item for item in self._pk_value] } + return { self.type_property_name: list(self._pk_value) } return { self.type_property_name: self._pk_value } @@ -75,10 +75,10 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartiti pk_value = data.get(cls.type_property_name) if not pk_value: return cls(_Empty(), feed_range) - if pk_value is [{}]: + if pk_value == [{}]: return cls(_Undefined(), feed_range) if isinstance(pk_value, list): - return cls([item for item in pk_value], feed_range) + return cls(list(pk_value), feed_range) return cls(data[cls.type_property_name], feed_range) raise ValueError(f"Can not parse FeedRangePartitionKey from the json," @@ -107,4 +107,4 @@ def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeEpk': if is_key_exists_and_not_none(data, cls.type_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) return cls(feed_range) - raise ValueError(f"Can not parse FeedRangeEPK from the json, there is no property {cls.type_property_name}") \ No newline at end of file + raise ValueError(f"Can not parse FeedRangeEPK from the json, there is no property {cls.type_property_name}") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index fb67b6b4603c..76bc9f02fde2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -170,4 +170,3 @@ def apply_not_modified_response(self) -> None: @property def feed_range(self) -> FeedRange: return self._feed_range - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index e70bdf9ab546..dd0e6f0e34eb 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -494,7 +494,8 @@ def query_items_change_feed( partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> AsyncItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. @@ -510,6 +511,7 @@ def query_items_change_feed( :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @overload @@ -521,7 +523,8 @@ def query_items_change_feed( start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> AsyncItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. @@ -537,6 +540,7 @@ def query_items_change_feed( :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @overload @@ -547,7 +551,8 @@ def query_items_change_feed( max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> AsyncItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. @@ -558,6 +563,7 @@ def query_items_change_feed( :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @distributed_trace @@ -565,8 +571,8 @@ def query_items_change_feed( self, *args: Any, **kwargs: Any - ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - + ) -> AsyncItemPaged[Dict[str, Any]]: + # pylint: disable=too-many-statements if is_key_exists_and_not_none(kwargs, "priority"): kwargs['priority'] = kwargs['priority'] feed_options = _build_options(kwargs) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 7a6e466bfc1d..94e1f3d7d9fb 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -134,7 +134,7 @@ def _set_partition_key( def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: container_properties = self._get_properties() - partition_key_definition = container_properties.get("partitionKey") + partition_key_definition = container_properties["partitionKey"] partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) return partition_key._get_epk_range_for_partition_key(partition_key_value) @@ -328,7 +328,8 @@ def query_items_change_feed( partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> ItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. @@ -346,6 +347,7 @@ def query_items_change_feed( :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @overload @@ -357,7 +359,9 @@ def query_items_change_feed( start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> ItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long + """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str feed_range: The feed range that is used to define the scope. @@ -375,6 +379,7 @@ def query_items_change_feed( :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @overload @@ -385,7 +390,8 @@ def query_items_change_feed( max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=line-too-long + ) -> ItemPaged[Dict[str, Any]]: + # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. @@ -396,6 +402,7 @@ def query_items_change_feed( :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @distributed_trace @@ -403,8 +410,8 @@ def query_items_change_feed( self, *args: Any, **kwargs: Any - ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - + ) -> ItemPaged[Dict[str, Any]]: + # pylint: disable=too-many-statements if is_key_exists_and_not_none(kwargs, "priority"): kwargs['priority'] = kwargs['priority'] feed_options = build_options(kwargs) @@ -1274,14 +1281,14 @@ def delete_all_items_by_partition_key( self.client_connection.DeleteAllItemsByPartitionKey( collection_link=self.container_link, options=request_options, **kwargs) - def read_feed_ranges( + def read_feed_ranges( # pylint: disable=unused-argument self, **kwargs: Any - ) -> List[str]: # pylint: disable=unused-argument + ) -> List[str]: partition_key_ranges =\ self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, - # default to full range - [Range("", "FF", True, False)]) + [Range("", "FF", True, False)]) # default to full range - return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] \ No newline at end of file + return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() + for partitionKeyRange in partition_key_ranges] From 3c569e89b02cedb2b66b861a65165e70a0ab9ab9 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 20 Aug 2024 11:04:37 -0700 Subject: [PATCH 06/59] added public surface methods --- .../azure-cosmos/azure/cosmos/container.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index a575d6a3304a..7c5e4cbf6098 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1140,3 +1140,19 @@ def delete_all_items_by_partition_key( self.client_connection.DeleteAllItemsByPartitionKey( collection_link=self.container_link, options=request_options, **kwargs) + + def mergeSessionTokens(self, + pks_to_session_tokens: List, + target_pk: PartitionKey + ) -> "Session Token": + """Merge session tokens from different clients to figure out which is the most up to date for a specific + partition key. This should only be used if maintaining own session token or else the sdk will keep track of + session token. + + :param pks_to_session_tokens: list of partition key and session token tuples. + :type pks_to_session_tokens: List[Tuple(string, PartitionKey)] + :param target_pk: partition key to get most up to date session token. + :type target_pk: PartitionKey + :rtype: string + """ + pass \ No newline at end of file From 7479b0c9663a8759e640a4f5e16d4117cd15878b Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 20 Aug 2024 15:39:59 -0700 Subject: [PATCH 07/59] pylint fix --- .../azure/cosmos/_change_feed/change_feed_state.py | 5 +++-- .../feed_range_composite_continuation_token.py | 3 ++- sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py | 1 - sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | 7 ++++--- sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py | 4 ++-- 5 files changed, 11 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 05d51daace59..2bb219eb5497 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -194,7 +194,7 @@ def __init__( self._container_rid = container_rid self._feed_range = feed_range self._change_feed_start_from = change_feed_start_from - if self._continuation is None: + if continuation is None: composite_continuation_token_queue: Deque = collections.deque() composite_continuation_token_queue.append( CompositeContinuationToken( @@ -374,7 +374,8 @@ def from_initial_state( container_link: str, collection_rid: str, change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV2': - + + feed_range: Optional[FeedRange] = None if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') feed_range_json = json.loads(feed_range_str) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index 76bc9f02fde2..2f73af4e5c47 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -23,7 +23,7 @@ database service. """ from collections import deque -from typing import Any, Deque, Dict +from typing import Any, Deque, Dict, Optional from azure.cosmos._change_feed.composite_continuation_token import CompositeContinuationToken from azure.cosmos._change_feed.feed_range import FeedRange, FeedRangeEpk, FeedRangePartitionKey @@ -89,6 +89,7 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': for child_range_continuation_token in continuation_data] # parsing feed range + feed_range: Optional[FeedRange] = None if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): feed_range = FeedRangeEpk.from_json(data) elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index cf4b4977ed44..7690383d375a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -82,4 +82,3 @@ def is_base64_encoded(data: str) -> bool: def is_key_exists_and_not_none(data: Dict[str, Any], key: str) -> bool: return key in data and data[key] is not None - diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index dd0e6f0e34eb..5cae847872d5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1213,14 +1213,15 @@ async def execute_item_batch( return await self.client_connection.Batch( collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) - async def read_feed_ranges( + async def read_feed_ranges( # pylint: disable=unused-argument self, **kwargs: Any - ) -> List[str]: # pylint: disable=unused-argument + ) -> List[str]: partition_key_ranges =\ await self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, # default to full range [Range("", "FF", True, False)]) - return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] + return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() + for partitionKeyRange in partition_key_ranges] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 21aca775cbd5..6a89da3c4f22 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -175,13 +175,13 @@ def _get_epk_range_for_prefix_partition_key( def _get_epk_range_for_partition_key( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + pk_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long ) -> _Range: if self._is_prefix_partition_key(pk_value): return self._get_epk_range_for_prefix_partition_key(pk_value) # else return point range - effective_partition_key_string = self._get_effective_partition_key_string(pk_value) + effective_partition_key_string = self._get_effective_partition_key_string([pk_value]) return _Range(effective_partition_key_string, effective_partition_key_string, True, True) def _get_effective_partition_key_for_hash_partitioning(self) -> str: From 2e766204285056245f6f970862d4312116ad8af5 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 20 Aug 2024 17:53:09 -0700 Subject: [PATCH 08/59] fix --- .../cosmos/_change_feed/change_feed_state.py | 6 ++--- .../azure-cosmos/azure/cosmos/container.py | 8 +++---- .../azure-cosmos/azure/cosmos/exceptions.py | 6 ++--- .../azure/cosmos/partition_key.py | 17 +++++++------ .../azure-cosmos/test/test_change_feed.py | 24 +++++++++++++++++++ .../test/test_change_feed_async.py | 24 +++++++++++++++++++ 6 files changed, 68 insertions(+), 17 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 2bb219eb5497..988932ff3af6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -332,8 +332,8 @@ async def handle_feed_range_gone_async( def apply_server_response_continuation(self, continuation: str) -> None: self._continuation.apply_server_response_continuation(continuation) - def should_retry_on_not_modified_response(self): - self._continuation.should_retry_on_not_modified_response() + def should_retry_on_not_modified_response(self) -> bool: + return self._continuation.should_retry_on_not_modified_response() def apply_not_modified_response(self) -> None: self._continuation.apply_not_modified_response() @@ -374,7 +374,7 @@ def from_initial_state( container_link: str, collection_rid: str, change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV2': - + feed_range: Optional[FeedRange] = None if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 94e1f3d7d9fb..b04972bdee6c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -470,9 +470,9 @@ def query_items_change_feed( feed_options["maxItemCount"] = args[3] if is_key_exists_and_not_none(kwargs, "partition_key"): - partition_key = kwargs.pop('partition_key') - change_feed_state_context["partitionKey"] = self._set_partition_key(partition_key) - change_feed_state_context["partitionKeyFeedRange"] = self._get_epk_range_for_partition_key(partition_key) + change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop('partition_key')) + change_feed_state_context["partitionKeyFeedRange"] =\ + self._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) if is_key_exists_and_not_none(kwargs, "feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') @@ -587,7 +587,7 @@ def query_items( # pylint:disable=docstring-missing-param kwargs["isPrefixPartitionQuery"] = True properties = self._get_properties() kwargs["partitionKeyDefinition"] = properties["partitionKey"] - kwargs["partitionKeyDefinition"]["partition_key"] = partition_key + kwargs["partitionKeyDefinition"]["partition_key"] = self._set_partition_key(partition_key) else: feed_options["partitionKey"] = self._set_partition_key(partition_key) if enable_scan_in_query is not None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py index 6913979cf81d..262de2ffbbbe 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py @@ -142,12 +142,12 @@ def __init__(self, message=None, response=None, **kwargs): """ :param int sub_status_code: HTTP response sub code. """ - self.status_code = StatusCodes.GONE + self.sub_status = SubStatusCodes.PARTITION_KEY_RANGE_GONE self.http_error_message = message - formatted_message = "Status code: %d Sub-status: %d\n%s" % (self.status_code, self.sub_status, str(message)) + formatted_message = "Status code: %d Sub-status: %d\n%s" % (StatusCodes.GONE, self.sub_status, str(message)) super(CosmosHttpResponseError, self).__init__(message=formatted_message, response=response, **kwargs) - + self.status_code = StatusCodes.GONE def _partition_range_is_gone(e): if (e.status_code == http_constants.StatusCodes.GONE diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 6a89da3c4f22..881cce0895e1 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 +from typing import IO, Sequence, Type, Union, overload, List, cast from typing_extensions import Literal from ._cosmos_integers import _UInt64, _UInt128 @@ -149,7 +149,7 @@ def version(self, value: int) -> None: def _get_epk_range_for_prefix_partition_key( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] ) -> _Range: if self.kind != "MultiHash": raise ValueError( @@ -175,13 +175,16 @@ 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: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined] # pylint: disable=line-too-long ) -> _Range: if self._is_prefix_partition_key(pk_value): - return self._get_epk_range_for_prefix_partition_key(pk_value) + return self._get_epk_range_for_prefix_partition_key( + cast(List[Union[str, int, float, bool]], pk_value)) # else return point range - effective_partition_key_string = self._get_effective_partition_key_string([pk_value]) + effective_partition_key_string =\ + self._get_effective_partition_key_string( + cast(List[Union[str, int, float, bool, _Empty, _Undefined]], [pk_value])) return _Range(effective_partition_key_string, effective_partition_key_string, True, True) def _get_effective_partition_key_for_hash_partitioning(self) -> str: @@ -190,7 +193,7 @@ def _get_effective_partition_key_for_hash_partitioning(self) -> str: def _get_effective_partition_key_string( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] ) -> Union[int, str]: if not pk_value: return _MinimumInclusiveEffectivePartitionKey @@ -278,7 +281,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: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> bool: # pylint: disable=line-too-long if self.kind!= "MultiHash": return False if isinstance(partition_key, list) and len(self.path) == len(partition_key): diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index a1d34262cae7..4b286d2b82f8 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -291,5 +291,29 @@ def test_query_change_feed_with_split(self, setup): assert actual_ids == expected_ids setup["created_db"].delete_container(created_collection.id) + def test_query_change_feed_with_multi_partition(self, setup): + created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=11000) + + # create one doc and make sure change feed query can return the document + new_documents = [ + {'pk': 'pk', 'id': 'doc1'}, + {'pk': 'pk2', 'id': 'doc2'}, + {'pk': 'pk3', 'id': 'doc3'}, + {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc1', 'doc2', 'doc3', 'doc4'] + + for document in new_documents: + created_collection.create_item(body=document) + + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") + it = query_iterable.__iter__() + actual_ids = [] + for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + if __name__ == "__main__": unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index b4165af0601c..886c1ffc1bcc 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -318,5 +318,29 @@ async def test_query_change_feed_with_split_async(self, setup): assert actual_ids == expected_ids setup["created_db"].delete_container(created_collection.id) + async def test_query_change_feed_with_multi_partition_async(self, setup): + created_collection = await setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=11000) + + # create one doc and make sure change feed query can return the document + new_documents = [ + {'pk': 'pk', 'id': 'doc1'}, + {'pk': 'pk2', 'id': 'doc2'}, + {'pk': 'pk3', 'id': 'doc3'}, + {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc1', 'doc2', 'doc3', 'doc4'] + + for document in new_documents: + await created_collection.create_item(body=document) + + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") + it = query_iterable.__aiter__() + actual_ids = [] + async for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + if __name__ == '__main__': unittest.main() From 56bbb9e6d23ad8b791a38d6472c6576e73a8512b Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 20 Aug 2024 22:41:01 -0700 Subject: [PATCH 09/59] added functionality for merging session tokens from logical pk --- .../azure/cosmos/_cosmos_client_connection.py | 28 +++++++++ .../azure/cosmos/_session_token_helpers.py | 3 + .../azure/cosmos/_vector_session_token.py | 5 ++ .../azure-cosmos/azure/cosmos/container.py | 7 ++- .../azure-cosmos/test/test_session_helpers.py | 57 +++++++++++++++++++ 5 files changed, 97 insertions(+), 3 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_session_helpers.py 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 cae20bc67ee7..a5ce5835dba9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -26,9 +26,12 @@ import os import urllib.parse from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast, Type + +from pycparser.ply.yacc import token from typing_extensions import TypedDict from urllib3.util.retry import Retry +from ._vector_session_token import VectorSessionToken from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged from azure.core import PipelineClient @@ -3290,3 +3293,28 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s partition_key_definition = collection.get("partitionKey") self.partition_key_definition_cache[collection_link] = partition_key_definition return partition_key_definition + + def MergeSessionTokens(self, pks_to_session_tokens, target_pk): + comparison_session_token = '' + comparison_pk_range_id = '' + for (pk, session_token) in pks_to_session_tokens: + if pk == target_pk: + token_pairs = session_token.split(",") + comparison_pk_range_id = token_pairs[0] + comparison_session_token = VectorSessionToken.create(token_pairs[1]) + break + for (pk, session_token) in pks_to_session_tokens: + token_pairs = session_token.split(",") + pk_range_id = token_pairs[0] + vector_session_token = VectorSessionToken.create(token_pairs[0]) + # This should not be necessary + # if pk_range_id == comparison_pk_range_id: + # comparison_session_token = comparison_session_token.merge(vector_session_token) + if pk == target_pk: + if pk_range_id == comparison_pk_range_id: + comparison_session_token = comparison_session_token.merge(vector_session_token) + elif session_token.is_greater(comparison_session_token): + comparison_pk_range_id = pk_range_id + comparison_session_token = session_token + return comparison_session_token + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py new file mode 100644 index 000000000000..b28b04f64312 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -0,0 +1,3 @@ + + + diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py index f8285f000b67..a5be01799447 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py @@ -113,6 +113,11 @@ def equals(self, other): and self.are_region_progress_equal(other.local_lsn_by_region) ) + def is_greater(self, other): + if self.global_lsn > other.global_lsn: + return True + return False + def merge(self, other): if other is None: raise ValueError("Invalid Session Token (should not be None)") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 7c5e4cbf6098..c29f138dcd64 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1150,9 +1150,10 @@ def mergeSessionTokens(self, session token. :param pks_to_session_tokens: list of partition key and session token tuples. - :type pks_to_session_tokens: List[Tuple(string, PartitionKey)] + :type pks_to_session_tokens: List[Tuple(str, PartitionKey)] :param target_pk: partition key to get most up to date session token. :type target_pk: PartitionKey - :rtype: string + :returns: a session token + :rtype: str """ - pass \ No newline at end of file + self.client_connection.MergeSessionTokens(pks_to_session_tokens, target_pk) \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/test/test_session_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_helpers.py new file mode 100644 index 000000000000..dc9d2347bf34 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_session_helpers.py @@ -0,0 +1,57 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest + +import azure.cosmos._synchronized_request as synchronized_request +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +import test_config +from azure.cosmos import DatabaseProxy +from azure.cosmos import _retry_utility +from azure.cosmos.http_constants import StatusCodes, SubStatusCodes, HttpHeaders + + +@pytest.mark.cosmosEmulator +class TestSession(unittest.TestCase): + """Test to ensure escaping of non-ascii characters from partition key""" + + created_db: DatabaseProxy = None + client: cosmos_client.CosmosClient = None + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + connectionPolicy = test_config.TestConfig.connectionPolicy + configs = test_config.TestConfig + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + TEST_COLLECTION_ID = configs.TEST_MULTI_PARTITION_CONTAINER_ID + + @classmethod + def setUpClass(cls): + # creates the database, collection, and insert all the documents + # we will gain some speed up in running the tests by creating the + # database, collection and inserting all the docs only once + + if cls.masterKey == '[YOUR_KEY_HERE]' or cls.host == '[YOUR_ENDPOINT_HERE]': + raise Exception("You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + + cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) + cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) + cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID) + + # have a test for split, merge, several pkrangeIds from different logical pks, and for normal case for several session tokens + def test_session_token_not_sent_for_master_resource_ops(self): + pass + + + + + + + +if __name__ == '__main__': + unittest.main() From 8c0aa4604015588c0e39af42c7bddde69a2942db Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 21 Aug 2024 09:41:54 -0700 Subject: [PATCH 10/59] fix mypy --- .../azure-cosmos/azure/cosmos/container.py | 16 ++++++++++------ .../azure-cosmos/azure/cosmos/partition_key.py | 4 ++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index b04972bdee6c..cc42b6f87411 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -132,7 +132,9 @@ def _set_partition_key( 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) - def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: + def _get_epk_range_for_partition_key( + self, + partition_key_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> Range: # pylint: disable=line-too-long container_properties = self._get_properties() partition_key_definition = container_properties["partitionKey"] partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) @@ -583,13 +585,14 @@ def query_items( # pylint:disable=docstring-missing-param if populate_index_metrics is not None: feed_options["populateIndexMetrics"] = populate_index_metrics if partition_key is not None: - if self.__is_prefix_partitionkey(partition_key): + partition_key_value = self._set_partition_key(partition_key) + if self.__is_prefix_partitionkey(partition_key_value): kwargs["isPrefixPartitionQuery"] = True properties = self._get_properties() kwargs["partitionKeyDefinition"] = properties["partitionKey"] - kwargs["partitionKeyDefinition"]["partition_key"] = self._set_partition_key(partition_key) + kwargs["partitionKeyDefinition"]["partition_key"] = partition_key_value else: - feed_options["partitionKey"] = self._set_partition_key(partition_key) + 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: @@ -616,8 +619,9 @@ def query_items( # pylint:disable=docstring-missing-param return items def __is_prefix_partitionkey( - self, partition_key: PartitionKeyType - ) -> bool: + self, + partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> bool: # pylint: disable=line-too-long + properties = self._get_properties() pk_properties = properties["partitionKey"] partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 881cce0895e1..e4d659a08fac 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -238,7 +238,7 @@ def _write_for_hashing_v2( def _get_effective_partition_key_for_hash_partitioning_v2( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] ) -> str: with BytesIO() as ms: for component in pk_value: @@ -257,7 +257,7 @@ def _get_effective_partition_key_for_hash_partitioning_v2( def _get_effective_partition_key_for_multi_hash_partitioning_v2( self, - pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] + pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] ) -> str: sb = [] for value in pk_value: From 28394b901628687bfe7cd1c8c8e0b1352342b977 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 21 Aug 2024 15:41:21 -0700 Subject: [PATCH 11/59] added tests for basic merge and split --- .../azure/cosmos/_cosmos_client_connection.py | 9 ++- .../azure-cosmos/azure/cosmos/container.py | 10 ++-- .../azure-cosmos/test/test_session_helpers.py | 57 ------------------- .../test/test_session_token_helpers.py | 0 4 files changed, 9 insertions(+), 67 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/test/test_session_helpers.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py 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 a5ce5835dba9..61b42ec4b96e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -27,7 +27,6 @@ import urllib.parse from typing import Callable, Dict, Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union, cast, Type -from pycparser.ply.yacc import token from typing_extensions import TypedDict from urllib3.util.retry import Retry @@ -3299,14 +3298,14 @@ def MergeSessionTokens(self, pks_to_session_tokens, target_pk): comparison_pk_range_id = '' for (pk, session_token) in pks_to_session_tokens: if pk == target_pk: - token_pairs = session_token.split(",") + token_pairs = session_token.split(":") comparison_pk_range_id = token_pairs[0] comparison_session_token = VectorSessionToken.create(token_pairs[1]) break for (pk, session_token) in pks_to_session_tokens: - token_pairs = session_token.split(",") + token_pairs = session_token.split(":") pk_range_id = token_pairs[0] - vector_session_token = VectorSessionToken.create(token_pairs[0]) + vector_session_token = VectorSessionToken.create(token_pairs[1]) # This should not be necessary # if pk_range_id == comparison_pk_range_id: # comparison_session_token = comparison_session_token.merge(vector_session_token) @@ -3316,5 +3315,5 @@ def MergeSessionTokens(self, pks_to_session_tokens, target_pk): elif session_token.is_greater(comparison_session_token): comparison_pk_range_id = pk_range_id comparison_session_token = session_token - return comparison_session_token + return comparison_pk_range_id + ":" + comparison_session_token.session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index c29f138dcd64..13038c991cbd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1141,10 +1141,10 @@ def delete_all_items_by_partition_key( self.client_connection.DeleteAllItemsByPartitionKey( collection_link=self.container_link, options=request_options, **kwargs) - def mergeSessionTokens(self, - pks_to_session_tokens: List, - target_pk: PartitionKey - ) -> "Session Token": + def merge_session_tokens(self, + pks_to_session_tokens: List, + target_pk: PartitionKey + ) -> str: """Merge session tokens from different clients to figure out which is the most up to date for a specific partition key. This should only be used if maintaining own session token or else the sdk will keep track of session token. @@ -1156,4 +1156,4 @@ def mergeSessionTokens(self, :returns: a session token :rtype: str """ - self.client_connection.MergeSessionTokens(pks_to_session_tokens, target_pk) \ No newline at end of file + return self.client_connection.MergeSessionTokens(pks_to_session_tokens, target_pk) \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/test/test_session_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_helpers.py deleted file mode 100644 index dc9d2347bf34..000000000000 --- a/sdk/cosmos/azure-cosmos/test/test_session_helpers.py +++ /dev/null @@ -1,57 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) Microsoft Corporation. All rights reserved. - -import unittest -import uuid - -import pytest - -import azure.cosmos._synchronized_request as synchronized_request -import azure.cosmos.cosmos_client as cosmos_client -import azure.cosmos.exceptions as exceptions -import test_config -from azure.cosmos import DatabaseProxy -from azure.cosmos import _retry_utility -from azure.cosmos.http_constants import StatusCodes, SubStatusCodes, HttpHeaders - - -@pytest.mark.cosmosEmulator -class TestSession(unittest.TestCase): - """Test to ensure escaping of non-ascii characters from partition key""" - - created_db: DatabaseProxy = None - client: cosmos_client.CosmosClient = None - host = test_config.TestConfig.host - masterKey = test_config.TestConfig.masterKey - connectionPolicy = test_config.TestConfig.connectionPolicy - configs = test_config.TestConfig - TEST_DATABASE_ID = configs.TEST_DATABASE_ID - TEST_COLLECTION_ID = configs.TEST_MULTI_PARTITION_CONTAINER_ID - - @classmethod - def setUpClass(cls): - # creates the database, collection, and insert all the documents - # we will gain some speed up in running the tests by creating the - # database, collection and inserting all the docs only once - - if cls.masterKey == '[YOUR_KEY_HERE]' or cls.host == '[YOUR_ENDPOINT_HERE]': - raise Exception("You must specify your Azure Cosmos account values for " - "'masterKey' and 'host' at the top of this class to run the " - "tests.") - - cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) - cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) - cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID) - - # have a test for split, merge, several pkrangeIds from different logical pks, and for normal case for several session tokens - def test_session_token_not_sent_for_master_resource_ops(self): - pass - - - - - - - -if __name__ == '__main__': - unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py new file mode 100644 index 000000000000..e69de29bb2d1 From 25c3363fc03c0f30eb2ab3820548eba720d44860 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 27 Aug 2024 13:53:26 -0700 Subject: [PATCH 12/59] resolve comments --- .../_change_feed/aio/change_feed_iterable.py | 13 +++++++------ .../cosmos/_change_feed/change_feed_iterable.py | 10 +++++----- .../cosmos/_change_feed/change_feed_state.py | 13 ++++++------- .../azure/cosmos/_change_feed/feed_range.py | 5 ++--- .../feed_range_composite_continuation_token.py | 6 ++---- sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py | 3 --- .../azure-cosmos/azure/cosmos/aio/_container.py | 17 ++++++++--------- .../azure-cosmos/azure/cosmos/container.py | 17 ++++++++--------- 8 files changed, 38 insertions(+), 46 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 83c12f59157c..cae3bc5c9bf7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -28,7 +28,8 @@ from azure.cosmos import PartitionKey from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateV1 -from azure.cosmos._utils import is_base64_encoded, is_key_exists_and_not_none +from azure.cosmos._utils import is_base64_encoded + # pylint: disable=protected-access @@ -63,7 +64,7 @@ def __init__( self._collection_link = collection_link self._change_feed_fetcher = None - if not is_key_exists_and_not_none(self._options, "changeFeedStateContext"): + if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") change_feed_state_context = self._options.pop("changeFeedStateContext") @@ -119,7 +120,7 @@ async def _fetch_next(self, *args): # pylint: disable=unused-argument async def _initialize_change_feed_fetcher(self): change_feed_state_context = self._options.pop("changeFeedStateContext") conn_properties = await self._options.pop("containerProperties") - if is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): + if change_feed_state_context.get("partitionKey"): change_feed_state_context["partitionKey"] = await change_feed_state_context.pop("partitionKey") pk_properties = conn_properties.get("partitionKey") partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) @@ -147,11 +148,11 @@ async def _initialize_change_feed_fetcher(self): def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: - if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + if change_feed_state_context.get("continuationPkRangeId"): # if continuation token is in v1 format, throw exception if feed_range is set - if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + if change_feed_state_context.get("feedRange"): raise ValueError("feed_range and continuation are incompatible") - elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + elif change_feed_state_context.get("continuationFeedRange"): # if continuation token is in v2 format, since the token itself contains the full change feed state # so we will ignore other parameters (including incompatible parameters) if they passed in pass diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 4b03e33d0122..7fc62684d9ff 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -27,7 +27,7 @@ from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedState -from azure.cosmos._utils import is_base64_encoded, is_key_exists_and_not_none +from azure.cosmos._utils import is_base64_encoded class ChangeFeedIterable(PageIterator): @@ -60,7 +60,7 @@ def __init__( self._collection_link = collection_link self._change_feed_fetcher = None - if not is_key_exists_and_not_none(self._options, "changeFeedStateContext"): + if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") change_feed_state_context = self._options.pop("changeFeedStateContext") @@ -140,11 +140,11 @@ def _initialize_change_feed_fetcher(self): def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: - if is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId"): + if change_feed_state_context.get("continuationPkRangeId"): # if continuation token is in v1 format, throw exception if feed_range is set - if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + if change_feed_state_context.get("feedRange"): raise ValueError("feed_range and continuation are incompatible") - elif is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + elif change_feed_state_context.get("continuationFeedRange"): # if continuation token is in v2 format, since the token itself contains the full change feed state # so we will ignore other parameters (including incompatible parameters) if they passed in pass diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 988932ff3af6..742c4891bfdf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -38,7 +38,6 @@ from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range -from azure.cosmos._utils import is_key_exists_and_not_none from azure.cosmos.exceptions import CosmosFeedRangeGoneError from azure.cosmos.partition_key import _Empty, _Undefined @@ -74,11 +73,11 @@ def from_json( container_rid: str, change_feed_state_context: Dict[str, Any]): - if (is_key_exists_and_not_none(change_feed_state_context, "partitionKeyRangeId") - or is_key_exists_and_not_none(change_feed_state_context, "continuationPkRangeId")): + if (change_feed_state_context.get("partitionKeyRangeId") + or change_feed_state_context.get("continuationPkRangeId")): return ChangeFeedStateV1.from_json(container_link, container_rid, change_feed_state_context) - if is_key_exists_and_not_none(change_feed_state_context, "continuationFeedRange"): + if change_feed_state_context.get("continuationFeedRange"): # get changeFeedState from continuation continuation_json_str = base64.b64decode(change_feed_state_context["continuationFeedRange"]).decode( 'utf-8') @@ -376,12 +375,12 @@ def from_initial_state( change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV2': feed_range: Optional[FeedRange] = None - if is_key_exists_and_not_none(change_feed_state_context, "feedRange"): + if change_feed_state_context.get("feedRange"): feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') feed_range_json = json.loads(feed_range_str) feed_range = FeedRangeEpk(Range.ParseFromDict(feed_range_json)) - elif is_key_exists_and_not_none(change_feed_state_context, "partitionKey"): - if is_key_exists_and_not_none(change_feed_state_context, "partitionKeyFeedRange"): + elif change_feed_state_context.get("partitionKey"): + if change_feed_state_context.get("partitionKeyFeedRange"): feed_range =\ FeedRangePartitionKey( change_feed_state_context["partitionKey"], diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index 856ccd6c5b48..481496159cf3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -26,7 +26,6 @@ from typing import Union, List, Dict, Any from azure.cosmos._routing.routing_range import Range -from azure.cosmos._utils import is_key_exists_and_not_none from azure.cosmos.partition_key import _Undefined, _Empty @@ -71,7 +70,7 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartitionKey': - if is_key_exists_and_not_none(data, cls.type_property_name): + if data.get(cls.type_property_name): pk_value = data.get(cls.type_property_name) if not pk_value: return cls(_Empty(), feed_range) @@ -104,7 +103,7 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeEpk': - if is_key_exists_and_not_none(data, cls.type_property_name): + if data.get(cls.type_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) return cls(feed_range) raise ValueError(f"Can not parse FeedRangeEPK from the json, there is no property {cls.type_property_name}") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index 2f73af4e5c47..e8bfe60ced3f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -30,8 +30,6 @@ from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range -from azure.cosmos._utils import is_key_exists_and_not_none - class FeedRangeCompositeContinuation(object): _version_property_name = "v" @@ -90,9 +88,9 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation': # parsing feed range feed_range: Optional[FeedRange] = None - if is_key_exists_and_not_none(data, FeedRangeEpk.type_property_name): + if data.get(FeedRangeEpk.type_property_name): feed_range = FeedRangeEpk.from_json(data) - elif is_key_exists_and_not_none(data, FeedRangePartitionKey.type_property_name): + elif data.get(FeedRangePartitionKey.type_property_name): feed_range = FeedRangePartitionKey.from_json(data, continuation[0].feed_range) else: raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 7690383d375a..6e3a8c67fcfe 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -79,6 +79,3 @@ def is_base64_encoded(data: str) -> bool: return True except (json.JSONDecodeError, ValueError): return False - -def is_key_exists_and_not_none(data: Dict[str, Any], key: str) -> bool: - return key in data and data[key] is not None diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 5cae847872d5..1e86ccb1aa44 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -43,7 +43,6 @@ ) from .._routing import routing_range from .._routing.routing_range import Range -from .._utils import is_key_exists_and_not_none from ..offer import ThroughputProperties from ..partition_key import ( NonePartitionKeyValue, @@ -573,13 +572,13 @@ def query_items_change_feed( **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - if is_key_exists_and_not_none(kwargs, "priority"): + if kwargs.get("priority"): kwargs['priority'] = kwargs['priority'] feed_options = _build_options(kwargs) change_feed_state_context = {} # Back compatibility with deprecation warnings for partition_key_range_id - if (args and args[0] is not None) or is_key_exists_and_not_none(kwargs, "partition_key_range_id"): + if (args and args[0] is not None) or kwargs.get("partition_key_range_id"): warnings.warn( "partition_key_range_id is deprecated. Please pass in feed_range instead.", DeprecationWarning @@ -591,7 +590,7 @@ def query_items_change_feed( change_feed_state_context['partitionKeyRangeId'] = args[0] # Back compatibility with deprecation warnings for is_start_from_beginning - if (len(args) >= 2 and args[1] is not None) or is_key_exists_and_not_none(kwargs, "is_start_from_beginning"): + if (len(args) >= 2 and args[1] is not None) or kwargs.get("is_start_from_beginning"): warnings.warn( "is_start_from_beginning is deprecated. Please pass in start_time instead.", DeprecationWarning @@ -606,7 +605,7 @@ def query_items_change_feed( change_feed_state_context["startTime"] = "Beginning" # parse start_time - if is_key_exists_and_not_none(kwargs, "start_time"): + if kwargs.get("start_time"): if change_feed_state_context.get("startTime") is not None: raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") @@ -617,23 +616,23 @@ def query_items_change_feed( change_feed_state_context["startTime"] = start_time # parse continuation token - if len(args) >= 3 and args[2] is not None or is_key_exists_and_not_none(feed_options, "continuation"): + if len(args) >= 3 and args[2] is not None or feed_options.get("continuation"): try: continuation = feed_options.pop('continuation') except KeyError: continuation = args[2] change_feed_state_context["continuation"] = continuation - if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): + if len(args) >= 4 and args[3] is not None or kwargs.get("max_item_count"): try: feed_options["maxItemCount"] = kwargs.pop('max_item_count') except KeyError: feed_options["maxItemCount"] = args[3] - if is_key_exists_and_not_none(kwargs, "partition_key"): + if kwargs.get("partition_key"): change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop("partition_key")) - if is_key_exists_and_not_none(kwargs, "feed_range"): + if kwargs.get("feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') feed_options["containerProperties"] = self._get_properties() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index cc42b6f87411..5e29302bcaae 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -41,7 +41,6 @@ from ._cosmos_client_connection import CosmosClientConnection from ._routing import routing_range from ._routing.routing_range import Range -from ._utils import is_key_exists_and_not_none from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -414,13 +413,13 @@ def query_items_change_feed( **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - if is_key_exists_and_not_none(kwargs, "priority"): + if kwargs.get("priority"): kwargs['priority'] = kwargs['priority'] feed_options = build_options(kwargs) change_feed_state_context = {} # Back compatibility with deprecation warnings for partition_key_range_id - if (args and args[0] is not None) or is_key_exists_and_not_none(kwargs, "partition_key_range_id"): + if (args and args[0] is not None) or kwargs.get("partition_key_range_id"): warnings.warn( "partition_key_range_id is deprecated. Please pass in feed_range instead.", DeprecationWarning @@ -432,7 +431,7 @@ def query_items_change_feed( change_feed_state_context['partitionKeyRangeId'] = args[0] # Back compatibility with deprecation warnings for is_start_from_beginning - if (len(args) >= 2 and args[1] is not None) or is_key_exists_and_not_none(kwargs, "is_start_from_beginning"): + if (len(args) >= 2 and args[1] is not None) or kwargs.get("is_start_from_beginning"): warnings.warn( "is_start_from_beginning is deprecated. Please pass in start_time instead.", DeprecationWarning @@ -447,7 +446,7 @@ def query_items_change_feed( change_feed_state_context["startTime"] = "Beginning" # parse start_time - if is_key_exists_and_not_none(kwargs, "start_time"): + if kwargs.get("start_time"): if change_feed_state_context.get("startTime") is not None: raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") @@ -458,25 +457,25 @@ def query_items_change_feed( change_feed_state_context["startTime"] = start_time # parse continuation token - if len(args) >= 3 and args[2] is not None or is_key_exists_and_not_none(feed_options, "continuation"): + if len(args) >= 3 and args[2] is not None or feed_options.get("continuation"): try: continuation = feed_options.pop('continuation') except KeyError: continuation = args[2] change_feed_state_context["continuation"] = continuation - if len(args) >= 4 and args[3] is not None or is_key_exists_and_not_none(kwargs, "max_item_count"): + if len(args) >= 4 and args[3] is not None or kwargs.get("max_item_count"): try: feed_options["maxItemCount"] = kwargs.pop('max_item_count') except KeyError: feed_options["maxItemCount"] = args[3] - if is_key_exists_and_not_none(kwargs, "partition_key"): + if kwargs.get("partition_key"): change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop('partition_key')) change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) - if is_key_exists_and_not_none(kwargs, "feed_range"): + if kwargs.get("feed_range"): change_feed_state_context["feedRange"] = kwargs.pop('feed_range') container_properties = self._get_properties() From cecdfa5ed2862d8ec14b4a14ae0b7c828fa53a37 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 27 Aug 2024 17:11:08 -0700 Subject: [PATCH 13/59] resolve comments --- .../_change_feed/aio/change_feed_fetcher.py | 52 ++++----- .../_change_feed/aio/change_feed_iterable.py | 44 ++++--- .../_change_feed/change_feed_fetcher.py | 46 ++++---- .../_change_feed/change_feed_iterable.py | 31 +++-- .../_change_feed/change_feed_start_from.py | 13 ++- .../cosmos/_change_feed/change_feed_state.py | 17 ++- .../composite_continuation_token.py | 20 ++-- .../azure/cosmos/_change_feed/feed_range.py | 4 +- ...feed_range_composite_continuation_token.py | 6 +- .../azure/cosmos/aio/_container.py | 110 ++++++++++-------- .../azure-cosmos/azure/cosmos/container.py | 96 +++++++++------ .../azure-cosmos/azure/cosmos/exceptions.py | 2 +- sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst | 8 ++ 13 files changed, 255 insertions(+), 194 deletions(-) create mode 100644 sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index ee926aa0e92c..90aa2d01adfa 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -23,14 +23,13 @@ database service. """ import base64 -import copy import json from abc import ABC, abstractmethod -from typing import Dict, Any, List +from typing import Dict, Any, List, Callable, Tuple, Awaitable from azure.cosmos import http_constants, exceptions -from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromPointInTime -from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 +from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromType +from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV2, ChangeFeedStateVersion from azure.cosmos.aio import _retry_utility_async from azure.cosmos.exceptions import CosmosHttpResponseError @@ -39,7 +38,7 @@ class ChangeFeedFetcher(ABC): @abstractmethod - async def fetch_next_block(self): + async def fetch_next_block(self) -> List[Dict[str, Any]]: pass class ChangeFeedFetcherV1(ChangeFeedFetcher): @@ -53,38 +52,38 @@ def __init__( client, resource_link: str, feed_options: Dict[str, Any], - fetch_function): + fetch_function: Callable[[Dict[str, Any]], Awaitable[Tuple[List[Dict[str, Any]], Dict[str, Any]]]] + ) -> None: self._client = client self._feed_options = feed_options self._change_feed_state = self._feed_options.pop("changeFeedState") - if not isinstance(self._change_feed_state, ChangeFeedStateV1): + if self._change_feed_state.version != ChangeFeedStateVersion.V1: raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version" f" {type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function - async def fetch_next_block(self): + async def fetch_next_block(self) -> List[Dict[str, Any]]: """Returns a block of results. :return: List of results. :rtype: list """ async def callback(): - return await self.fetch_change_feed_items(self._fetch_function) + return await self.fetch_change_feed_items() return await _retry_utility_async.ExecuteAsync(self._client, self._client._global_endpoint_manager, callback) - async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: - new_options = copy.deepcopy(self._feed_options) - new_options["changeFeedState"] = self._change_feed_state + async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: + self._feed_options["changeFeedState"] = self._change_feed_state - self._change_feed_state.populate_feed_options(new_options) + self._change_feed_state.populate_feed_options(self._feed_options) is_s_time_first_fetch = self._change_feed_state._continuation is None while True: - (fetched_items, response_headers) = await fetch_function(new_options) + (fetched_items, response_headers) = await self._fetch_function(self._feed_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. @@ -96,7 +95,7 @@ async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: # When processing from point in time, there will be no initial results being returned, # so we will retry with the new continuation token again - if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): is_s_time_first_fetch = False else: @@ -113,20 +112,21 @@ def __init__( client, resource_link: str, feed_options: Dict[str, Any], - fetch_function): + fetch_function: Callable[[Dict[str, Any]], Awaitable[Tuple[List[Dict[str, Any]], Dict[str, Any]]]] + ) -> None: self._client = client self._feed_options = feed_options self._change_feed_state: ChangeFeedStateV2 = self._feed_options.pop("changeFeedState") - if not isinstance(self._change_feed_state, ChangeFeedStateV2): + if self._change_feed_state.version != ChangeFeedStateVersion.V2: raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version " - f"{type(self._change_feed_state)}") + f"{type(self._change_feed_state.version)}") self._resource_link = resource_link self._fetch_function = fetch_function - async def fetch_next_block(self): + async def fetch_next_block(self) -> List[Dict[str, Any]]: """Returns a block of results. :return: List of results. @@ -134,7 +134,7 @@ async def fetch_next_block(self): """ async def callback(): - return await self.fetch_change_feed_items(self._fetch_function) + return await self.fetch_change_feed_items() try: return await _retry_utility_async.ExecuteAsync( @@ -152,15 +152,14 @@ async def callback(): return await self.fetch_next_block() - async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: - new_options = copy.deepcopy(self._feed_options) - new_options["changeFeedState"] = self._change_feed_state + async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: + self._feed_options["changeFeedState"] = self._change_feed_state - self._change_feed_state.populate_feed_options(new_options) + self._change_feed_state.populate_feed_options(self._feed_options) is_s_time_first_fetch = True while True: - (fetched_items, response_headers) = await fetch_function(new_options) + (fetched_items, response_headers) = await self._fetch_function(self._feed_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether @@ -180,8 +179,7 @@ async def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) - #TODO: can this part logic be simplified - if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): response_headers[continuation_key] = self._get_base64_encoded_continuation() is_s_time_first_fetch = False diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index cae3bc5c9bf7..745ed19279c7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,13 +21,13 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any +from typing import Dict, Any, Optional, Callable, Coroutine, Tuple, List, AsyncIterator from azure.core.async_paging import AsyncPageIterator from azure.cosmos import PartitionKey from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 -from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateV1 +from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion from azure.cosmos._utils import is_base64_encoded @@ -42,21 +42,20 @@ class ChangeFeedIterable(AsyncPageIterator): def __init__( self, client, - options, - fetch_function=None, - collection_link=None, - continuation_token=None, - ): + options: Dict[str, Any], + fetch_function=Optional[Callable[[Dict[str, Any]], Coroutine[Tuple[List[Dict[str, Any]], Dict[str, Any]]]]], + collection_link=Optional[str], + continuation_token=Optional[str], + ) -> None: """Instantiates a ChangeFeedIterable for non-client side partitioning queries. - ChangeFeedFetcher will be used as the internal query execution - context. - - :param CosmosClient client: Instance of document client. - :param dict options: The request options for the request. - :param method fetch_function: - + :param CosmosClient client: Instance of document client. + :param dict options: The request options for the request. + :param fetch_function: The fetch function. + :param collection_link: The collection resource link. + :param continuation_token: The continuation token passed in from by_page """ + self._client = client self.retry_options = client.connection_policy.RetryOptions self._options = options @@ -90,7 +89,7 @@ def __init__( super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) - async def _unpack(self, block): + async def _unpack(self, block) -> Tuple[str, AsyncIterator[List[Dict[str, Any]]]]: continuation = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get('etag') @@ -99,12 +98,9 @@ async def _unpack(self, block): self._did_a_call_already = False return continuation, block - async def _fetch_next(self, *args): # pylint: disable=unused-argument + async def _fetch_next(self, *args) -> List[Dict[str, Any]]: # pylint: disable=unused-argument """Return a block of results with respecting retry policy. - This method only exists for backward compatibility reasons. (Because - QueryIterable has exposed fetch_next_block api). - :param Any args: :return: List of results. :rtype: list @@ -117,7 +113,7 @@ async def _fetch_next(self, *args): # pylint: disable=unused-argument raise StopAsyncIteration return block - async def _initialize_change_feed_fetcher(self): + async def _initialize_change_feed_fetcher(self) -> None: change_feed_state_context = self._options.pop("changeFeedStateContext") conn_properties = await self._options.pop("containerProperties") if change_feed_state_context.get("partitionKey"): @@ -131,7 +127,7 @@ async def _initialize_change_feed_fetcher(self): ChangeFeedState.from_json(self._collection_link, conn_properties["_rid"], change_feed_state_context) self._options["changeFeedState"] = change_feed_state - if isinstance(change_feed_state, ChangeFeedStateV1): + if change_feed_state.version != ChangeFeedStateVersion.V1: self._change_feed_fetcher = ChangeFeedFetcherV1( self._client, self._collection_link, @@ -148,11 +144,11 @@ async def _initialize_change_feed_fetcher(self): def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: - if change_feed_state_context.get("continuationPkRangeId"): + if change_feed_state_context.get("continuationPkRangeId") is not None: # if continuation token is in v1 format, throw exception if feed_range is set - if change_feed_state_context.get("feedRange"): + if change_feed_state_context.get("feedRange") is not None: raise ValueError("feed_range and continuation are incompatible") - elif change_feed_state_context.get("continuationFeedRange"): + elif change_feed_state_context.get("continuationFeedRange") is not None: # if continuation token is in v2 format, since the token itself contains the full change feed state # so we will ignore other parameters (including incompatible parameters) if they passed in pass diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index 92f0b2446f74..35ae9a15a08a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -23,14 +23,13 @@ database service. """ import base64 -import copy import json from abc import ABC, abstractmethod -from typing import Dict, Any, List +from typing import Dict, Any, List, Callable, Tuple from azure.cosmos import _retry_utility, http_constants, exceptions -from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromPointInTime -from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2 +from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromType +from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedStateV2, ChangeFeedStateVersion from azure.cosmos.exceptions import CosmosHttpResponseError # pylint: disable=protected-access @@ -52,38 +51,38 @@ def __init__( client, resource_link: str, feed_options: Dict[str, Any], - fetch_function): + fetch_function: Callable[[Dict[str, Any]], Tuple[List[Dict[str, Any]], Dict[str, Any]]] + ) -> None: self._client = client self._feed_options = feed_options self._change_feed_state: ChangeFeedStateV1 = self._feed_options.pop("changeFeedState") - if not isinstance(self._change_feed_state, ChangeFeedStateV1): + if self._change_feed_state.version != ChangeFeedStateVersion.V1: raise ValueError(f"ChangeFeedFetcherV1 can not handle change feed state version" f" {type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function - def fetch_next_block(self): + def fetch_next_block(self) -> List[Dict[str, Any]]: """Returns a block of results. :return: List of results. :rtype: list """ def callback(): - return self.fetch_change_feed_items(self._fetch_function) + return self.fetch_change_feed_items() return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback) - def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: - new_options = copy.deepcopy(self._feed_options) - new_options["changeFeedState"] = self._change_feed_state + def fetch_change_feed_items(self) -> List[Dict[str, Any]]: + self._feed_options["changeFeedState"] = self._change_feed_state - self._change_feed_state.populate_feed_options(new_options) + self._change_feed_state.populate_feed_options(self._feed_options) is_s_time_first_fetch = self._change_feed_state._continuation is None while True: - (fetched_items, response_headers) = fetch_function(new_options) + (fetched_items, response_headers) = self._fetch_function(self._feed_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. @@ -95,7 +94,7 @@ def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: # When processing from point in time, there will be no initial results being returned, # so we will retry with the new continuation token again - if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): is_s_time_first_fetch = False else: @@ -112,20 +111,20 @@ def __init__( client, resource_link: str, feed_options: Dict[str, Any], - fetch_function): + fetch_function: Callable[[Dict[str, Any]], Tuple[List[Dict[str, Any]], Dict[str, Any]]]): self._client = client self._feed_options = feed_options self._change_feed_state: ChangeFeedStateV2 = self._feed_options.pop("changeFeedState") - if not isinstance(self._change_feed_state, ChangeFeedStateV2): + if self._change_feed_state.version != ChangeFeedStateVersion.V2: raise ValueError(f"ChangeFeedFetcherV2 can not handle change feed state version " f"{type(self._change_feed_state)}") self._resource_link = resource_link self._fetch_function = fetch_function - def fetch_next_block(self): + def fetch_next_block(self) -> List[Dict[str, Any]]: """Returns a block of results. :return: List of results. @@ -133,7 +132,7 @@ def fetch_next_block(self): """ def callback(): - return self.fetch_change_feed_items(self._fetch_function) + return self.fetch_change_feed_items() try: return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback) @@ -146,15 +145,14 @@ def callback(): return self.fetch_next_block() - def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: - new_options = copy.deepcopy(self._feed_options) - new_options["changeFeedState"] = self._change_feed_state + def fetch_change_feed_items(self) -> List[Dict[str, Any]]: + self._feed_options["changeFeedState"] = self._change_feed_state - self._change_feed_state.populate_feed_options(new_options) + self._change_feed_state.populate_feed_options(self._feed_options) is_s_time_first_fetch = self._change_feed_state._continuation.current_token.token is None while True: - (fetched_items, response_headers) = fetch_function(new_options) + (fetched_items, response_headers) = self._fetch_function(self._feed_options) continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. @@ -174,7 +172,7 @@ def fetch_change_feed_items(self, fetch_function) -> List[Dict[str, Any]]: self._change_feed_state.apply_server_response_continuation( response_headers.get(continuation_key)) - if (isinstance(self._change_feed_state._change_feed_start_from, ChangeFeedStartFromPointInTime) + if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): response_headers[continuation_key] = self._get_base64_encoded_continuation() is_s_time_first_fetch = False diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 7fc62684d9ff..6eaaa31bfd8f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -21,12 +21,12 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any +from typing import Dict, Any, Tuple, List, Optional, Callable from azure.core.paging import PageIterator from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 -from azure.cosmos._change_feed.change_feed_state import ChangeFeedStateV1, ChangeFeedState +from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion from azure.cosmos._utils import is_base64_encoded @@ -39,11 +39,11 @@ class ChangeFeedIterable(PageIterator): def __init__( self, client, - options, - fetch_function=None, - collection_link=None, - continuation_token=None, - ): + options: Dict[str, Any], + fetch_function=Optional[Callable[[Dict[str, Any]], Tuple[List[Dict[str, Any]], Dict[str, Any]]]], + collection_link=Optional[str], + continuation_token=Optional[str], + ) -> None: """Instantiates a ChangeFeedIterable for non-client side partitioning queries. :param CosmosClient client: Instance of document client. @@ -85,7 +85,7 @@ def __init__( super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) - def _unpack(self, block): + def _unpack(self, block) -> Tuple[str, List[Dict[str, Any]]]: continuation = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get('etag') @@ -94,12 +94,9 @@ def _unpack(self, block): self._did_a_call_already = False return continuation, block - def _fetch_next(self, *args): # pylint: disable=unused-argument + def _fetch_next(self, *args) -> List[Dict[str, Any]]: # pylint: disable=unused-argument """Return a block of results with respecting retry policy. - This method only exists for backward compatibility reasons. (Because - QueryIterable has exposed fetch_next_block api). - :param Any args: :return: List of results. :rtype: list @@ -113,7 +110,7 @@ def _fetch_next(self, *args): # pylint: disable=unused-argument raise StopIteration return block - def _initialize_change_feed_fetcher(self): + def _initialize_change_feed_fetcher(self) -> None: change_feed_state_context = self._options.pop("changeFeedStateContext") change_feed_state = \ ChangeFeedState.from_json( @@ -123,7 +120,7 @@ def _initialize_change_feed_fetcher(self): self._options["changeFeedState"] = change_feed_state - if isinstance(change_feed_state, ChangeFeedStateV1): + if change_feed_state.version == ChangeFeedStateVersion.V1: self._change_feed_fetcher = ChangeFeedFetcherV1( self._client, self._collection_link, @@ -140,11 +137,11 @@ def _initialize_change_feed_fetcher(self): def _validate_change_feed_state_context(self, change_feed_state_context: Dict[str, Any]) -> None: - if change_feed_state_context.get("continuationPkRangeId"): + if change_feed_state_context.get("continuationPkRangeId") is not None: # if continuation token is in v1 format, throw exception if feed_range is set - if change_feed_state_context.get("feedRange"): + if change_feed_state_context.get("feedRange") is not None: raise ValueError("feed_range and continuation are incompatible") - elif change_feed_state_context.get("continuationFeedRange"): + elif change_feed_state_context.get("continuationFeedRange") is not None: # if continuation token is in v2 format, since the token itself contains the full change feed state # so we will ignore other parameters (including incompatible parameters) if they passed in pass diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py index 30d0ce787983..dc255eced586 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_start_from.py @@ -42,6 +42,9 @@ class ChangeFeedStartFromInternal(ABC): type_property_name = "Type" + def __init__(self, start_from_type: ChangeFeedStartFromType) -> None: + self.version = start_from_type + @abstractmethod def to_dict(self) -> Dict[str, Any]: pass @@ -86,6 +89,9 @@ class ChangeFeedStartFromBeginning(ChangeFeedStartFromInternal): """Class for change feed start from beginning implementation in the Azure Cosmos database service. """ + def __init__(self) -> None: + super().__init__(ChangeFeedStartFromType.BEGINNING) + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.BEGINNING.value @@ -106,12 +112,13 @@ class ChangeFeedStartFromETagAndFeedRange(ChangeFeedStartFromInternal): _etag_property_name = "Etag" _feed_range_property_name = "FeedRange" - def __init__(self, etag, feed_range): + def __init__(self, etag, feed_range) -> None: if feed_range is None: raise ValueError("feed_range is missing") self._etag = etag self._feed_range = feed_range + super().__init__(ChangeFeedStartFromType.LEASE) def to_dict(self) -> Dict[str, Any]: return { @@ -142,6 +149,9 @@ class ChangeFeedStartFromNow(ChangeFeedStartFromInternal): """Class for change feed start from etag and feed range implementation in the Azure Cosmos database service. """ + def __init__(self) -> None: + super().__init__(ChangeFeedStartFromType.NOW) + def to_dict(self) -> Dict[str, Any]: return { self.type_property_name: ChangeFeedStartFromType.NOW.value @@ -166,6 +176,7 @@ def __init__(self, start_time: datetime): raise ValueError("start_time is missing") self._start_time = start_time + super().__init__(ChangeFeedStartFromType.POINT_IN_TIME) def to_dict(self) -> Dict[str, Any]: return { diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 742c4891bfdf..95d5624017a8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -27,6 +27,7 @@ import collections import json from abc import ABC, abstractmethod +from enum import Enum from typing import Optional, Union, List, Any, Dict, Deque from azure.cosmos import http_constants @@ -41,10 +42,16 @@ from azure.cosmos.exceptions import CosmosFeedRangeGoneError from azure.cosmos.partition_key import _Empty, _Undefined +class ChangeFeedStateVersion(Enum): + V1 = "v1" + V2 = "v2" class ChangeFeedState(ABC): version_property_name = "v" + def __init__(self, version: ChangeFeedStateVersion) -> None: + self.version = version + @abstractmethod def populate_feed_options(self, feed_options: Dict[str, Any]) -> None: pass @@ -106,7 +113,7 @@ def __init__( change_feed_start_from: ChangeFeedStartFromInternal, partition_key_range_id: Optional[str] = None, partition_key: Optional[Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]] = None, # pylint: disable=line-too-long - continuation: Optional[str] = None): + continuation: Optional[str] = None) -> None: self._container_link = container_link self._container_rid = container_rid @@ -114,6 +121,7 @@ def __init__( self._partition_key_range_id = partition_key_range_id self._partition_key = partition_key self._continuation = continuation + super(ChangeFeedStateV1).__init__(ChangeFeedStateVersion.V1) @property def container_rid(self): @@ -187,7 +195,8 @@ def __init__( container_rid: str, feed_range: FeedRange, change_feed_start_from: ChangeFeedStartFromInternal, - continuation: Optional[FeedRangeCompositeContinuation]): + continuation: Optional[FeedRangeCompositeContinuation] + ) -> None: self._container_link = container_link self._container_rid = container_rid @@ -207,13 +216,15 @@ def __init__( else: self._continuation = continuation + super(ChangeFeedStateV2).__init__(ChangeFeedStateVersion.V2) + @property def container_rid(self) -> str : return self._container_rid def to_dict(self) -> Dict[str, Any]: return { - self.version_property_name: "V2", + self.version_property_name: ChangeFeedStateVersion.V2.value, self.container_rid_property_name: self._container_rid, self.change_feed_mode_property_name: "Incremental", self.change_feed_start_from_property_name: self._change_feed_start_from.to_dict(), diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py index 9945405e4b57..90d3d6132822 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py @@ -22,41 +22,41 @@ """Internal class for change feed composite continuation token in the Azure Cosmos database service. """ -from typing import Optional +from typing import Optional, Dict, Any from azure.cosmos._routing.routing_range import Range -class CompositeContinuationToken(object): +class CompositeContinuationToken: token_property_name = "token" feed_range_property_name = "range" - def __init__(self, feed_range: Range, token: Optional[str] = None): + def __init__(self, feed_range: Range, token: Optional[str] = None) -> None: if feed_range is None: raise ValueError("Missing required parameter feed_range") self._token = token self._feed_range = feed_range - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: return { self.token_property_name: self._token, - self.feed_range_property_name: self._feed_range.to_dict() + self.feed_range_property_name: self.feed_range.to_dict() } @property - def feed_range(self): + def feed_range(self) -> Range: return self._feed_range @property - def token(self): + def token(self) -> str: return self._token - def update_token(self, etag): + def update_token(self, etag) -> None: self._token = etag @classmethod - def from_json(cls, data): + def from_json(cls, data) -> 'CompositeContinuationToken': token = data.get(cls.token_property_name) if token is None: raise ValueError(f"Invalid composite token [Missing {cls.token_property_name}]") @@ -69,4 +69,4 @@ def from_json(cls, data): return cls(feed_range=feed_range, token=token) def __repr__(self): - return f"CompositeContinuationToken(token={self.token}, range={self._feed_range})" + return f"CompositeContinuationToken(token={self.token}, range={self.feed_range})" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index 481496159cf3..b4f731f2c2ef 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -45,7 +45,7 @@ class FeedRangePartitionKey(FeedRange): def __init__( self, pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], - feed_range: Range): # pylint: disable=line-too-long + feed_range: Range) -> None: # pylint: disable=line-too-long if pk_value is None: raise ValueError("PartitionKey cannot be None") @@ -87,7 +87,7 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartiti class FeedRangeEpk(FeedRange): type_property_name = "Range" - def __init__(self, feed_range: Range): + def __init__(self, feed_range: Range) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index e8bfe60ced3f..0aaebb616249 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -31,7 +31,7 @@ from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range -class FeedRangeCompositeContinuation(object): +class FeedRangeCompositeContinuation: _version_property_name = "v" _container_rid_property_name = "rid" _continuation_property_name = "continuation" @@ -40,7 +40,7 @@ def __init__( self, container_rid: str, feed_range: FeedRange, - continuation: Deque[CompositeContinuationToken]): + continuation: Deque[CompositeContinuationToken]) -> None: if container_rid is None: raise ValueError("container_rid is missing") @@ -51,7 +51,7 @@ def __init__( self._initial_no_result_range = None @property - def current_token(self): + def current_token(self) -> CompositeContinuationToken: return self._current_token def to_dict(self) -> Dict[str, Any]: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 1e86ccb1aa44..bc4f975a19e8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -490,68 +490,97 @@ def query_items_change_feed( *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, - partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long + """Get a sorted list of items that were changed in the entire container, + in the order in which they were modified. + + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :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 Callable response_hook: A callable invoked with the response metadata. + :returns: An AsyncItemPaged of items (dicts). + :rtype: AsyncItemPaged[Dict[str, Any]] + """ + ... + + @overload + def query_items_change_feed( + self, + *, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: PartitionKeyType, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> AsyncItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. Now: Processing change feed from the current time, so only events for all future changes will be retrieved. ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) - :keyword PartitionKeyType partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :keyword partition_key: The partition key that is used to define the scope + (logical partition or a subset of a container) + :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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 Callable response_hook: A callable invoked with the response metadata. :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @overload def query_items_change_feed( self, *, - feed_range: Optional[str] = None, + feed_range: str, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword str feed_range: The feed range that is used to define the scope. By default, the scope will be the entire container. + :keyword str feed_range: The feed range that is used to define the scope. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: The start time to start processing chang feed items. + :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. Now: Processing change feed from the current time, so only events for all future changes will be retrieved. ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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 Callable response_hook: A callable invoked with the response metadata. :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @overload def query_items_change_feed( self, *, - continuation: Optional[str] = None, + continuation: str, max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. @@ -559,6 +588,7 @@ def query_items_change_feed( :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 Callable response_hook: A callable invoked with the response metadata. :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ @@ -566,73 +596,59 @@ def query_items_change_feed( ... @distributed_trace - def query_items_change_feed( + def query_items_change_feed( # pylint: disable=unused-argument self, *args: Any, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - if kwargs.get("priority"): + if kwargs.get("priority") is not None: kwargs['priority'] = kwargs['priority'] feed_options = _build_options(kwargs) change_feed_state_context = {} # Back compatibility with deprecation warnings for partition_key_range_id - if (args and args[0] is not None) or kwargs.get("partition_key_range_id"): + if kwargs.get("partition_key_range_id") is not None: warnings.warn( "partition_key_range_id is deprecated. Please pass in feed_range instead.", DeprecationWarning ) - try: - change_feed_state_context["partitionKeyRangeId"] = kwargs.pop('partition_key_range_id') - except KeyError: - change_feed_state_context['partitionKeyRangeId'] = args[0] + change_feed_state_context["partitionKeyRangeId"] = kwargs.pop('partition_key_range_id') # Back compatibility with deprecation warnings for is_start_from_beginning - if (len(args) >= 2 and args[1] is not None) or kwargs.get("is_start_from_beginning"): + if kwargs.get("is_start_from_beginning") is not None: warnings.warn( "is_start_from_beginning is deprecated. Please pass in start_time instead.", DeprecationWarning ) - try: - is_start_from_beginning = kwargs.pop('is_start_from_beginning') - except KeyError: - is_start_from_beginning = args[1] + if kwargs.get("start_time") is not None: + raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") - if is_start_from_beginning: + is_start_from_beginning = kwargs.pop('is_start_from_beginning') + if is_start_from_beginning is True: change_feed_state_context["startTime"] = "Beginning" # parse start_time - if kwargs.get("start_time"): - if change_feed_state_context.get("startTime") is not None: - raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") - + if kwargs.get("start_time") is not None: start_time = kwargs.pop('start_time') if not isinstance(start_time, (datetime, str)): raise TypeError( - "'start_time' must be either a datetime object, or either the values 'now' or 'beginning'.") + "'start_time' must be either a datetime object, or either the values 'Now' or 'Beginning'.") change_feed_state_context["startTime"] = start_time # parse continuation token - if len(args) >= 3 and args[2] is not None or feed_options.get("continuation"): - try: - continuation = feed_options.pop('continuation') - except KeyError: - continuation = args[2] - change_feed_state_context["continuation"] = continuation - - if len(args) >= 4 and args[3] is not None or kwargs.get("max_item_count"): - try: - feed_options["maxItemCount"] = kwargs.pop('max_item_count') - except KeyError: - feed_options["maxItemCount"] = args[3] - - if kwargs.get("partition_key"): + if feed_options.get("continuation") is not None: + change_feed_state_context["continuation"] = feed_options.pop('continuation') + + if kwargs.get("max_item_count") is not None: + feed_options["maxItemCount"] = kwargs.pop('max_item_count') + + if kwargs.get("partition_key") is not None: change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop("partition_key")) - if kwargs.get("feed_range"): + if kwargs.get("feed_range") is not None: change_feed_state_context["feedRange"] = kwargs.pop('feed_range') feed_options["containerProperties"] = self._get_properties() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 5e29302bcaae..cee4bcb4843b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -326,73 +326,98 @@ def query_items_change_feed( *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, - partition_key: Optional[PartitionKeyType] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long + """Get a sorted list of items that were changed in the entire container, + in the order in which they were modified, + + :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword start_time:The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :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 Callable response_hook: A callable invoked with the response metadata. + :returns: An Iterable of items (dicts). + :rtype: Iterable[Dict[str, Any]] + """ + ... + + @overload + def query_items_change_feed( + self, + *, + max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: PartitionKeyType, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any + ) -> ItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Union[datetime, Literal["Now", "Beginning"]] start_time: - The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) - :keyword PartitionKeyType partition_key: The partition key that is used to define the scope + :keyword start_time:The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :keyword partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) + :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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 Callable response_hook: A callable invoked with the response metadata. :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @overload def query_items_change_feed( self, *, - feed_range: Optional[str] = None, + feed_range: str, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str feed_range: The feed range that is used to define the scope. - By default, the scope will be the entire container. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword Union[datetime, Literal["Now", "Beginning"]] - start_time: The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current (NOW) + :keyword start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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 Callable response_hook: A callable invoked with the response metadata. :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @overload def query_items_change_feed( self, *, - continuation: Optional[str] = None, + continuation: str, max_item_count: Optional[int] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - # pylint: disable=line-too-long """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. @@ -400,10 +425,10 @@ def query_items_change_feed( :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 Callable response_hook: A callable invoked with the response metadata. :returns: An Iterable of items (dicts). :rtype: Iterable[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @distributed_trace @@ -413,13 +438,13 @@ def query_items_change_feed( **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: # pylint: disable=too-many-statements - if kwargs.get("priority"): + if kwargs.get("priority") is not None: kwargs['priority'] = kwargs['priority'] feed_options = build_options(kwargs) change_feed_state_context = {} # Back compatibility with deprecation warnings for partition_key_range_id - if (args and args[0] is not None) or kwargs.get("partition_key_range_id"): + if (args and args[0] is not None) or kwargs.get("partition_key_range_id") is not None: warnings.warn( "partition_key_range_id is deprecated. Please pass in feed_range instead.", DeprecationWarning @@ -431,51 +456,52 @@ def query_items_change_feed( change_feed_state_context['partitionKeyRangeId'] = args[0] # Back compatibility with deprecation warnings for is_start_from_beginning - if (len(args) >= 2 and args[1] is not None) or kwargs.get("is_start_from_beginning"): + if (len(args) >= 2 and args[1] is not None) or kwargs.get("is_start_from_beginning") is not None: warnings.warn( "is_start_from_beginning is deprecated. Please pass in start_time instead.", DeprecationWarning ) + if kwargs.get("start_time") is not None: + raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") + try: is_start_from_beginning = kwargs.pop('is_start_from_beginning') except KeyError: is_start_from_beginning = args[1] - if is_start_from_beginning: + if is_start_from_beginning is True: change_feed_state_context["startTime"] = "Beginning" # parse start_time - if kwargs.get("start_time"): - if change_feed_state_context.get("startTime") is not None: - raise ValueError("is_start_from_beginning and start_time are exclusive, please only set one of them") + if kwargs.get("start_time") is not None: start_time = kwargs.pop('start_time') if not isinstance(start_time, (datetime, str)): raise TypeError( - "'start_time' must be either a datetime object, or either the values 'now' or 'beginning'.") + "'start_time' must be either a datetime object, or either the values 'Now' or 'Beginning'.") change_feed_state_context["startTime"] = start_time # parse continuation token - if len(args) >= 3 and args[2] is not None or feed_options.get("continuation"): + if len(args) >= 3 and args[2] is not None or feed_options.get("continuation") is not None: try: continuation = feed_options.pop('continuation') except KeyError: continuation = args[2] change_feed_state_context["continuation"] = continuation - if len(args) >= 4 and args[3] is not None or kwargs.get("max_item_count"): + if len(args) >= 4 and args[3] is not None or kwargs.get("max_item_count") is not None: try: feed_options["maxItemCount"] = kwargs.pop('max_item_count') except KeyError: feed_options["maxItemCount"] = args[3] - if kwargs.get("partition_key"): + if kwargs.get("partition_key") is not None: change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop('partition_key')) change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) - if kwargs.get("feed_range"): + if kwargs.get("feed_range") is not None: change_feed_state_context["feedRange"] = kwargs.pop('feed_range') container_properties = self._get_properties() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py index 262de2ffbbbe..ed6e6b114869 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py @@ -137,7 +137,7 @@ def __init__(self, **kwargs): class CosmosFeedRangeGoneError(CosmosHttpResponseError): - """An HTTP error response with status code 404.""" + """An HTTP error response with status code 410.""" def __init__(self, message=None, response=None, **kwargs): """ :param int sub_status_code: HTTP response sub code. diff --git a/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst b/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst new file mode 100644 index 000000000000..5b7484884dd7 --- /dev/null +++ b/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst @@ -0,0 +1,8 @@ +.. toctree:: + :maxdepth: 5 + :glob: + :caption: Developer Documentation + + ref/azure.common + /Users/annie-mac/dev/git/azure-sdk-for-python/sdk/cosmos/azure-mgmt-cosmosdb/.tox/sphinx/tmp/dist/unzipped/docgen/azure.mgmt.cosmosdb.rst + ref/azure.servicemanagement From 65ed1329052930c6508da2cea505f9a2994aebf8 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 27 Aug 2024 22:21:43 -0700 Subject: [PATCH 14/59] resolve comments --- .../_routing/aio/routing_map_provider.py | 8 +++---- .../cosmos/_routing/routing_map_provider.py | 8 +++---- .../azure/cosmos/aio/_container.py | 18 +++++++++++++-- .../azure-cosmos/azure/cosmos/container.py | 22 +++++++++++++++---- 4 files changed, 42 insertions(+), 14 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py index ebf1ee82b005..ba0b5ca3a3e6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py @@ -49,7 +49,7 @@ def __init__(self, client): # keeps the cached collection routing map by collection id self._collection_routing_map_by_item = {} - async def get_overlapping_ranges(self, collection_link, partition_key_ranges): + async def get_overlapping_ranges(self, collection_link, partition_key_ranges, **kwargs): """Given a partition key range and a collection, return the list of overlapping partition key ranges. @@ -64,7 +64,7 @@ async def get_overlapping_ranges(self, collection_link, partition_key_ranges): collection_routing_map = self._collection_routing_map_by_item.get(collection_id) if collection_routing_map is None: - collection_pk_ranges = [pk async for pk in cl._ReadPartitionKeyRanges(collection_link)] + collection_pk_ranges = [pk async for pk in cl._ReadPartitionKeyRanges(collection_link, **kwargs)] # for large collections, a split may complete between the read partition key ranges query page responses, # causing the partitionKeyRanges to have both the children ranges and their parents. Therefore, we need # to discard the parent ranges to have a valid routing map. @@ -131,7 +131,7 @@ class SmartRoutingMapProvider(PartitionKeyRangeCache): invocation of CollectionRoutingMap.get_overlapping_ranges() """ - async def get_overlapping_ranges(self, collection_link, partition_key_ranges): + async def get_overlapping_ranges(self, collection_link, partition_key_ranges, **kwargs): """ Given the sorted ranges and a collection, Returns the list of overlapping partition key ranges @@ -166,7 +166,7 @@ async def get_overlapping_ranges(self, collection_link, partition_key_ranges): queryRange = currentProvidedRange overlappingRanges = await PartitionKeyRangeCache.get_overlapping_ranges(self, - collection_link, queryRange) + collection_link, queryRange, **kwargs) assert overlappingRanges, "code bug: returned overlapping ranges for queryRange {} is empty".format( queryRange ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py index 59c609dec7ea..5a6bb304b5c8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py @@ -50,7 +50,7 @@ def __init__(self, client): # keeps the cached collection routing map by collection id self._collection_routing_map_by_item = {} - def get_overlapping_ranges(self, collection_link, partition_key_ranges): + def get_overlapping_ranges(self, collection_link, partition_key_ranges, **kwargs): """Given a partition key range and a collection, return the list of overlapping partition key ranges. @@ -65,7 +65,7 @@ def get_overlapping_ranges(self, collection_link, partition_key_ranges): collection_routing_map = self._collection_routing_map_by_item.get(collection_id) if collection_routing_map is None: - collection_pk_ranges = list(cl._ReadPartitionKeyRanges(collection_link)) + collection_pk_ranges = list(cl._ReadPartitionKeyRanges(collection_link, **kwargs)) # for large collections, a split may complete between the read partition key ranges query page responses, # causing the partitionKeyRanges to have both the children ranges and their parents. Therefore, we need # to discard the parent ranges to have a valid routing map. @@ -132,7 +132,7 @@ class SmartRoutingMapProvider(PartitionKeyRangeCache): invocation of CollectionRoutingMap.get_overlapping_ranges() """ - def get_overlapping_ranges(self, collection_link, partition_key_ranges): + def get_overlapping_ranges(self, collection_link, partition_key_ranges, **kwargs): """ Given the sorted ranges and a collection, Returns the list of overlapping partition key ranges @@ -166,7 +166,7 @@ def get_overlapping_ranges(self, collection_link, partition_key_ranges): else: queryRange = currentProvidedRange - overlappingRanges = PartitionKeyRangeCache.get_overlapping_ranges(self, collection_link, queryRange) + overlappingRanges = PartitionKeyRangeCache.get_overlapping_ranges(self, collection_link, queryRange, **kwargs) assert overlappingRanges, "code bug: returned overlapping ranges for queryRange {} is empty".format( queryRange ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index bc4f975a19e8..f289442cf3b3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1228,15 +1228,29 @@ async def execute_item_batch( return await self.client_connection.Batch( collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) - async def read_feed_ranges( # pylint: disable=unused-argument + async def read_feed_ranges( self, + *, + force_refresh: Optional[bool] = False, **kwargs: Any ) -> List[str]: + """ Obtains a list of feed ranges that can be used to parallelize feed operations. + + :param bool force_refresh: + Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. + :returns: A list representing the feed ranges in base64 encoded string + :rtype: List[str] + + """ + if force_refresh is True: + self.client_connection.refresh_routing_map_provider() + partition_key_ranges =\ await self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, # default to full range - [Range("", "FF", True, False)]) + [Range("", "FF", True, False)], + **kwargs) return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index cee4bcb4843b..017c58b8b492 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1310,14 +1310,28 @@ def delete_all_items_by_partition_key( self.client_connection.DeleteAllItemsByPartitionKey( collection_link=self.container_link, options=request_options, **kwargs) - def read_feed_ranges( # pylint: disable=unused-argument + def read_feed_ranges( self, - **kwargs: Any - ) -> List[str]: + *, + force_refresh: Optional[bool] = False, + **kwargs: Any) -> List[str]: + + """ Obtains a list of feed ranges that can be used to parallelize feed operations. + + :param bool force_refresh: + Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. + :returns: A list representing the feed ranges in base64 encoded string + :rtype: List[str] + + """ + if force_refresh is True: + self.client_connection.refresh_routing_map_provider() + partition_key_ranges =\ self.client_connection._routing_map_provider.get_overlapping_ranges( self.container_link, - [Range("", "FF", True, False)]) # default to full range + [Range("", "FF", True, False)], # default to full range + **kwargs) return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() for partitionKeyRange in partition_key_ranges] From 4bb30d27fe3e936f70a1cb632545eea6b65632f1 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Tue, 27 Aug 2024 22:43:45 -0700 Subject: [PATCH 15/59] resolve comments --- .../_change_feed/aio/change_feed_iterable.py | 19 ++++---- .../_change_feed/change_feed_iterable.py | 7 ++- .../cosmos/_change_feed/change_feed_state.py | 43 ++++++++----------- .../azure/cosmos/_cosmos_client_connection.py | 5 +-- .../azure/cosmos/_routing/routing_range.py | 4 ++ .../azure-cosmos/azure/cosmos/_utils.py | 10 ----- .../azure/cosmos/aio/_container.py | 22 +++++++--- .../aio/_cosmos_client_connection_async.py | 5 +-- .../azure-cosmos/azure/cosmos/container.py | 16 +++---- .../azure-cosmos/azure/cosmos/exceptions.py | 14 ------ .../azure/cosmos/partition_key.py | 16 +++---- 11 files changed, 68 insertions(+), 93 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 745ed19279c7..f265805d390c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,14 +21,13 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any, Optional, Callable, Coroutine, Tuple, List, AsyncIterator +from collections.abc import Awaitable +from typing import Dict, Any, Optional, Callable, Tuple, List, AsyncIterator from azure.core.async_paging import AsyncPageIterator -from azure.cosmos import PartitionKey from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion -from azure.cosmos._utils import is_base64_encoded # pylint: disable=protected-access @@ -43,7 +42,7 @@ def __init__( self, client, options: Dict[str, Any], - fetch_function=Optional[Callable[[Dict[str, Any]], Coroutine[Tuple[List[Dict[str, Any]], Dict[str, Any]]]]], + fetch_function=Optional[Callable[[Dict[str, Any]], Awaitable[Tuple[List[Dict[str, Any]], Dict[str, Any]]]]], collection_link=Optional[str], continuation_token=Optional[str], ) -> None: @@ -79,10 +78,10 @@ def __init__( # v2 version: the continuation token will be base64 encoded composition token # which includes full change feed state if continuation is not None: - if is_base64_encoded(continuation): - change_feed_state_context["continuationFeedRange"] = continuation - else: + if continuation.isdigit() or continuation.strip('\'"').isdigit(): change_feed_state_context["continuationPkRangeId"] = continuation + else: + change_feed_state_context["continuationFeedRange"] = continuation self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context @@ -118,16 +117,14 @@ async def _initialize_change_feed_fetcher(self) -> None: conn_properties = await self._options.pop("containerProperties") if change_feed_state_context.get("partitionKey"): change_feed_state_context["partitionKey"] = await change_feed_state_context.pop("partitionKey") - pk_properties = conn_properties.get("partitionKey") - partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) change_feed_state_context["partitionKeyFeedRange"] =\ - partition_key_definition._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) + await change_feed_state_context.pop("partitionKeyFeedRange") change_feed_state =\ ChangeFeedState.from_json(self._collection_link, conn_properties["_rid"], change_feed_state_context) self._options["changeFeedState"] = change_feed_state - if change_feed_state.version != ChangeFeedStateVersion.V1: + if change_feed_state.version == ChangeFeedStateVersion.V1: self._change_feed_fetcher = ChangeFeedFetcherV1( self._client, self._collection_link, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 6eaaa31bfd8f..e8f3e414bc4f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -27,7 +27,6 @@ from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion -from azure.cosmos._utils import is_base64_encoded class ChangeFeedIterable(PageIterator): @@ -75,10 +74,10 @@ def __init__( # v2 version: the continuation token will be base64 encoded composition token # which includes full change feed state if continuation is not None: - if is_base64_encoded(continuation): - change_feed_state_context["continuationFeedRange"] = continuation - else: + if continuation.isdigit() or continuation.strip('\'"').isdigit(): change_feed_state_context["continuationPkRangeId"] = continuation + else: + change_feed_state_context["continuationFeedRange"] = continuation self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 95d5624017a8..77e603b3d834 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -39,7 +39,8 @@ from azure.cosmos._routing.aio.routing_map_provider import SmartRoutingMapProvider as AsyncSmartRoutingMapProvider from azure.cosmos._routing.routing_map_provider import SmartRoutingMapProvider from azure.cosmos._routing.routing_range import Range -from azure.cosmos.exceptions import CosmosFeedRangeGoneError +from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.cosmos.http_constants import StatusCodes, SubStatusCodes from azure.cosmos.partition_key import _Empty, _Undefined class ChangeFeedStateVersion(Enum): @@ -93,7 +94,7 @@ def from_json( if version is None: raise ValueError("Invalid base64 encoded continuation string [Missing version]") - if version == "V2": + if version == ChangeFeedStateVersion.V2.value: return ChangeFeedStateV2.from_continuation(container_link, container_rid, continuation_json) raise ValueError("Invalid base64 encoded continuation string [Invalid version]") @@ -121,7 +122,7 @@ def __init__( self._partition_key_range_id = partition_key_range_id self._partition_key = partition_key self._continuation = continuation - super(ChangeFeedStateV1).__init__(ChangeFeedStateVersion.V1) + super(ChangeFeedStateV1, self).__init__(ChangeFeedStateVersion.V1) @property def container_rid(self): @@ -148,11 +149,6 @@ def populate_request_headers( request_headers: Dict[str, Any]) -> None: request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. So when reading the changeFeed by LSN, - # it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, - # we will need to always pass the start time in the header. self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation: request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation @@ -164,11 +160,6 @@ async def populate_request_headers_async( request_headers[http_constants.HttpHeaders.AIM] = http_constants.HttpHeaders.IncrementalFeedHeaderValue - # When a merge happens, the child partition will contain documents ordered by LSN but the _ts/creation time - # of the documents may not be sequential. - # So when reading the changeFeed by LSN, it is possible to encounter documents with lower _ts. - # In order to guarantee we always get the documents after customer's point start time, - # we will need to always pass the start time in the header. self._change_feed_start_from.populate_request_headers(request_headers) if self._continuation: request_headers[http_constants.HttpHeaders.IfNoneMatch] = self._continuation @@ -216,7 +207,7 @@ def __init__( else: self._continuation = continuation - super(ChangeFeedStateV2).__init__(ChangeFeedStateVersion.V2) + super(ChangeFeedStateV2, self).__init__(ChangeFeedStateVersion.V2) @property def container_rid(self) -> str : @@ -258,11 +249,7 @@ def populate_request_headers( [self._continuation.current_token.feed_range]) if len(over_lapping_ranges) > 1: - raise CosmosFeedRangeGoneError( - message= - f"Range {self._continuation.current_token.feed_range}" - f" spans {len(over_lapping_ranges)}" - f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + raise self.get_feed_range_gone_error(over_lapping_ranges) overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) if overlapping_feed_range == self._continuation.current_token.feed_range: @@ -304,11 +291,7 @@ async def populate_request_headers_async( [self._continuation.current_token.feed_range]) if len(over_lapping_ranges) > 1: - raise CosmosFeedRangeGoneError( - message= - f"Range {self._continuation.current_token.feed_range}" - f" spans {len(over_lapping_ranges)}" - f" physical partitions: {[child_range['id'] for child_range in over_lapping_ranges]}") + raise self.get_feed_range_gone_error(over_lapping_ranges) overlapping_feed_range = Range.PartitionKeyRangeToRange(over_lapping_ranges[0]) if overlapping_feed_range == self._continuation.current_token.feed_range: @@ -348,6 +331,18 @@ def should_retry_on_not_modified_response(self) -> bool: def apply_not_modified_response(self) -> None: self._continuation.apply_not_modified_response() + def get_feed_range_gone_error(self, over_lapping_ranges: list[Dict[str, Any]]) -> CosmosHttpResponseError: + formatted_message =\ + (f"Status code: {StatusCodes.GONE} " + f"Sub-status: {SubStatusCodes.PARTITION_KEY_RANGE_GONE}. " + f"Range {self._continuation.current_token.feed_range}" + f" spans {len(over_lapping_ranges)} physical partitions:" + f" {[child_range['id'] for child_range in over_lapping_ranges]}") + + response_error = CosmosHttpResponseError(status_code=StatusCodes.GONE, message=formatted_message) + response_error.sub_status = SubStatusCodes.PARTITION_KEY_RANGE_GONE + return response_error + @classmethod def from_continuation( cls, 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 87e72391cf0a..1f44c491ae46 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3025,9 +3025,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: partition_key_range_id ) - change_feed_state = options.get("changeFeedState", None) - if change_feed_state and isinstance(change_feed_state, ChangeFeedState): - change_feed_state.populate_request_headers(self._routing_map_provider, headers) + if options.get("changeFeedState") is not None: + options.pop("changeFeedState").populate_request_headers(self._routing_map_provider, headers) result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index a2d789f20644..f2e7576bf376 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -25,8 +25,12 @@ import base64 import binascii import json +from typing import Dict, Any +def partition_key_range_to_range_string(partition_key_range: Dict[str, Any]) -> str: + return Range.PartitionKeyRangeToRange(partition_key_range).to_base64_encoded_string() + class PartitionKeyRange(object): """Partition Key Range Constants""" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 6e3a8c67fcfe..1b3d0370e6ef 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -69,13 +69,3 @@ def get_index_metrics_info(delimited_string: Optional[str]) -> Dict[str, Any]: return result except (json.JSONDecodeError, ValueError): return {} - - -def is_base64_encoded(data: str) -> bool: - if data is None: - return False - try: - base64.b64decode(data, validate=True).decode('utf-8') - return True - except (json.JSONDecodeError, ValueError): - return False diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index f289442cf3b3..fc6503aa630e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -41,14 +41,13 @@ GenerateGuidId, _set_properties_cache ) -from .._routing import routing_range -from .._routing.routing_range import Range +from .._routing.routing_range import Range, partition_key_range_to_range_string from ..offer import ThroughputProperties from ..partition_key import ( NonePartitionKeyValue, _return_undefined_or_empty_partition_key, _Empty, - _Undefined + _Undefined, PartitionKey ) __all__ = ("ContainerProxy",) @@ -136,6 +135,16 @@ async def _set_partition_key( return _return_undefined_or_empty_partition_key(await self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) + async def _get_epk_range_for_partition_key( + self, + partition_key_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> Range: # pylint: disable=line-too-long + + container_properties = await self._get_properties() + partition_key_definition = container_properties["partitionKey"] + partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) + + return partition_key._get_epk_range_for_partition_key(partition_key_value) + @distributed_trace_async async def read( self, @@ -646,7 +655,9 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["maxItemCount"] = kwargs.pop('max_item_count') if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop("partition_key")) + change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.get("partition_key")) + change_feed_state_context["partitionKeyFeedRange"] = \ + self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: change_feed_state_context["feedRange"] = kwargs.pop('feed_range') @@ -1252,5 +1263,4 @@ async def read_feed_ranges( [Range("", "FF", True, False)], **kwargs) - return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() - for partitionKeyRange in partition_key_ranges] + return [partition_key_range_to_range_string(partitionKeyRange) for partitionKeyRange in partition_key_ranges] 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 b2a9aaff9ec1..47e83f3e31df 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 @@ -2814,9 +2814,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: ) headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) - change_feed_state = options.get("changeFeedState", None) - if change_feed_state and isinstance(change_feed_state, ChangeFeedState): - await change_feed_state.populate_request_headers_async(self._routing_map_provider, headers) + if options.get("changeFeedState") is not None: + await options.pop("changeFeedState").populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) if response_hook: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 017c58b8b492..7ecac5391407 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -39,8 +39,7 @@ _set_properties_cache ) from ._cosmos_client_connection import CosmosClientConnection -from ._routing import routing_range -from ._routing.routing_range import Range +from ._routing.routing_range import Range, partition_key_range_to_range_string from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -133,7 +132,7 @@ def _set_partition_key( def _get_epk_range_for_partition_key( self, - partition_key_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> Range: # pylint: disable=line-too-long + partition_key_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> Range: # pylint: disable=line-too-long container_properties = self._get_properties() partition_key_definition = container_properties["partitionKey"] partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) @@ -497,9 +496,9 @@ def query_items_change_feed( feed_options["maxItemCount"] = args[3] if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.pop('partition_key')) + change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.get('partition_key')) change_feed_state_context["partitionKeyFeedRange"] =\ - self._get_epk_range_for_partition_key(change_feed_state_context["partitionKey"]) + self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: change_feed_state_context["feedRange"] = kwargs.pop('feed_range') @@ -644,9 +643,7 @@ def query_items( # pylint:disable=docstring-missing-param return items def __is_prefix_partitionkey( - self, - partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> bool: # pylint: disable=line-too-long - + self, partition_key: PartitionKeyType) -> bool: properties = self._get_properties() pk_properties = properties["partitionKey"] partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) @@ -1333,5 +1330,4 @@ def read_feed_ranges( [Range("", "FF", True, False)], # default to full range **kwargs) - return [routing_range.Range.PartitionKeyRangeToRange(partitionKeyRange).to_base64_encoded_string() - for partitionKeyRange in partition_key_ranges] + return [partition_key_range_to_range_string(partitionKeyRange) for partitionKeyRange in partition_key_ranges] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py index ed6e6b114869..7170a4d1dc39 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/exceptions.py @@ -135,20 +135,6 @@ def __init__(self, **kwargs): self.history = None super(CosmosClientTimeoutError, self).__init__(message, **kwargs) - -class CosmosFeedRangeGoneError(CosmosHttpResponseError): - """An HTTP error response with status code 410.""" - def __init__(self, message=None, response=None, **kwargs): - """ - :param int sub_status_code: HTTP response sub code. - """ - - self.sub_status = SubStatusCodes.PARTITION_KEY_RANGE_GONE - self.http_error_message = message - formatted_message = "Status code: %d Sub-status: %d\n%s" % (StatusCodes.GONE, self.sub_status, str(message)) - super(CosmosHttpResponseError, self).__init__(message=formatted_message, response=response, **kwargs) - self.status_code = StatusCodes.GONE - def _partition_range_is_gone(e): if (e.status_code == http_constants.StatusCodes.GONE and e.sub_status == http_constants.SubStatusCodes.PARTITION_KEY_RANGE_GONE): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index e4d659a08fac..5870e7519e8b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -149,7 +149,7 @@ def version(self, value: int) -> None: def _get_epk_range_for_prefix_partition_key( self, - pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> _Range: if self.kind != "MultiHash": raise ValueError( @@ -175,16 +175,16 @@ def _get_epk_range_for_prefix_partition_key( def _get_epk_range_for_partition_key( self, - pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined] # pylint: disable=line-too-long + pk_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long ) -> _Range: if self._is_prefix_partition_key(pk_value): return self._get_epk_range_for_prefix_partition_key( - cast(List[Union[str, int, float, bool]], pk_value)) + cast(Sequence[Union[None, bool, int, float, str, Type[NonePartitionKeyValue]]], pk_value)) # else return point range effective_partition_key_string =\ self._get_effective_partition_key_string( - cast(List[Union[str, int, float, bool, _Empty, _Undefined]], [pk_value])) + cast(List[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]], [pk_value])) return _Range(effective_partition_key_string, effective_partition_key_string, True, True) def _get_effective_partition_key_for_hash_partitioning(self) -> str: @@ -193,7 +193,7 @@ def _get_effective_partition_key_for_hash_partitioning(self) -> str: def _get_effective_partition_key_string( self, - pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> Union[int, str]: if not pk_value: return _MinimumInclusiveEffectivePartitionKey @@ -238,7 +238,7 @@ def _write_for_hashing_v2( def _get_effective_partition_key_for_hash_partitioning_v2( self, - pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> str: with BytesIO() as ms: for component in pk_value: @@ -257,7 +257,7 @@ def _get_effective_partition_key_for_hash_partitioning_v2( def _get_effective_partition_key_for_multi_hash_partitioning_v2( self, - pk_value: Sequence[Union[str, int, float, bool, _Empty, _Undefined]] + pk_value: Sequence[Union[None, bool, int, float, str, _Undefined, Type[NonePartitionKeyValue]]] ) -> str: sb = [] for value in pk_value: @@ -281,7 +281,7 @@ def _get_effective_partition_key_for_multi_hash_partitioning_v2( def _is_prefix_partition_key( self, - partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]) -> bool: # pylint: disable=line-too-long + 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": return False if isinstance(partition_key, list) and len(self.path) == len(partition_key): From 5addcdcc0275d876a335f07215c347327ca54197 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Thu, 29 Aug 2024 09:42:28 -0700 Subject: [PATCH 16/59] fix pylint --- .../_change_feed/aio/change_feed_fetcher.py | 6 +-- .../_change_feed/aio/change_feed_iterable.py | 14 +++--- .../_change_feed/change_feed_fetcher.py | 8 ++-- .../_change_feed/change_feed_iterable.py | 13 +++--- .../cosmos/_change_feed/change_feed_state.py | 2 +- .../composite_continuation_token.py | 2 +- ...feed_range_composite_continuation_token.py | 2 +- .../azure/cosmos/_cosmos_client_connection.py | 3 +- .../_routing/aio/routing_map_provider.py | 8 +++- .../cosmos/_routing/routing_map_provider.py | 3 +- .../azure/cosmos/aio/_container.py | 5 ++- .../aio/_cosmos_client_connection_async.py | 3 +- .../azure-cosmos/azure/cosmos/container.py | 45 ++++++++++--------- .../azure-cosmos/test/test_change_feed.py | 1 + .../test/test_change_feed_async.py | 1 + 15 files changed, 64 insertions(+), 52 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index 90aa2d01adfa..376fd0fac397 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -25,7 +25,7 @@ import base64 import json from abc import ABC, abstractmethod -from typing import Dict, Any, List, Callable, Tuple, Awaitable +from typing import Dict, Any, List, Callable, Tuple, Awaitable, cast from azure.cosmos import http_constants, exceptions from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromType @@ -166,7 +166,7 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # there is any items in the response or not. if fetched_items: self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key))) response_headers[continuation_key] = self._get_base64_encoded_continuation() break @@ -177,7 +177,7 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # then we will read from the next feed range until we have looped through all physical partitions self._change_feed_state.apply_not_modified_response() self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key))) if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index f265805d390c..1da44219daa4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,12 +21,12 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from collections.abc import Awaitable -from typing import Dict, Any, Optional, Callable, Tuple, List, AsyncIterator +from typing import Dict, Any, Optional, Callable, Tuple, List, AsyncIterator, Awaitable from azure.core.async_paging import AsyncPageIterator -from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 +from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2, \ + ChangeFeedFetcher from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion @@ -60,7 +60,7 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link - self._change_feed_fetcher = None + self._change_feed_fetcher: Optional[ChangeFeedFetcher] = None if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") @@ -88,7 +88,10 @@ def __init__( super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) - async def _unpack(self, block) -> Tuple[str, AsyncIterator[List[Dict[str, Any]]]]: + async def _unpack( + self, + block: AsyncIterator[List[Dict[str, Any]]] + ) -> Tuple[Optional[str], AsyncIterator[List[Dict[str, Any]]]]: continuation = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get('etag') @@ -107,6 +110,7 @@ async def _fetch_next(self, *args) -> List[Dict[str, Any]]: # pylint: disable=u if self._change_feed_fetcher is None: await self._initialize_change_feed_fetcher() + assert self._change_feed_fetcher is not None block = await self._change_feed_fetcher.fetch_next_block() if not block: raise StopAsyncIteration diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index 35ae9a15a08a..2417eff46259 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -25,7 +25,7 @@ import base64 import json from abc import ABC, abstractmethod -from typing import Dict, Any, List, Callable, Tuple +from typing import Dict, Any, List, Callable, Tuple, cast from azure.cosmos import _retry_utility, http_constants, exceptions from azure.cosmos._change_feed.change_feed_start_from import ChangeFeedStartFromType @@ -87,7 +87,7 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key))) if fetched_items: break @@ -158,7 +158,7 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # In change feed queries, the continuation token is always populated. if fetched_items: self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key))) self._change_feed_state._continuation._move_to_next_token() response_headers[continuation_key] = self._get_base64_encoded_continuation() break @@ -170,7 +170,7 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # then we will read from the next feed range until we have looped through all physical partitions self._change_feed_state.apply_not_modified_response() self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key))) if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index e8f3e414bc4f..c1174db5f93d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -21,11 +21,11 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any, Tuple, List, Optional, Callable +from typing import Dict, Any, Tuple, List, Optional, Callable, cast from azure.core.paging import PageIterator -from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 +from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2, ChangeFeedFetcher from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion @@ -57,7 +57,7 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link - self._change_feed_fetcher = None + self._change_feed_fetcher: Optional[ChangeFeedFetcher] = None if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") @@ -84,8 +84,8 @@ def __init__( super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) - def _unpack(self, block) -> Tuple[str, List[Dict[str, Any]]]: - continuation = None + def _unpack(self, block: List[Dict[str, Any]]) -> Tuple[Optional[str], List[Dict[str, Any]]]: + continuation: Optional[str] = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get('etag') @@ -104,6 +104,7 @@ def _fetch_next(self, *args) -> List[Dict[str, Any]]: # pylint: disable=unused- if self._change_feed_fetcher is None: self._initialize_change_feed_fetcher() + assert self._change_feed_fetcher is not None block = self._change_feed_fetcher.fetch_next_block() if not block: raise StopIteration @@ -114,7 +115,7 @@ def _initialize_change_feed_fetcher(self) -> None: change_feed_state = \ ChangeFeedState.from_json( self._collection_link, - self._options.get("containerRID"), + cast(str, self._options.get("containerRID")), change_feed_state_context) self._options["changeFeedState"] = change_feed_state diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 77e603b3d834..d7c52bf89b88 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -331,7 +331,7 @@ def should_retry_on_not_modified_response(self) -> bool: def apply_not_modified_response(self) -> None: self._continuation.apply_not_modified_response() - def get_feed_range_gone_error(self, over_lapping_ranges: list[Dict[str, Any]]) -> CosmosHttpResponseError: + def get_feed_range_gone_error(self, over_lapping_ranges: List[Dict[str, Any]]) -> CosmosHttpResponseError: formatted_message =\ (f"Status code: {StatusCodes.GONE} " f"Sub-status: {SubStatusCodes.PARTITION_KEY_RANGE_GONE}. " diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py index 90d3d6132822..f0d433fd966e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/composite_continuation_token.py @@ -49,7 +49,7 @@ def feed_range(self) -> Range: return self._feed_range @property - def token(self) -> str: + def token(self) -> Optional[str]: return self._token def update_token(self, etag) -> None: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index 0aaebb616249..fc9b94f27eef 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -48,7 +48,7 @@ def __init__( self._feed_range = feed_range self._continuation = continuation self._current_token = self._continuation[0] - self._initial_no_result_range = None + self._initial_no_result_range: Optional[Range] = None @property def current_token(self) -> CompositeContinuationToken: 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 1f44c491ae46..010b91c76dd6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -57,7 +57,6 @@ from ._auth_policy import CosmosBearerTokenCredentialPolicy from ._base import _set_properties_cache from ._change_feed.change_feed_iterable import ChangeFeedIterable -from ._change_feed.change_feed_state import ChangeFeedState from ._constants import _Constants as Constants from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy from ._range_partition_resolver import RangePartitionResolver @@ -3026,7 +3025,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: ) if options.get("changeFeedState") is not None: - options.pop("changeFeedState").populate_request_headers(self._routing_map_provider, headers) + options.get("changeFeedState").populate_request_headers(self._routing_map_provider, headers) result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py index ba0b5ca3a3e6..e70ae355c495 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py @@ -165,8 +165,12 @@ async def get_overlapping_ranges(self, collection_link, partition_key_ranges, ** else: queryRange = currentProvidedRange - overlappingRanges = await PartitionKeyRangeCache.get_overlapping_ranges(self, - collection_link, queryRange, **kwargs) + overlappingRanges =\ + await PartitionKeyRangeCache.get_overlapping_ranges( + self, + collection_link, + [queryRange], + **kwargs) assert overlappingRanges, "code bug: returned overlapping ranges for queryRange {} is empty".format( queryRange ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py index 5a6bb304b5c8..8dacb5190e07 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py @@ -166,7 +166,8 @@ def get_overlapping_ranges(self, collection_link, partition_key_ranges, **kwargs else: queryRange = currentProvidedRange - overlappingRanges = PartitionKeyRangeCache.get_overlapping_ranges(self, collection_link, queryRange, **kwargs) + overlappingRanges = ( + PartitionKeyRangeCache.get_overlapping_ranges(self, collection_link, [queryRange], **kwargs)) assert overlappingRanges, "code bug: returned overlapping ranges for queryRange {} is empty".format( queryRange ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index fc6503aa630e..4d867a952179 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -655,7 +655,8 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["maxItemCount"] = kwargs.pop('max_item_count') if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.get("partition_key")) + change_feed_state_context["partitionKey"] =\ + self._set_partition_key(cast(PartitionKeyType, kwargs.get("partition_key"))) change_feed_state_context["partitionKeyFeedRange"] = \ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) @@ -1247,7 +1248,7 @@ async def read_feed_ranges( ) -> List[str]: """ Obtains a list of feed ranges that can be used to parallelize feed operations. - :param bool force_refresh: + :keyword bool force_refresh: Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. :returns: A list representing the feed ranges in base64 encoded string :rtype: List[str] 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 47e83f3e31df..675cbbe9dc69 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 @@ -51,7 +51,6 @@ from .._base import _set_properties_cache from .. import documents from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable -from .._change_feed.change_feed_state import ChangeFeedState from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants @@ -2815,7 +2814,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) if options.get("changeFeedState") is not None: - await options.pop("changeFeedState").populate_request_headers_async(self._routing_map_provider, headers) + await options.get("changeFeedState").populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) if response_hook: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 7ecac5391407..b8dcc9a45dc0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -325,11 +325,11 @@ def query_items_change_feed( *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: PartitionKeyType, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - """Get a sorted list of items that were changed in the entire container, - in the order in which they were modified, + """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword start_time:The start time to start processing chang feed items. @@ -338,6 +338,9 @@ def query_items_change_feed( ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current ("Now") :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :keyword partition_key: The partition key that is used to define the scope + (logical partition or a subset of a container) + :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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. @@ -351,24 +354,23 @@ def query_items_change_feed( def query_items_change_feed( self, *, + feed_range: str, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, - partition_key: PartitionKeyType, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: + """Get a sorted list of items that were changed, in the order in which they were modified. + :keyword str feed_range: The feed range that is used to define the scope. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword start_time:The start time to start processing chang feed items. + :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. Now: Processing change feed from the current time, so only events for all future changes will be retrieved. ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current ("Now") :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] - :keyword partition_key: The partition key that is used to define the scope - (logical partition or a subset of a container) - :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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. @@ -382,23 +384,15 @@ def query_items_change_feed( def query_items_change_feed( self, *, - feed_range: str, + continuation: str, max_item_count: Optional[int] = None, - start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword str continuation: The continuation token retrieved from previous response. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword start_time: The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current ("Now") - :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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. @@ -412,15 +406,21 @@ def query_items_change_feed( def query_items_change_feed( self, *, - continuation: str, max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: - """Get a sorted list of items that were changed, in the order in which they were modified. + """Get a sorted list of items that were changed in the entire container, + in the order in which they were modified, - :keyword str continuation: The continuation token retrieved from previous response. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword start_time:The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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. @@ -496,7 +496,8 @@ def query_items_change_feed( feed_options["maxItemCount"] = args[3] if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = self._set_partition_key(kwargs.get('partition_key')) + change_feed_state_context["partitionKey"] =\ + self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) @@ -1315,7 +1316,7 @@ def read_feed_ranges( """ Obtains a list of feed ranges that can be used to parallelize feed operations. - :param bool force_refresh: + :keyword bool force_refresh: Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. :returns: A list representing the feed ranges in base64 encoded string :rtype: List[str] diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index 4b286d2b82f8..6b96355bb126 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -237,6 +237,7 @@ def create_random_items(container, batch_size): setup["created_db"].delete_container(created_collection.id) + @pytest.mark.skip def test_query_change_feed_with_split(self, setup): created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), PartitionKey(path="/pk"), diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index 886c1ffc1bcc..c1a1e633d62e 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -264,6 +264,7 @@ async def create_random_items(container, batch_size): await setup["created_db"].delete_container(created_collection.id) + @pytest.mark.skip async def test_query_change_feed_with_split_async(self, setup): created_collection = await setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), PartitionKey(path="/pk"), From 59814d7e368420bfdc806fabede945d61af9420d Mon Sep 17 00:00:00 2001 From: annie-mac Date: Thu, 29 Aug 2024 13:30:37 -0700 Subject: [PATCH 17/59] fix mypy --- .../_change_feed/aio/change_feed_iterable.py | 13 +++---- .../_change_feed/change_feed_iterable.py | 4 +- .../azure/cosmos/_cosmos_client_connection.py | 6 ++- .../azure/cosmos/aio/_container.py | 38 +++++++++---------- .../aio/_cosmos_client_connection_async.py | 6 ++- .../azure-cosmos/azure/cosmos/container.py | 2 +- 6 files changed, 36 insertions(+), 33 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 1da44219daa4..6a02d82c0b93 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -21,12 +21,11 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any, Optional, Callable, Tuple, List, AsyncIterator, Awaitable +from typing import Dict, Any, Optional, Callable, Tuple, List, Awaitable, Union from azure.core.async_paging import AsyncPageIterator -from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2, \ - ChangeFeedFetcher +from azure.cosmos._change_feed.aio.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion @@ -60,7 +59,7 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link - self._change_feed_fetcher: Optional[ChangeFeedFetcher] = None + self._change_feed_fetcher: Optional[Union[ChangeFeedFetcherV1, ChangeFeedFetcherV2]] = None if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") @@ -90,9 +89,9 @@ def __init__( async def _unpack( self, - block: AsyncIterator[List[Dict[str, Any]]] - ) -> Tuple[Optional[str], AsyncIterator[List[Dict[str, Any]]]]: - continuation = None + block: List[Dict[str, Any]] + ) -> Tuple[Optional[str], List[Dict[str, Any]]]: + continuation: Optional[str] = None if self._client.last_response_headers: continuation = self._client.last_response_headers.get('etag') diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index c1174db5f93d..a7590b4442df 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -21,7 +21,7 @@ """Iterable change feed results in the Azure Cosmos database service. """ -from typing import Dict, Any, Tuple, List, Optional, Callable, cast +from typing import Dict, Any, Tuple, List, Optional, Callable, cast, Union from azure.core.paging import PageIterator @@ -57,7 +57,7 @@ def __init__( self._options = options self._fetch_function = fetch_function self._collection_link = collection_link - self._change_feed_fetcher: Optional[ChangeFeedFetcher] = None + self._change_feed_fetcher: Optional[Union[ChangeFeedFetcherV1, ChangeFeedFetcherV2]] = None if self._options.get("changeFeedStateContext") is None: raise ValueError("Missing changeFeedStateContext in feed options") 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 010b91c76dd6..49198910b772 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -57,6 +57,7 @@ from ._auth_policy import CosmosBearerTokenCredentialPolicy from ._base import _set_properties_cache from ._change_feed.change_feed_iterable import ChangeFeedIterable +from ._change_feed.change_feed_state import ChangeFeedState from ._constants import _Constants as Constants from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy from ._range_partition_resolver import RangePartitionResolver @@ -3024,8 +3025,9 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: partition_key_range_id ) - if options.get("changeFeedState") is not None: - options.get("changeFeedState").populate_request_headers(self._routing_map_provider, headers) + change_feed_state: Optional[ChangeFeedState] = options.get("changeFeedState") + if change_feed_state is not None: + change_feed_state.populate_request_headers(self._routing_map_provider, headers) result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 4d867a952179..1ba8bf4b2b47 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -499,11 +499,11 @@ def query_items_change_feed( *, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, + partition_key: PartitionKeyType, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: - """Get a sorted list of items that were changed in the entire container, - in the order in which they were modified. + """Get a sorted list of items that were changed, in the order in which they were modified. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword start_time: The start time to start processing chang feed items. @@ -512,6 +512,9 @@ def query_items_change_feed( ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current ("Now") :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] + :keyword partition_key: The partition key that is used to define the scope + (logical partition or a subset of a container) + :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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. @@ -525,14 +528,15 @@ def query_items_change_feed( def query_items_change_feed( self, *, + feed_range: str, max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, - partition_key: PartitionKeyType, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. + :keyword str feed_range: The feed range that is used to define the scope. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. @@ -540,9 +544,6 @@ def query_items_change_feed( ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. By default, it is start from current ("Now") :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] - :keyword partition_key: The partition key that is used to define the scope - (logical partition or a subset of a container) - :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :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. @@ -556,22 +557,15 @@ def query_items_change_feed( def query_items_change_feed( self, *, - feed_range: str, + continuation: str, max_item_count: Optional[int] = None, - start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword str continuation: The continuation token retrieved from previous response. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. - :keyword start_time: The start time to start processing chang feed items. - Beginning: Processing the change feed items from the beginning of the change feed. - Now: Processing change feed from the current time, so only events for all future changes will be retrieved. - ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. - By default, it is start from current ("Now") - :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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. @@ -579,21 +573,28 @@ def query_items_change_feed( :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ + # pylint: enable=line-too-long ... @overload def query_items_change_feed( self, *, - continuation: str, max_item_count: Optional[int] = None, + start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: - """Get a sorted list of items that were changed, in the order in which they were modified. + """Get a sorted list of items that were changed in the entire container, + in the order in which they were modified. - :keyword str continuation: The continuation token retrieved from previous response. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. + :keyword start_time: The start time to start processing chang feed items. + Beginning: Processing the change feed items from the beginning of the change feed. + Now: Processing change feed from the current time, so only events for all future changes will be retrieved. + ~datetime.datetime: processing change feed from a point of time. Provided value will be converted to UTC. + By default, it is start from current ("Now") + :type start_time: Union[~datetime.datetime, Literal["Now", "Beginning"]] :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. @@ -601,7 +602,6 @@ def query_items_change_feed( :returns: An AsyncItemPaged of items (dicts). :rtype: AsyncItemPaged[Dict[str, Any]] """ - # pylint: enable=line-too-long ... @distributed_trace 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 675cbbe9dc69..9e73445e2063 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 @@ -51,6 +51,7 @@ from .._base import _set_properties_cache from .. import documents from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable +from .._change_feed.change_feed_state import ChangeFeedState from .._routing import routing_range from ..documents import ConnectionPolicy, DatabaseAccount from .._constants import _Constants as Constants @@ -2813,8 +2814,9 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: ) headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id) - if options.get("changeFeedState") is not None: - await options.get("changeFeedState").populate_request_headers_async(self._routing_map_provider, headers) + change_feed_state: Optional[ChangeFeedState] = options.get("changeFeedState") + if change_feed_state is not None: + await change_feed_state.populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) if response_hook: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index b8dcc9a45dc0..3f52b43d9994 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -611,7 +611,7 @@ def query_items( # pylint:disable=docstring-missing-param feed_options["populateIndexMetrics"] = populate_index_metrics if partition_key is not None: partition_key_value = self._set_partition_key(partition_key) - if self.__is_prefix_partitionkey(partition_key_value): + if self.__is_prefix_partitionkey(partition_key): kwargs["isPrefixPartitionQuery"] = True properties = self._get_properties() kwargs["partitionKeyDefinition"] = properties["partitionKey"] From 66c3f7bf98f6679e8ced39fda1877e1e1e7b1349 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 4 Sep 2024 11:30:44 -0700 Subject: [PATCH 18/59] fix tests --- sdk/cosmos/azure-cosmos/test/test_change_feed_async.py | 1 - sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py | 2 +- .../azure-cosmos/test/test_container_properties_cache_async.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index c1a1e633d62e..b65694d6c138 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -32,7 +32,6 @@ async def setup(): } yield created_db_data - await test_client.delete_database(config.TEST_DATABASE_ID) await test_client.close() @pytest.mark.cosmosEmulator diff --git a/sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py b/sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py index 6ced2c6d0cd9..fbac47dfb215 100644 --- a/sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py +++ b/sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py @@ -599,7 +599,7 @@ def test_container_recreate_change_feed(self): client.client_connection._CosmosClientConnection__container_properties_cache = copy.deepcopy(old_cache) # Query change feed for the new items - change_feed = list(created_container.query_items_change_feed()) + change_feed = list(created_container.query_items_change_feed(start_time='Beginning')) self.assertEqual(len(change_feed), 2) # Verify that the change feed contains the new items diff --git a/sdk/cosmos/azure-cosmos/test/test_container_properties_cache_async.py b/sdk/cosmos/azure-cosmos/test/test_container_properties_cache_async.py index 88fd6e20cc14..8cf3b9f39ba0 100644 --- a/sdk/cosmos/azure-cosmos/test/test_container_properties_cache_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_container_properties_cache_async.py @@ -612,7 +612,7 @@ async def test_container_recreate_change_feed(self): client.client_connection._CosmosClientConnection__container_properties_cache = copy.deepcopy(old_cache) # Query change feed for the new items - change_feed = [item async for item in created_container.query_items_change_feed()] + change_feed = [item async for item in created_container.query_items_change_feed(start_time='Beginning')] assert len(change_feed) == 2 # Verify that the change feed contains the new items From 3a2e4e159bcf88086e46f8991b9f4d0d2b1185fc Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 4 Sep 2024 22:41:37 -0700 Subject: [PATCH 19/59] add tests --- .../_change_feed/aio/change_feed_fetcher.py | 15 +-- .../_change_feed/aio/change_feed_iterable.py | 2 +- .../_change_feed/change_feed_fetcher.py | 13 ++- .../_change_feed/change_feed_iterable.py | 4 +- .../cosmos/_change_feed/change_feed_state.py | 8 +- ...feed_range_composite_continuation_token.py | 6 +- .../azure-cosmos/test/test_change_feed.py | 55 ----------- .../test/test_change_feed_async.py | 56 ----------- .../test/test_change_feed_split.py | 81 ++++++++++++++++ .../test/test_change_feed_split_async.py | 94 +++++++++++++++++++ 10 files changed, 201 insertions(+), 133 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/test/test_change_feed_split.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index 376fd0fac397..c1db30d41fa2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -88,7 +88,8 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. self._change_feed_state.apply_server_response_continuation( - response_headers.get(continuation_key)) + cast(str, response_headers.get(continuation_key)), + (True if fetched_items else False)) if fetched_items: break @@ -164,9 +165,13 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. + + self._change_feed_state.apply_server_response_continuation( + cast(str, response_headers.get(continuation_key)), + (True if fetched_items else False)) + if fetched_items: - self._change_feed_state.apply_server_response_continuation( - cast(str, response_headers.get(continuation_key))) + self._change_feed_state._continuation._move_to_next_token() response_headers[continuation_key] = self._get_base64_encoded_continuation() break @@ -175,10 +180,6 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # so we will retry with the new continuation token # 2. if the feed range of the changeFeedState span multiple physical partitions # then we will read from the next feed range until we have looped through all physical partitions - self._change_feed_state.apply_not_modified_response() - self._change_feed_state.apply_server_response_continuation( - cast(str, response_headers.get(continuation_key))) - if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): response_headers[continuation_key] = self._get_base64_encoded_continuation() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 6a02d82c0b93..8d7b4eaf699b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -85,7 +85,7 @@ def __init__( self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context - super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) + super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) # type: ignore[arg-type] async def _unpack( self, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index 2417eff46259..846861c704cc 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -87,7 +87,8 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # In change feed queries, the continuation token is always populated. The hasNext() test is whether # there is any items in the response or not. self._change_feed_state.apply_server_response_continuation( - cast(str, response_headers.get(continuation_key))) + cast(str, response_headers.get(continuation_key)), + (True if fetched_items else False)) if fetched_items: break @@ -156,9 +157,11 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: continuation_key = http_constants.HttpHeaders.ETag # In change feed queries, the continuation token is always populated. + self._change_feed_state.apply_server_response_continuation( + cast(str, response_headers.get(continuation_key)), + (True if fetched_items else False)) + if fetched_items: - self._change_feed_state.apply_server_response_continuation( - cast(str, response_headers.get(continuation_key))) self._change_feed_state._continuation._move_to_next_token() response_headers[continuation_key] = self._get_base64_encoded_continuation() break @@ -168,10 +171,6 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # so we will retry with the new continuation token # 2. if the feed range of the changeFeedState span multiple physical partitions # then we will read from the next feed range until we have looped through all physical partitions - self._change_feed_state.apply_not_modified_response() - self._change_feed_state.apply_server_response_continuation( - cast(str, response_headers.get(continuation_key))) - if (self._change_feed_state._change_feed_start_from.version == ChangeFeedStartFromType.POINT_IN_TIME and is_s_time_first_fetch): response_headers[continuation_key] = self._get_base64_encoded_continuation() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index a7590b4442df..00193ec3da72 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -25,7 +25,7 @@ from azure.core.paging import PageIterator -from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2, ChangeFeedFetcher +from azure.cosmos._change_feed.change_feed_fetcher import ChangeFeedFetcherV1, ChangeFeedFetcherV2 from azure.cosmos._change_feed.change_feed_state import ChangeFeedState, ChangeFeedStateVersion @@ -82,7 +82,7 @@ def __init__( self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context - super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) + super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) # type: ignore[arg-type] def _unpack(self, block: List[Dict[str, Any]]) -> Tuple[Optional[str], List[Dict[str, Any]]]: continuation: Optional[str] = None diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index d7c52bf89b88..46dd1afddcfe 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -72,7 +72,7 @@ async def populate_request_headers_async( pass @abstractmethod - def apply_server_response_continuation(self, continuation: str) -> None: + def apply_server_response_continuation(self, continuation: str, has_modified_response: bool) -> None: pass @staticmethod @@ -170,7 +170,7 @@ def populate_feed_options(self, feed_options: Dict[str, Any]) -> None: if self._partition_key is not None: feed_options["partitionKey"] = self._partition_key - def apply_server_response_continuation(self, continuation: str) -> None: + def apply_server_response_continuation(self, continuation: str, has_modified_response) -> None: self._continuation = continuation class ChangeFeedStateV2(ChangeFeedState): @@ -322,8 +322,8 @@ async def handle_feed_range_gone_async( resource_link: str) -> None: await self._continuation.handle_feed_range_gone_async(routing_provider, resource_link) - def apply_server_response_continuation(self, continuation: str) -> None: - self._continuation.apply_server_response_continuation(continuation) + def apply_server_response_continuation(self, continuation: str, has_modified_response: bool) -> None: + self._continuation.apply_server_response_continuation(continuation, has_modified_response) def should_retry_on_not_modified_response(self) -> bool: return self._continuation.should_retry_on_not_modified_response() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py index fc9b94f27eef..f5967b6bf34b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py @@ -159,8 +159,12 @@ def _move_to_next_token(self) -> None: self._continuation.append(first_composition_token) self._current_token = self._continuation[0] - def apply_server_response_continuation(self, etag) -> None: + def apply_server_response_continuation(self, etag, has_modified_response: bool) -> None: self._current_token.update_token(etag) + if has_modified_response: + self._initial_no_result_range = None + else: + self.apply_not_modified_response() def apply_not_modified_response(self) -> None: if self._initial_no_result_range is None: diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index 6b96355bb126..bd6e1b7c4faa 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -237,61 +237,6 @@ def create_random_items(container, batch_size): setup["created_db"].delete_container(created_collection.id) - @pytest.mark.skip - def test_query_change_feed_with_split(self, setup): - created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk"), - offer_throughput=400) - - # initial change feed query returns empty result - query_iterable = created_collection.query_items_change_feed(start_time="Beginning") - iter_list = list(query_iterable) - assert len(iter_list) == 0 - continuation = created_collection.client_connection.last_response_headers['etag'] - assert continuation != '' - - # create one doc and make sure change feed query can return the document - document_definition = {'pk': 'pk', 'id': 'doc1'} - created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed(continuation=continuation) - iter_list = list(query_iterable) - assert len(iter_list) == 1 - continuation = created_collection.client_connection.last_response_headers['etag'] - - print("Triggering a split in test_query_change_feed_with_split") - created_collection.replace_throughput(11000) - print("changed offer to 11k") - print("--------------------------------") - print("Waiting for split to complete") - start_time = time.time() - - while True: - offer = created_collection.get_throughput() - if offer.properties['content'].get('isOfferReplacePending', False): - if time.time() - start_time > 60 * 25: # timeout test at 25 minutes - unittest.skip("Partition split didn't complete in time.") - else: - print("Waiting for split to complete") - time.sleep(60) - else: - break - - print("Split in test_query_change_feed_with_split has completed") - print("creating few more documents") - new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] - expected_ids = ['doc2', 'doc3', 'doc4'] - for document in new_documents: - created_collection.create_item(body=document) - - query_iterable = created_collection.query_items_change_feed(continuation=continuation) - it = query_iterable.__iter__() - actual_ids = [] - for item in it: - actual_ids.append(item['id']) - - assert actual_ids == expected_ids - setup["created_db"].delete_container(created_collection.id) - def test_query_change_feed_with_multi_partition(self, setup): created_collection = setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), PartitionKey(path="/pk"), diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index b65694d6c138..2d574f2f8ee7 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -1,7 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import time import unittest import uuid from asyncio import sleep @@ -263,61 +262,6 @@ async def create_random_items(container, batch_size): await setup["created_db"].delete_container(created_collection.id) - @pytest.mark.skip - async def test_query_change_feed_with_split_async(self, setup): - created_collection = await setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), - PartitionKey(path="/pk"), - offer_throughput=400) - - # initial change feed query returns empty result - query_iterable = created_collection.query_items_change_feed(start_time="Beginning") - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 0 - continuation = created_collection.client_connection.last_response_headers['etag'] - assert continuation != '' - - # create one doc and make sure change feed query can return the document - document_definition = {'pk': 'pk', 'id': 'doc1'} - await created_collection.create_item(body=document_definition) - query_iterable = created_collection.query_items_change_feed(continuation=continuation) - iter_list = [item async for item in query_iterable] - assert len(iter_list) == 1 - continuation = created_collection.client_connection.last_response_headers['etag'] - - print("Triggering a split in test_query_change_feed_with_split") - await created_collection.replace_throughput(11000) - print("changed offer to 11k") - print("--------------------------------") - print("Waiting for split to complete") - start_time = time.time() - - while True: - offer = await created_collection.get_throughput() - if offer.properties['content'].get('isOfferReplacePending', False): - if time.time() - start_time > 60 * 25: # timeout test at 25 minutes - unittest.skip("Partition split didn't complete in time.") - else: - print("Waiting for split to complete") - time.sleep(60) - else: - break - - print("Split in test_query_change_feed_with_split has completed") - print("creating few more documents") - new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] - expected_ids = ['doc2', 'doc3', 'doc4'] - for document in new_documents: - await created_collection.create_item(body=document) - - query_iterable = created_collection.query_items_change_feed(continuation=continuation) - it = query_iterable.__aiter__() - actual_ids = [] - async for item in it: - actual_ids.append(item['id']) - - assert actual_ids == expected_ids - setup["created_db"].delete_container(created_collection.id) - async def test_query_change_feed_with_multi_partition_async(self, setup): created_collection = await setup["created_db"].create_container("change_feed_test_" + str(uuid.uuid4()), PartitionKey(path="/pk"), diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_split.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_split.py new file mode 100644 index 000000000000..8ecb7da9cff3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_split.py @@ -0,0 +1,81 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import time +import unittest +import uuid + +import azure.cosmos.cosmos_client as cosmos_client +import test_config +from azure.cosmos import DatabaseProxy, PartitionKey + + +class TestPartitionSplitChangeFeed(unittest.TestCase): + database: DatabaseProxy = None + client: cosmos_client.CosmosClient = None + configs = test_config.TestConfig + host = configs.host + masterKey = configs.masterKey + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) + cls.database = cls.client.get_database_client(cls.TEST_DATABASE_ID) + + def test_query_change_feed_with_split(self): + created_collection = self.database.create_container("change_feed_split_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + + # initial change feed query returns empty result + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") + iter_list = list(query_iterable) + assert len(iter_list) == 0 + continuation = created_collection.client_connection.last_response_headers['etag'] + assert continuation != '' + + # create one doc and make sure change feed query can return the document + document_definition = {'pk': 'pk', 'id': 'doc1'} + created_collection.create_item(body=document_definition) + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + iter_list = list(query_iterable) + assert len(iter_list) == 1 + continuation = created_collection.client_connection.last_response_headers['etag'] + + print("Triggering a split in test_query_change_feed_with_split") + created_collection.replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = created_collection.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in test_query_change_feed_with_split has completed") + print("creating few more documents") + new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc2', 'doc3', 'doc4'] + for document in new_documents: + created_collection.create_item(body=document) + + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + it = query_iterable.__iter__() + actual_ids = [] + for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + self.database.delete_container(created_collection.id) + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py new file mode 100644 index 000000000000..60f7b2810884 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py @@ -0,0 +1,94 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import time +import unittest +import uuid + +import test_config +from azure.cosmos import PartitionKey +from azure.cosmos.aio import CosmosClient, DatabaseProxy + + +class TestPartitionSplitChangeFeedAsync(unittest.IsolatedAsyncioTestCase): + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + connectionPolicy = test_config.TestConfig.connectionPolicy + + client: CosmosClient = None + created_database: DatabaseProxy = None + + TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + if (cls.masterKey == '[YOUR_KEY_HERE]' or + cls.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + + async def asyncSetUp(self): + self.client = CosmosClient(self.host, self.masterKey) + self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID) + + async def tearDown(self): + await self.client.close() + + async def test_query_change_feed_with_split_async(self): + created_collection = await self.created_database.create_container("change_feed_test_" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + + # initial change feed query returns empty result + query_iterable = created_collection.query_items_change_feed(start_time="Beginning") + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 0 + continuation = created_collection.client_connection.last_response_headers['etag'] + assert continuation != '' + + # create one doc and make sure change feed query can return the document + document_definition = {'pk': 'pk', 'id': 'doc1'} + await created_collection.create_item(body=document_definition) + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + iter_list = [item async for item in query_iterable] + assert len(iter_list) == 1 + continuation = created_collection.client_connection.last_response_headers['etag'] + + print("Triggering a split in test_query_change_feed_with_split") + await created_collection.replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = await created_collection.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in test_query_change_feed_with_split has completed") + print("creating few more documents") + new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}] + expected_ids = ['doc2', 'doc3', 'doc4'] + for document in new_documents: + await created_collection.create_item(body=document) + + query_iterable = created_collection.query_items_change_feed(continuation=continuation) + it = query_iterable.__aiter__() + actual_ids = [] + async for item in it: + actual_ids.append(item['id']) + + assert actual_ids == expected_ids + self.created_database.delete_container(created_collection.id) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file From 0883dac536e9940fe281b64e0ffb8bc33c21e304 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Thu, 5 Sep 2024 08:13:06 -0700 Subject: [PATCH 20/59] fix pylint --- .../azure/cosmos/_change_feed/aio/change_feed_fetcher.py | 4 ++-- .../azure/cosmos/_change_feed/aio/change_feed_iterable.py | 5 ++++- .../azure/cosmos/_change_feed/change_feed_fetcher.py | 4 ++-- .../azure/cosmos/_change_feed/change_feed_iterable.py | 5 ++++- 4 files changed, 12 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py index c1db30d41fa2..d997360e4c41 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_fetcher.py @@ -89,7 +89,7 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # there is any items in the response or not. self._change_feed_state.apply_server_response_continuation( cast(str, response_headers.get(continuation_key)), - (True if fetched_items else False)) + bool(fetched_items)) if fetched_items: break @@ -168,7 +168,7 @@ async def fetch_change_feed_items(self) -> List[Dict[str, Any]]: self._change_feed_state.apply_server_response_continuation( cast(str, response_headers.get(continuation_key)), - (True if fetched_items else False)) + bool(fetched_items)) if fetched_items: self._change_feed_state._continuation._move_to_next_token() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py index 8d7b4eaf699b..3f73050dfc7a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/aio/change_feed_iterable.py @@ -85,7 +85,10 @@ def __init__( self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context - super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) # type: ignore[arg-type] + super(ChangeFeedIterable, self).__init__( + self._fetch_next, + self._unpack, # type: ignore[arg-type] + continuation_token=continuation_token) async def _unpack( self, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py index 846861c704cc..c3ff6472af28 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_fetcher.py @@ -88,7 +88,7 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # there is any items in the response or not. self._change_feed_state.apply_server_response_continuation( cast(str, response_headers.get(continuation_key)), - (True if fetched_items else False)) + bool(fetched_items)) if fetched_items: break @@ -159,7 +159,7 @@ def fetch_change_feed_items(self) -> List[Dict[str, Any]]: # In change feed queries, the continuation token is always populated. self._change_feed_state.apply_server_response_continuation( cast(str, response_headers.get(continuation_key)), - (True if fetched_items else False)) + bool(fetched_items)) if fetched_items: self._change_feed_state._continuation._move_to_next_token() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py index 00193ec3da72..bd37b60926cf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_iterable.py @@ -82,7 +82,10 @@ def __init__( self._validate_change_feed_state_context(change_feed_state_context) self._options["changeFeedStateContext"] = change_feed_state_context - super(ChangeFeedIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token) # type: ignore[arg-type] + super(ChangeFeedIterable, self).__init__( + self._fetch_next, + self._unpack, # type: ignore[arg-type] + continuation_token=continuation_token) def _unpack(self, block: List[Dict[str, Any]]) -> Tuple[Optional[str], List[Dict[str, Any]]]: continuation: Optional[str] = None From 195c47cd3783db43a5b4ef062de6a898887fe41e Mon Sep 17 00:00:00 2001 From: annie-mac Date: Fri, 6 Sep 2024 09:01:30 -0700 Subject: [PATCH 21/59] fix and resolve comments --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | 4 +--- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 4 +--- sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py | 2 +- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 1ba8bf4b2b47..d7d66738b4ee 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -135,9 +135,7 @@ async def _set_partition_key( return _return_undefined_or_empty_partition_key(await self.is_system_key) return cast(Union[str, int, float, bool, List[Union[str, int, float, bool]]], partition_key) - async def _get_epk_range_for_partition_key( - self, - partition_key_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> Range: # pylint: disable=line-too-long + async def _get_epk_range_for_partition_key(self, partition_key_value: PartitionKeyType) -> Range: container_properties = await self._get_properties() partition_key_definition = container_properties["partitionKey"] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 3f52b43d9994..e6a6ac7b36b9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -130,9 +130,7 @@ def _set_partition_key( 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) - def _get_epk_range_for_partition_key( - self, - partition_key_value: Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]]) -> Range: # pylint: disable=line-too-long + def _get_epk_range_for_partition_key( self, partition_key_value: PartitionKeyType) -> Range: container_properties = self._get_properties() partition_key_definition = container_properties["partitionKey"] partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py index 5870e7519e8b..7fa093aa15e1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py @@ -284,7 +284,7 @@ def _is_prefix_partition_key( 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": return False - if isinstance(partition_key, list) and len(self.path) == len(partition_key): + if isinstance(partition_key, list) and len(self['paths']) == len(partition_key): return False return True From 246b1be3cd8e2163b2853ad9595ae4d8f822ca4c Mon Sep 17 00:00:00 2001 From: annie-mac Date: Fri, 6 Sep 2024 14:32:40 -0700 Subject: [PATCH 22/59] fix and resolve comments --- .../azure-cosmos/test/test_change_feed.py | 19 +++++-------------- .../test/test_change_feed_async.py | 12 +----------- .../azure-cosmos/test/test_vector_policy.py | 2 +- .../test/test_vector_policy_async.py | 2 +- 4 files changed, 8 insertions(+), 27 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index bd6e1b7c4faa..01e2dc21ddb6 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -1,7 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import time import unittest import uuid from datetime import datetime, timedelta, timezone @@ -13,21 +12,21 @@ import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.exceptions as exceptions import test_config -from azure.cosmos import DatabaseProxy from azure.cosmos.partition_key import PartitionKey @pytest.fixture(scope="class") def setup(): - if (TestChangeFeed.masterKey == '[YOUR_KEY_HERE]' or - TestChangeFeed.host == '[YOUR_ENDPOINT_HERE]'): + config = test_config.TestConfig() + if (config.masterKey == '[YOUR_KEY_HERE]' or + config.host == '[YOUR_ENDPOINT_HERE]'): raise Exception( "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(test_config.TestConfig.host, test_config.TestConfig.masterKey), + test_client = cosmos_client.CosmosClient(config.host, config.masterKey), return { - "created_db": test_client[0].get_database_client(TestChangeFeed.TEST_DATABASE_ID) + "created_db": test_client[0].get_database_client(config.TEST_DATABASE_ID) } @pytest.mark.cosmosEmulator @@ -36,14 +35,6 @@ def setup(): class TestChangeFeed: """Test to ensure escaping of non-ascii characters from partition key""" - created_db: DatabaseProxy = None - client: cosmos_client.CosmosClient = None - config = test_config.TestConfig - host = config.host - masterKey = config.masterKey - connectionPolicy = config.connectionPolicy - TEST_DATABASE_ID = config.TEST_DATABASE_ID - def test_get_feed_ranges(self, setup): created_collection = setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index 2d574f2f8ee7..2ef61ee5c8a3 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -12,7 +12,7 @@ import azure.cosmos.exceptions as exceptions import test_config -from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy +from azure.cosmos.aio import CosmosClient from azure.cosmos.partition_key import PartitionKey @@ -39,16 +39,6 @@ async def setup(): class TestChangeFeedAsync: """Test to ensure escaping of non-ascii characters from partition key""" - created_db: DatabaseProxy = None - created_container: ContainerProxy = None - client: CosmosClient = None - config = test_config.TestConfig - TEST_CONTAINER_ID = config.TEST_MULTI_PARTITION_CONTAINER_ID - TEST_DATABASE_ID = config.TEST_DATABASE_ID - host = config.host - masterKey = config.masterKey - connectionPolicy = config.connectionPolicy - async def test_get_feed_ranges(self, setup): created_collection = await setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/test/test_vector_policy.py b/sdk/cosmos/azure-cosmos/test/test_vector_policy.py index da0aeb8ec6a8..e44f8f21c5fd 100644 --- a/sdk/cosmos/azure-cosmos/test/test_vector_policy.py +++ b/sdk/cosmos/azure-cosmos/test/test_vector_policy.py @@ -163,7 +163,7 @@ def test_fail_replace_vector_indexing_policy(self): pytest.fail("Container replace should have failed for indexing policy.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert "Vector Indexing Policy cannot be changed in Collection Replace" in e.http_error_message + assert "vector indexing policy cannot be modified in Collection Replace" in e.http_error_message self.test_db.delete_container(container_id) def test_fail_create_vector_embedding_policy(self): diff --git a/sdk/cosmos/azure-cosmos/test/test_vector_policy_async.py b/sdk/cosmos/azure-cosmos/test/test_vector_policy_async.py index 19dd48268417..71c4997e3179 100644 --- a/sdk/cosmos/azure-cosmos/test/test_vector_policy_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_vector_policy_async.py @@ -173,7 +173,7 @@ async def test_fail_replace_vector_indexing_policy_async(self): pytest.fail("Container replace should have failed for indexing policy.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert "Vector Indexing Policy cannot be changed in Collection Replace" in e.http_error_message + assert "vector indexing policy cannot be modified in Collection Replace" in e.http_error_message await self.test_db.delete_container(container_id) async def test_fail_create_vector_embedding_policy_async(self): From 10fe3875df5a85d0d8ec1227e55ecefd1b5ab596 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 9 Sep 2024 10:23:39 -0700 Subject: [PATCH 23/59] Added isSubsetFeedRange logic --- .../azure/cosmos/_cosmos_client_connection.py | 5 +- .../azure/cosmos/_request_context.py | 13 ++ .../azure-cosmos/azure/cosmos/container.py | 35 ++++- .../samples/merge_session_tokens.py | 121 ++++++++++++++++++ .../azure-cosmos/test/test_change_feed.py | 1 - .../azure-cosmos/test/test_feed_range.py | 90 +++++++++++++ .../test/test_session_token_helpers.py | 16 ++- 7 files changed, 268 insertions(+), 13 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py create mode 100644 sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_feed_range.py 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 f8ee1575bfb6..2949bc295dd2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -46,7 +46,7 @@ from azure.core.pipeline.transport import HttpRequest, \ HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import -from . import _base as base +from . import _base as base, _request_context from . import _global_endpoint_manager as global_endpoint_manager from . import _query_iterable as query_iterable from . import _runtime_constants as runtime_constants @@ -2620,6 +2620,7 @@ def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) + self.last_response_headers = _request_context.add_request_context(last_response_headers, options) if response_hook: response_hook(last_response_headers, result) return result @@ -3328,7 +3329,7 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s self.__container_properties_cache[collection_link] = _set_properties_cache(container) return partition_key_definition - def MergeSessionTokens(self, feed_ranges_to_session_tokens, target_feed_range): + def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_range): comparison_session_token = '' comparison_pk_range_id = '' for (pk, session_token) in feed_ranges_to_session_tokens: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py new file mode 100644 index 000000000000..4dbbb42085f7 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py @@ -0,0 +1,13 @@ +session_token_header = 'x-ms-session-token' +request_context_header = 'request-context' +session_token_request_context = 'session-token' +feed_range_request_context = 'feed-range' + + +def add_request_context(last_response_headers, options) -> dict[str, str]: + request_context = {} + request_context[session_token_request_context] = last_response_headers[session_token_header] + if 'partitionKey' in options: + last_response_headers['partitionKey'] + last_response_headers[request_context_header] = request_context + return last_response_headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 9e057675662b..d26b1e30bca7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -38,6 +38,7 @@ GenerateGuidId, _set_properties_cache ) +from ._change_feed.feed_range import FeedRange, FeedRangeEpk from ._cosmos_client_connection import CosmosClientConnection from ._routing.routing_range import Range, partition_key_range_to_range_string from .offer import Offer, ThroughputProperties @@ -1333,19 +1334,43 @@ def read_feed_ranges( return [partition_key_range_to_range_string(partitionKeyRange) for partitionKeyRange in partition_key_ranges] - def merge_session_tokens(self, + def get_updated_session_token(self, feed_ranges_to_session_tokens: List, target_feed_range: Range ) -> "Session Token": - """Merge session tokens from different clients to figure out which is the most up to date for a specific - feed range. The feed range can be obtained from the response from crud operations. + """Gets the best session token from the list of session token and feed range tuples + to figure out which is the most up to date for a specific + feed range. The feed range can be obtained from a response from any crud operation. This should only be used if maintaining own session token or else the sdk will keep track of session token. - :param feed_ranges_to_session_tokens: list of partition key and session token tuples. + :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple(str, Range)] :param target_feed_range: feed range to get most up to date session token. :type target_feed_range: Range :returns: a session token :rtype: str """ - self.client_connection.MergeSessionTokens(feed_ranges_to_session_tokens, target_feed_range) + self.client_connection._get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + + def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: + """Gets the feed range for a given partition key. + :param partition_key: partition key to get feed range. + :type partition_key: PartitionKey + :returns: a feed range + :rtype: Range + """ + return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) + + def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + """Checks if child feed range is a subset of parent feed range. + :param parent_feed_range: left feed range + :type parent_feed_range: Range + :param child_feed_range: right feed range + :type child_feed_range: Range + :returns: a boolean indicating if child feed range is a subset of parent feed range + :rtype: bool + """ + normalized_parent_range = parent_feed_range.get_normalized_range() + normalized_child_range = child_feed_range.get_normalized_range() + return normalized_parent_range.contains(normalized_child_range.min) and \ + (normalized_parent_range.contains(normalized_child_range.max) or normalized_parent_range.max == normalized_child_range.max) diff --git a/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py b/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py new file mode 100644 index 000000000000..726cbbf8f3f4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py @@ -0,0 +1,121 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- +import os + +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +from azure.cosmos.aio import CosmosClient +from azure.cosmos.http_constants import StatusCodes +from azure.cosmos.partition_key import PartitionKey +import datetime + +import config + +# ---------------------------------------------------------------------------------------------------------- +# Prerequisites - +# +# 1. An Azure Cosmos account - +# https:#azure.microsoft.com/en-us/documentation/articles/documentdb-create-account/ +# +# 2. Microsoft Azure Cosmos PyPi package - +# https://pypi.python.org/pypi/azure-cosmos/ +# ---------------------------------------------------------------------------------------------------------- +# Sample - demonstrates how to merge session tokens using different strategies for storing the session token +# ---------------------------------------------------------------------------------------------------------- + + + +# url = os.environ["ACCOUNT_URI"] +# key = os.environ["ACCOUNT_KEY"] +# client = CosmosClient(url, key) +# # Create a database in the account using the CosmosClient, +# # specifying that the operation shouldn't throw an exception +# # if a database with the given ID already exists. +# # [START create_database] +# database_name = "testDatabase" +# try: +# database = client.create_database(id=database_name) +# except exceptions.CosmosResourceExistsError: +# database = client.get_database_client(database=database_name) +# # [END create_database] +# +# # Create a container, handling the exception if a container with the +# # same ID (name) already exists in the database. +# # [START create_container] +# container_name = "products" +# try: +# container = database.create_container( +# id=container_name, partition_key=PartitionKey(path="/productName") +# ) +# except exceptions.CosmosResourceExistsError: +# container = database.get_container_client(container_name) +# # [END create_container] +# +# # This would be happening through different clients +# # Using physical partition model for read operations +# cache = {} +# session_token = "" +# feed_range = container.feed_range_for_logical_partition(logical_pk) +# for stored_feed_range, stored_session_token in cache: +# if is_feed_range_subset(stored_feed_range, feed_range): +# session_token = stored_session_token +# read_item = container.read_item(doc_to_read, logical_pk, session_token) +# +# # the feed range returned in the request context will correspond to the logical partition key +# logical_pk_feed_range = container.client_connection.last_response_headers["request-context"]["feed-range"] +# session_token = container.client_connection.last_response_headers["request-context"]["session-token"] +# feed_ranges_and_session_tokens = [] +# +# # Get feed ranges for physical partitions +# container_feed_ranges = container.read_feed_ranges() +# target_feed_range = "" +# +# # which feed range maps to the logical pk from the operation +# for feed_range in container_feed_ranges: +# if is_feed_range_subset(feed_range, logical_pk_feed_range): +# target_feed_range = feed_range +# break +# for cached_feed_range, cached_session_token in cache: +# feed_ranges_and_session_tokens.append((cached_feed_range, cached_session_token)) +# # Add the target feed range and session token from the operation +# feed_ranges_and_session_tokens.append((target_feed_range, session_token)) +# cache[feed_range] = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) +# +# +# +# # Different ways of storing the session token and how to get most updated session token +# +# # ---------------------1. using logical partition key --------------------------------------------------- +# # could also use the one stored from the responses headers +# target_feed_range = container.feed_range_for_logical_partition(logical_pk) +# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) +# # ---------------------2. using artificial feed range ---------------------------------------------------- +# # Get four artificial feed ranges +# container_feed_ranges = container.read_feed_ranges(4) +# +# pk_feed_range = container.feed_range_for_logical_partition(logical_pk) +# target_feed_range = "" +# # which feed range maps to the logical pk from the operation +# for feed_range in container_feed_ranges: +# if is_feed_range_subset(feed_range, pk_feed_range): +# target_feed_range = feed_range +# break +# +# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) +# # ---------------------3. using physical partitions ----------------------------------------------------- +# # Get feed ranges for physical partitions +# container_feed_ranges = container.read_feed_ranges() +# +# pk_feed_range = container.feed_range_for_logical_partition(logical_pk) +# target_feed_range = "" +# # which feed range maps to the logical pk from the operation +# for feed_range in container_feed_ranges: +# if is_feed_range_subset(feed_range, pk_feed_range): +# target_feed_range = feed_range +# break +# +# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) +# # ------------------------------------------------------------------------------------------------------ \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index bd6e1b7c4faa..292aa7850f78 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -1,7 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import time import unittest import uuid from datetime import datetime, timedelta, timezone diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py new file mode 100644 index 000000000000..621b353f7987 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -0,0 +1,90 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest + +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.partition_key as partition_key +import test_config +from azure.cosmos._change_feed.feed_range import FeedRangeEpk +from azure.cosmos._routing.routing_range import Range +from test.test_config import TestConfig + + + +@pytest.fixture(scope="class") +def setup(): + if (TestFeedRange.masterKey == '[YOUR_KEY_HERE]' or + TestFeedRange.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.credential), + created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID) + return { + "created_db": created_db, + "created_collection": created_db.get_container_client(TestFeedRange.TEST_CONTAINER_ID) + } + +@pytest.mark.cosmosEmulator +@pytest.mark.unittest +@pytest.mark.usefixtures("setup") +class TestFeedRange: + """Tests to verify methods for operations on feed ranges + """ + + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID + TEST_CONTAINER_ID = test_config.TestConfig.TEST_MULTI_PARTITION_CONTAINER_ID + + + def test_partition_key_to_feed_range(self, setup): + created_container = setup["created_db"].create_container( + id='container_' + str(uuid.uuid4()), + partition_key=partition_key.PartitionKey(path="/id") + ) + feed_range = created_container.feed_range_from_partition_key("1") + print(feed_range.get_normalized_range()) + setup["created_db"].delete_container(created_container) + + + test_ranges = [(Range("", "FFFFFFFFFFFFFFFF", True, False), + Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + True), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + Range("", "FFFFFFFFFFFFFFFF", True, False), + False), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + Range("", "5FFFFFFFFFFFFFFF", True, False), + False), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + True), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", False, True), + Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + False), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + False), + (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + Range("", "2FFFFFFFFFFFFFFF", True, False), + False)] + + @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_ranges) + def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset): + epk_parent_feed_range = FeedRangeEpk(parent_feed_range) + epk_child_feed_range = FeedRangeEpk(child_feed_range) + assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) == is_subset + + def test_feed_range_is_subset_from_pk(self, setup): + epk_parent_feed_range = FeedRangeEpk(Range("", "FFFFFFFFFFFFFFFF", True, False)) + epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") + assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 0a9c5c8e4ec0..748977a3739a 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -2,12 +2,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. import unittest +from time import sleep import pytest import azure.cosmos.cosmos_client as cosmos_client import test_config -from azure.cosmos import DatabaseProxy, ThroughputProperties +from azure.cosmos import DatabaseProxy, ThroughputProperties, PartitionKey +from test.test_config import TestConfig @pytest.mark.cosmosEmulator @@ -34,17 +36,21 @@ def setUpClass(cls): "'masterKey' and 'host' at the top of this class to run the " "tests.") - cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) + cls.client = cosmos_client.CosmosClient(cls.host, TestConfig.credential) cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) - cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID, ThroughputProperties(11000)) + cls.created_collection = cls.created_db.create_container("TestColl", PartitionKey("/mypk"), offer_throughput=ThroughputProperties(11000)) # have a test for split, merge, different versions, compound session tokens, hpk, and for normal case for several session tokens # have tests for all the scenarios in the testing plan + # should try once with 35000 session tokens def test_session_token_merge(self): + sleep(5) for i in range(10): - self.created_collection.create_item(body={'id': 'doc' + str(i)}) - self.created_collection.read_all_items() + self.created_collection.create_item(body={'id': 'doc' + str(i), 'mypk': str(i)}) session_token = self.created_collection.client_connection.last_response_headers['x-ms-session-token'] + self.created_collection.read_all_items(session_token=session_token) + session_token = self.created_collection.client_connection.last_response_headers['x-ms-session-token'] + print("\n" + "Session Token: " + session_token) self.client.delete_database(self.TEST_DATABASE_ID) From 6498311d59b119da8f3394c02f1820196bedf567 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 11 Sep 2024 16:32:21 -0700 Subject: [PATCH 24/59] Added request context to crud operations, session token helpers --- .../azure/cosmos/_change_feed/feed_range.py | 17 +++- .../azure/cosmos/_cosmos_client_connection.py | 98 ++++++++++++++----- .../azure/cosmos/_request_context.py | 13 --- .../azure/cosmos/_routing/routing_range.py | 21 ++++ .../azure-cosmos/azure/cosmos/container.py | 63 ++++++++---- .../azure-cosmos/test/test_feed_range.py | 37 +++---- .../azure-cosmos/test/test_request_context.py | 65 ++++++++++++ 7 files changed, 239 insertions(+), 75 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py create mode 100644 sdk/cosmos/azure-cosmos/test/test_request_context.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py index b4f731f2c2ef..293178436cdf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py @@ -25,7 +25,7 @@ from abc import ABC, abstractmethod from typing import Union, List, Dict, Any -from azure.cosmos._routing.routing_range import Range +from azure.cosmos._routing.routing_range import Range, PartitionKeyRange from azure.cosmos.partition_key import _Undefined, _Empty @@ -39,6 +39,21 @@ def get_normalized_range(self) -> Range: def to_dict(self) -> Dict[str, Any]: pass + + + @staticmethod + def _get_range(pk_ranges) -> Range: + max_range = pk_ranges[0][PartitionKeyRange.MaxExclusive] + min_range = pk_ranges[0][PartitionKeyRange.MinInclusive] + for i in range(1, len(pk_ranges)): + pk_ranges_min = pk_ranges[i][PartitionKeyRange.MinInclusive] + pk_ranges_max = pk_ranges[i][PartitionKeyRange.MaxExclusive] + if pk_ranges_min < min_range: + min_range = pk_ranges_min + if pk_ranges_max > max_range: + max_range = pk_ranges_max + return Range(min_range, max_range, True, False) + class FeedRangePartitionKey(FeedRange): type_property_name = "PK" 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 2949bc295dd2..b9d2e6d7b0ed 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -30,6 +30,8 @@ from urllib3.util.retry import Retry from azure.core import PipelineClient + +from ._change_feed.feed_range import FeedRange from ._vector_session_token import VectorSessionToken from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged @@ -46,7 +48,7 @@ from azure.core.pipeline.transport import HttpRequest, \ HttpResponse # pylint: disable=no-legacy-azure-core-http-response-import -from . import _base as base, _request_context +from . import _base as base from . import _global_endpoint_manager as global_endpoint_manager from . import _query_iterable as query_iterable from . import _runtime_constants as runtime_constants @@ -1259,6 +1261,7 @@ def CreateItem( self, database_or_container_link: str, document: Dict[str, Any], + request_context: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -1289,7 +1292,9 @@ def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) - return self.Create(document, path, "docs", collection_id, None, options, **kwargs) + request_context["partitionKey"] = options["partitionKey"] + result = self.Create(document, path, "docs", collection_id, None, options, **kwargs) + return result def UpsertItem( self, @@ -1971,6 +1976,7 @@ def ReplaceItem( self, document_link: str, new_document: Dict[str, Any], + request_context: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2620,7 +2626,6 @@ def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) - self.last_response_headers = _request_context.add_request_context(last_response_headers, options) if response_hook: response_hook(last_response_headers, result) return result @@ -3330,26 +3335,67 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s return partition_key_definition def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_range): - comparison_session_token = '' - comparison_pk_range_id = '' - for (pk, session_token) in feed_ranges_to_session_tokens: - if pk == target_feed_range: - token_pairs = session_token.split(":") - comparison_pk_range_id = token_pairs[0] - comparison_session_token = VectorSessionToken.create(token_pairs[1]) - break - for (pk, session_token) in feed_ranges_to_session_tokens: - token_pairs = session_token.split(":") - pk_range_id = token_pairs[0] - vector_session_token = VectorSessionToken.create(token_pairs[1]) - # This should not be necessary - # if pk_range_id == comparison_pk_range_id: - # comparison_session_token = comparison_session_token.merge(vector_session_token) - if pk == target_feed_range: - if pk_range_id == comparison_pk_range_id: - comparison_session_token = comparison_session_token.merge(vector_session_token) - elif session_token.is_greater(comparison_session_token): - comparison_pk_range_id = pk_range_id - comparison_session_token = session_token - return comparison_pk_range_id + ":" + comparison_session_token.session_token - + target_feed_range_normalized = target_feed_range.get_normalized_range() + # filter out tuples that overlap with target_feed_range and normalizes all the ranges + overlapping_ranges = [(feed_range[0].get_normalized_range(), feed_range[1]) for feed_range in feed_ranges_to_session_tokens if + feed_range[0].get_normalized_range().overlaps(target_feed_range_normalized)] + # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples + # with compound session tokens? + if overlapping_ranges == 0: + raise ValueError('There were no overlapping feed ranges with the target.') + + + # Clean this up + for i in range(len(overlapping_ranges)): + for j in range(i + 1, len(overlapping_ranges)): + session_token = overlapping_ranges[i][1] + session_token_1 = overlapping_ranges[j][1] + if (not CosmosClientConnection.is_compound_session_token(session_token) and + not CosmosClientConnection.is_compound_session_token(overlapping_ranges[j][1]) and + overlapping_ranges[i][0] == overlapping_ranges[j][0]): + session_token = CosmosClientConnection.merge_session_tokens(session_token, session_token_1) + overlapping_ranges.append((overlapping_ranges[i][0], session_token)) + overlapping_ranges.remove(overlapping_ranges[i]) + overlapping_ranges.remove(overlapping_ranges[j]) + + + i = 0 + session_token = "" + while i < len(overlapping_ranges): + feed_range_cmp, session_token_cmp = overlapping_ranges[i] + subsets = [] + for j in range(i + 1, len(overlapping_ranges)): + feed_range = overlapping_ranges[j][0] + session_token = overlapping_ranges[j][1] + if not CosmosClientConnection.is_compound_session_token(feed_range) and \ + feed_range.is_subset(feed_range_cmp): + subsets.append(overlapping_ranges[j]) + for j in range(len(subsets)): + merged_range = subsets[j][0] + for k in range(len(subsets)): + if j == k: + continue + if merged_range.can_merge(subsets[k][0]): + merged_range = merged_range.merge(subsets[k][0]) + if feed_range_cmp == merged_range: + session_token = subsets[j][1] + return session_token + + @staticmethod + def merge_session_tokens(session_token1, session_token2): + token_pairs1 = session_token1.split(",") + pk_range_id1 = token_pairs1[0] + vector_session_token1 = VectorSessionToken.create(token_pairs1[1]) + token_pairs2 = session_token2.split(",") + pk_range_id2 = token_pairs2[0] + vector_session_token2 = VectorSessionToken.create(token_pairs2[1]) + pk_range_id = pk_range_id1 + if pk_range_id1 != pk_range_id2: + pk_range_id = pk_range_id1 \ + if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 + vector_session_token = vector_session_token1.merge(vector_session_token2) + return pk_range_id + "," + vector_session_token.session_token + + @staticmethod + def is_compound_session_token(session_token): + return "," in session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py deleted file mode 100644 index 4dbbb42085f7..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_request_context.py +++ /dev/null @@ -1,13 +0,0 @@ -session_token_header = 'x-ms-session-token' -request_context_header = 'request-context' -session_token_request_context = 'session-token' -feed_range_request_context = 'feed-range' - - -def add_request_context(last_response_headers, options) -> dict[str, str]: - request_context = {} - request_context[session_token_request_context] = last_response_headers[session_token_header] - if 'partitionKey' in options: - last_response_headers['partitionKey'] - last_response_headers[request_context_header] = request_context - return last_response_headers diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index f2e7576bf376..e485139ae17f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -206,3 +206,24 @@ def overlaps(range1, range2): return False return True return False + + def can_merge(self, other): + if self.isSingleValue() and other.isSingleValue(): + return self.min == other.min + return self.overlaps(self, other) + + def merge(self, other): + if not self.can_merge(other): + raise ValueError("Ranges do not overlap") + min_val = self.min if self.min < other.min else other.min + max_val = self.max if self.max > other.max else other.max + is_min_inclusive = self.isMinInclusive if self.min < other.min else other.isMinInclusive + is_max_inclusive = self.isMaxInclusive if self.max > other.max else other.isMaxInclusive + return Range(min_val, max_val, is_min_inclusive, is_max_inclusive) + + def is_subset(self, parent_range) -> bool: + normalized_parent_range = parent_range.to_normalized_range() + normalized_child_range = self.to_normalized_range() + return normalized_parent_range.contains(normalized_child_range.min) and \ + (normalized_parent_range.contains(normalized_child_range.max) + or normalized_parent_range.max == normalized_child_range.max) \ 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 d26b1e30bca7..c41241cb9823 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -36,11 +36,13 @@ _deserialize_throughput, _replace_throughput, GenerateGuidId, - _set_properties_cache + _set_properties_cache, + ParsePaths, + TrimBeginningAndEndingSlashes ) from ._change_feed.feed_range import FeedRange, FeedRangeEpk from ._cosmos_client_connection import CosmosClientConnection -from ._routing.routing_range import Range, partition_key_range_to_range_string +from ._routing.routing_range import Range, partition_key_range_to_range_string, PartitionKeyRange from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -259,6 +261,7 @@ def read_item( # pylint:disable=docstring-missing-param request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + self._add_request_context(request_options) return self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) @distributed_trace @@ -316,6 +319,8 @@ def read_all_items( # pylint:disable=docstring-missing-param items = self.client_connection.ReadItems( collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs) + # Change to be full range + self._add_request_context({}) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -499,6 +504,7 @@ def query_items_change_feed( if kwargs.get("partition_key") is not None: change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) + request_context = {"partitionKey": change_feed_state_context["partitionKey"]} change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) @@ -516,6 +522,7 @@ def query_items_change_feed( result = self.client_connection.QueryItemsChangeFeed( self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) + self._add_request_context(request_context) if response_hook: response_hook(self.client_connection.last_response_headers, result) return result @@ -640,6 +647,8 @@ def query_items( # pylint:disable=docstring-missing-param response_hook=response_hook, **kwargs ) + request_context = {"partitionKey": feed_options["partitionKey"]} + self._add_request_context(request_context) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -717,10 +726,12 @@ def replace_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - - return self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, options=request_options, **kwargs + request_context = {} + result = self.client_connection.ReplaceItem( + document_link=item_link, new_document=body, request_context=request_context, options=request_options, **kwargs ) + self._add_request_context(request_context) + return result @distributed_trace def upsert_item( # pylint:disable=docstring-missing-param @@ -783,13 +794,16 @@ def upsert_item( # pylint:disable=docstring-missing-param request_options["populateQueryMetrics"] = populate_query_metrics if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - - return self.client_connection.UpsertItem( + request_context = {} + result = self.client_connection.UpsertItem( database_or_container_link=self.container_link, document=body, + request_context=request_context, options=request_options, **kwargs ) + self._add_request_context(request_context) + return result @distributed_trace def create_item( # pylint:disable=docstring-missing-param @@ -860,8 +874,11 @@ def create_item( # pylint:disable=docstring-missing-param request_options["indexingDirective"] = indexing_directive if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - return self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, options=request_options, **kwargs) + request_context = {} + result = self.client_connection.CreateItem( + database_or_container_link=self.container_link, document=body, request_context=request_context, options=request_options, **kwargs) + self._add_request_context(request_context) + return result @distributed_trace def patch_item( @@ -927,8 +944,11 @@ def patch_item( if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] item_link = self._get_document_link(item) - return self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, options=request_options, **kwargs) + request_context = {} + result = self.client_connection.PatchItem( + document_link=item_link, operations=patch_operations, request_context=request_context, options=request_options, **kwargs) + self._add_request_context(request_context) + return result @distributed_trace def execute_item_batch( @@ -979,10 +999,12 @@ def execute_item_batch( kwargs['priority'] = priority request_options = build_options(kwargs) request_options["partitionKey"] = self._set_partition_key(partition_key) + request_context = {"partitionKey": request_options["partitionKey"]} request_options["disableAutomaticIdGeneration"] = True - - return self.client_connection.Batch( + result = self.client_connection.Batch( collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) + self._add_request_context(request_context) + return result @distributed_trace def delete_item( # pylint:disable=docstring-missing-param @@ -1049,7 +1071,9 @@ def delete_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) + request_context = {"partitionKey": request_options["partitionKey"]} self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) + self._add_request_context(request_context) @distributed_trace def read_offer(self, **kwargs: Any) -> Offer: @@ -1336,7 +1360,7 @@ def read_feed_ranges( def get_updated_session_token(self, feed_ranges_to_session_tokens: List, - target_feed_range: Range + target_feed_range: FeedRange ) -> "Session Token": """Gets the best session token from the list of session token and feed range tuples to figure out which is the most up to date for a specific @@ -1370,7 +1394,10 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - normalized_parent_range = parent_feed_range.get_normalized_range() - normalized_child_range = child_feed_range.get_normalized_range() - return normalized_parent_range.contains(normalized_child_range.min) and \ - (normalized_parent_range.contains(normalized_child_range.max) or normalized_parent_range.max == normalized_child_range.max) + return child_feed_range.get_normalized_range().is_subset(parent_feed_range.get_normalized_range()) + + def _add_request_context(self, request_context): + request_context['session_token'] = self.client_connection.last_response_headers['x-ms-session-token'] + if 'partitionKey' in request_context: + request_context["feed_range"] = self._get_epk_range_for_partition_key(request_context['partitionKey']) + self.client_connection.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 621b353f7987..922e67b7d1e2 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -9,7 +9,7 @@ import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.partition_key as partition_key import test_config -from azure.cosmos._change_feed.feed_range import FeedRangeEpk +from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._routing.routing_range import Range from test.test_config import TestConfig @@ -49,30 +49,31 @@ def test_partition_key_to_feed_range(self, setup): partition_key=partition_key.PartitionKey(path="/id") ) feed_range = created_container.feed_range_from_partition_key("1") - print(feed_range.get_normalized_range()) + assert feed_range.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", + "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) setup["created_db"].delete_container(created_container) - test_ranges = [(Range("", "FFFFFFFFFFFFFFFF", True, False), - Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), + test_ranges = [(Range("", "FF", True, False), + Range("3FFFFFFFFFFFFFFF", "7F", True, False), True), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), - Range("", "FFFFFFFFFFFFFFFF", True, False), + (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + Range("", "FF", True, False), False), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), - Range("", "5FFFFFFFFFFFFFFF", True, False), + (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + Range("", "5F", True, False), False), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), - Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + (Range("3FFFFFFFFFFFFFFF", "7F", True, True), + Range("3FFFFFFFFFFFFFFF", "7F", True, True), True), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", False, True), - Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + (Range("3FFFFFFFFFFFFFFF", "7F", False, True), + Range("3FFFFFFFFFFFFFFF", "7F", True, True), False), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), - Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, True), + (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + Range("3FFFFFFFFFFFFFFF", "7F", True, True), False), - (Range("3FFFFFFFFFFFFFFF", "7FFFFFFFFFFFFFFF", True, False), - Range("", "2FFFFFFFFFFFFFFF", True, False), + (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + Range("", "2F", True, False), False)] @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_ranges) @@ -82,9 +83,11 @@ def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) == is_subset def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FFFFFFFFFFFFFFFF", True, False)) + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + child_feed_range = FeedRangePartitionKey("1", Range("", "CC", True, False)) + assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, child_feed_range) if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py new file mode 100644 index 000000000000..2a34ba2806e2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -0,0 +1,65 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest +from requests import session + +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.partition_key as partition_key +import test_config +from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey +from azure.cosmos._routing.routing_range import Range +from test.test_config import TestConfig + + + +@pytest.fixture(scope="class") +def setup(): + if (TestFeedRange.masterKey == '[YOUR_KEY_HERE]' or + TestFeedRange.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.credential), + created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID) + return { + "created_db": created_db, + "created_collection": created_db.get_container_client(TestFeedRange.TEST_CONTAINER_ID) + } + +@pytest.mark.cosmosEmulator +@pytest.mark.unittest +@pytest.mark.usefixtures("setup") +class TestFeedRange: + """Tests to verify methods for operations on feed ranges + """ + + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID + TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + + test_operations = [] + + #@pytest.mark.parametrize("operation", test_operations) + # test out all operations + def test_crud_request_context(self, setup): + keys_expected = ["session_token", "feed_range", "partitionKey"] + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'key': 'A' + } + setup["created_collection"].create_item(item) + request_context = setup["created_collection"].client_connection.last_response_headers["request_context"] + for key in keys_expected: + assert request_context is not None + assert request_context[key] is not None + + +if __name__ == '__main__': + unittest.main() From 5cde59b96df8aa3219684b6b341891d9d96f9622 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Fri, 13 Sep 2024 15:02:22 -0700 Subject: [PATCH 25/59] revert unnecessary change --- sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst | 8 -------- 1 file changed, 8 deletions(-) delete mode 100644 sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst diff --git a/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst b/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst deleted file mode 100644 index 5b7484884dd7..000000000000 --- a/sdk/cosmos/azure-mgmt-cosmosdb/toc_tree.rst +++ /dev/null @@ -1,8 +0,0 @@ -.. toctree:: - :maxdepth: 5 - :glob: - :caption: Developer Documentation - - ref/azure.common - /Users/annie-mac/dev/git/azure-sdk-for-python/sdk/cosmos/azure-mgmt-cosmosdb/.tox/sphinx/tmp/dist/unzipped/docgen/azure.mgmt.cosmosdb.rst - ref/azure.servicemanagement From a494346f4eeb603e2326910364e569b2ed5284bf Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 20 Sep 2024 09:10:59 -0700 Subject: [PATCH 26/59] Added more tests --- .../azure/cosmos/_cosmos_client_connection.py | 146 ++++++++++++------ .../azure/cosmos/_routing/routing_range.py | 3 + .../azure/cosmos/_session_token_helpers.py | 44 ++++++ .../azure-cosmos/azure/cosmos/container.py | 2 +- .../azure-cosmos/test/test_feed_range.py | 21 ++- .../azure-cosmos/test/test_request_context.py | 14 +- .../test/test_session_token_helpers.py | 91 +++++++---- 7 files changed, 230 insertions(+), 91 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 b9d2e6d7b0ed..1ccb52da8cfa 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -32,6 +32,7 @@ from azure.core import PipelineClient from ._change_feed.feed_range import FeedRange +from ._session_token_helpers import is_compound_session_token, merge_session_tokens from ._vector_session_token import VectorSessionToken from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged @@ -68,6 +69,7 @@ from ._retry_utility import ConnectionRetryPolicy from ._routing import routing_map_provider, routing_range from .documents import ConnectionPolicy, DatabaseAccount +from azure.cosmos._routing.routing_range import Range from .partition_key import ( _Undefined, _Empty, @@ -3338,64 +3340,118 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ target_feed_range_normalized = target_feed_range.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges overlapping_ranges = [(feed_range[0].get_normalized_range(), feed_range[1]) for feed_range in feed_ranges_to_session_tokens if - feed_range[0].get_normalized_range().overlaps(target_feed_range_normalized)] + Range.overlaps(target_feed_range_normalized, feed_range[0].get_normalized_range())] # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples # with compound session tokens? if overlapping_ranges == 0: raise ValueError('There were no overlapping feed ranges with the target.') - - # Clean this up - for i in range(len(overlapping_ranges)): - for j in range(i + 1, len(overlapping_ranges)): - session_token = overlapping_ranges[i][1] - session_token_1 = overlapping_ranges[j][1] - if (not CosmosClientConnection.is_compound_session_token(session_token) and - not CosmosClientConnection.is_compound_session_token(overlapping_ranges[j][1]) and - overlapping_ranges[i][0] == overlapping_ranges[j][0]): - session_token = CosmosClientConnection.merge_session_tokens(session_token, session_token_1) - overlapping_ranges.append((overlapping_ranges[i][0], session_token)) - overlapping_ranges.remove(overlapping_ranges[i]) - overlapping_ranges.remove(overlapping_ranges[j]) - - i = 0 - session_token = "" - while i < len(overlapping_ranges): - feed_range_cmp, session_token_cmp = overlapping_ranges[i] + j = 1 + while i < len(overlapping_ranges) and j < len(overlapping_ranges): + cur_feed_range = overlapping_ranges[i][0] + session_token = overlapping_ranges[i][1] + session_token_1 = overlapping_ranges[j][1] + if (not is_compound_session_token(session_token) and + not is_compound_session_token(overlapping_ranges[j][1]) and + cur_feed_range == overlapping_ranges[j][0]): + session_token = merge_session_tokens(session_token, session_token_1) + feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + overlapping_ranges.append((cur_feed_range, session_token)) + else: + j += 1 + if j == len(overlapping_ranges): + i += 1 + j = i + 1 + + + updated_session_token = "" + remaining_session_tokens = [] + done_overlapping_ranges = [] + while len(overlapping_ranges) != 0: + feed_range_cmp, session_token_cmp = overlapping_ranges[0] + if is_compound_session_token(session_token_cmp): + done_overlapping_ranges.append(overlapping_ranges[0]) + overlapping_ranges.remove(overlapping_ranges[0]) + continue + tokens_cmp = session_token_cmp.split(":") + vector_session_token_cmp = VectorSessionToken.create(tokens_cmp[1]) subsets = [] - for j in range(i + 1, len(overlapping_ranges)): + # creating subsets of feed ranges that are subsets of the current feed range + for j in range(1, len(overlapping_ranges)): feed_range = overlapping_ranges[j][0] - session_token = overlapping_ranges[j][1] - if not CosmosClientConnection.is_compound_session_token(feed_range) and \ + if not is_compound_session_token(overlapping_ranges[j][1]) and \ feed_range.is_subset(feed_range_cmp): - subsets.append(overlapping_ranges[j]) - for j in range(len(subsets)): + subsets.append(overlapping_ranges[j] + (j,)) + + # go through subsets to see if can create current feed range from the subsets + not_found = True + j = 0 + while not_found and j < len(subsets): merged_range = subsets[j][0] + session_tokens = [subsets[j][1]] + merged_indices = [subsets[j][2]] for k in range(len(subsets)): if j == k: continue if merged_range.can_merge(subsets[k][0]): merged_range = merged_range.merge(subsets[k][0]) + session_tokens.append(subsets[k][1]) + merged_indices.append(subsets[k][2]) if feed_range_cmp == merged_range: - session_token = subsets[j][1] - return session_token - - @staticmethod - def merge_session_tokens(session_token1, session_token2): - token_pairs1 = session_token1.split(",") - pk_range_id1 = token_pairs1[0] - vector_session_token1 = VectorSessionToken.create(token_pairs1[1]) - token_pairs2 = session_token2.split(",") - pk_range_id2 = token_pairs2[0] - vector_session_token2 = VectorSessionToken.create(token_pairs2[1]) - pk_range_id = pk_range_id1 - if pk_range_id1 != pk_range_id2: - pk_range_id = pk_range_id1 \ - if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 - vector_session_token = vector_session_token1.merge(vector_session_token2) - return pk_range_id + "," + vector_session_token.session_token - - @staticmethod - def is_compound_session_token(session_token): - return "," in session_token + all_global_lsns_larger = True + for session_token in session_tokens: + tokens = session_token.split(":") + vector_session_token = VectorSessionToken.create(tokens[1]) + if vector_session_token.global_lsn < vector_session_token_cmp.global_lsn: + all_global_lsns_larger = False + break + feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + if all_global_lsns_larger: + overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) + overlapping_ranges.remove(overlapping_ranges[0]) + not_found = False + break + j += 1 + + done_overlapping_ranges.append(overlapping_ranges[0]) + overlapping_ranges.remove(overlapping_ranges[0]) + + for _, session_token in done_overlapping_ranges: + # here break up session tokens that are compound + if is_compound_session_token(session_token): + tokens = session_token.split(",") + for token in tokens: + remaining_session_tokens.append(token) + else: + remaining_session_tokens.append(session_token) + + if len(remaining_session_tokens) == 1: + return remaining_session_tokens[0] + new_session_tokens = [] + # merging any session tokens with same pkrangeid + for i in range(len(remaining_session_tokens)): + for j in range(i + 1, len(remaining_session_tokens)): + tokens1 = remaining_session_tokens[i].split(":") + tokens2 = remaining_session_tokens[j].split(":") + pk_range_id1 = tokens1[0] + pk_range_id2 = tokens2[0] + if pk_range_id1 == pk_range_id2: + vector_session_token1 = VectorSessionToken.create(tokens1[1]) + vector_session_token2 = VectorSessionToken.create(tokens2[1]) + vector_session_token = vector_session_token1.merge(vector_session_token2) + new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) + remaining_session_tokens.remove(remaining_session_tokens[i]) + remaining_session_tokens.remove(remaining_session_tokens[j]) + new_session_tokens.extend(remaining_session_tokens) + for i in range(len(new_session_tokens)): + if i == len(new_session_tokens) - 1: + updated_session_token += new_session_tokens[i] + else: + updated_session_token += new_session_tokens[i] + "," + + return updated_session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index e485139ae17f..e5ada7030c23 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -210,6 +210,9 @@ def overlaps(range1, range2): def can_merge(self, other): if self.isSingleValue() and other.isSingleValue(): return self.min == other.min + # if share the same boundary, they can merge + if (self.max == other.min and self.isMaxInclusive or other.isMinInclusive) or (other.max == self.min and other.isMaxInclusive or self.isMinInclusive): + return True return self.overlaps(self, other) def merge(self, other): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index b28b04f64312..5447d26d38ee 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -1,3 +1,47 @@ +# The MIT License (MIT) +# Copyright (c) 2014 Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Internal Helper functions for manipulating session tokens. +""" + +from azure.cosmos._vector_session_token import VectorSessionToken + + +def merge_session_tokens(session_token1, session_token2): + # TODO method for splitting the session token into pk_range_id and vector_session_token + token_pairs1 = session_token1.split(":") + pk_range_id1 = token_pairs1[0] + vector_session_token1 = VectorSessionToken.create(token_pairs1[1]) + token_pairs2 = session_token2.split(":") + pk_range_id2 = token_pairs2[0] + vector_session_token2 = VectorSessionToken.create(token_pairs2[1]) + pk_range_id = pk_range_id1 + if pk_range_id1 != pk_range_id2: + pk_range_id = pk_range_id1 \ + if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 + vector_session_token = vector_session_token1.merge(vector_session_token2) + return pk_range_id + ":" + vector_session_token.session_token + +def is_compound_session_token(session_token): + return "," in session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index db6a263ac63e..35b12bf7b57c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1372,7 +1372,7 @@ def get_updated_session_token(self, :returns: a session token :rtype: str """ - self.client_connection._get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return self.client_connection._get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 922e67b7d1e2..139e266f1d15 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -53,26 +53,25 @@ def test_partition_key_to_feed_range(self, setup): "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) setup["created_db"].delete_container(created_container) - test_ranges = [(Range("", "FF", True, False), - Range("3FFFFFFFFFFFFFFF", "7F", True, False), + Range("3F", "7F", True, False), True), - (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + (Range("3F", "7F", True, False), Range("", "FF", True, False), False), - (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + (Range("3F", "7F", True, False), Range("", "5F", True, False), False), - (Range("3FFFFFFFFFFFFFFF", "7F", True, True), - Range("3FFFFFFFFFFFFFFF", "7F", True, True), + (Range("3F", "7F", True, True), + Range("3F", "7F", True, True), True), - (Range("3FFFFFFFFFFFFFFF", "7F", False, True), - Range("3FFFFFFFFFFFFFFF", "7F", True, True), + (Range("3F", "7F", False, True), + Range("3F", "7F", True, True), False), - (Range("3FFFFFFFFFFFFFFF", "7F", True, False), - Range("3FFFFFFFFFFFFFFF", "7F", True, True), + (Range("3F", "7F", True, False), + Range("3F", "7F", True, True), False), - (Range("3FFFFFFFFFFFFFFF", "7F", True, False), + (Range("3F", "7F", True, False), Range("", "2F", True, False), False)] diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index 2a34ba2806e2..ad46108c131c 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -18,23 +18,23 @@ @pytest.fixture(scope="class") def setup(): - if (TestFeedRange.masterKey == '[YOUR_KEY_HERE]' or - TestFeedRange.host == '[YOUR_ENDPOINT_HERE]'): + if (TestRequestContext.masterKey == '[YOUR_KEY_HERE]' or + TestRequestContext.host == '[YOUR_ENDPOINT_HERE]'): raise Exception( "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.credential), - created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID) + test_client = cosmos_client.CosmosClient(TestRequestContext.host, TestConfig.credential), + created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) return { "created_db": created_db, - "created_collection": created_db.get_container_client(TestFeedRange.TEST_CONTAINER_ID) + "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) } @pytest.mark.cosmosEmulator @pytest.mark.unittest @pytest.mark.usefixtures("setup") -class TestFeedRange: +class TestRequestContext: """Tests to verify methods for operations on feed ranges """ @@ -46,7 +46,7 @@ class TestFeedRange: test_operations = [] #@pytest.mark.parametrize("operation", test_operations) - # test out all operations + # test out all operations and using response hook def test_crud_request_context(self, setup): keys_expected = ["session_token", "feed_range", "partitionKey"] item = { diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 748977a3739a..fd87c05ba043 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -1,19 +1,53 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import random import unittest -from time import sleep import pytest import azure.cosmos.cosmos_client as cosmos_client import test_config -from azure.cosmos import DatabaseProxy, ThroughputProperties, PartitionKey +from azure.cosmos import DatabaseProxy +from azure.cosmos._change_feed.feed_range import FeedRangeEpk +from azure.cosmos._routing.routing_range import Range from test.test_config import TestConfig +@pytest.fixture(scope="class") +def setup(): + if (TestSessionTokenHelpers.masterKey == '[YOUR_KEY_HERE]' or + TestSessionTokenHelpers.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, TestConfig.credential), + created_db = test_client[0].get_database_client(TestSessionTokenHelpers.TEST_DATABASE_ID) + return { + "created_db": created_db, + "created_collection": created_db.get_container_client(TestSessionTokenHelpers.TEST_COLLECTION_ID) + } + +def create_split_ranges(): + test_params = [([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], + ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), + ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], + ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), + ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], + ("AA", "DD"), "0:1#55#3=52")] + actual_test_params = [] + for test_param in test_params: + split_ranges = [] + for feed_range, session_token in test_param[0]: + split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], + True, False)), session_token)) + target_feed_range = FeedRangeEpk(Range(test_param[1][0], test_param[1][1], True, False)) + actual_test_params.append((split_ranges, target_feed_range, test_param[2])) + return actual_test_params @pytest.mark.cosmosEmulator -class TestSessionTokenMerge(unittest.TestCase): +@pytest.mark.unittest +@pytest.mark.usefixtures("setup") +class TestSessionTokenHelpers: """Test to ensure escaping of non-ascii characters from partition key""" created_db: DatabaseProxy = None @@ -25,33 +59,36 @@ class TestSessionTokenMerge(unittest.TestCase): TEST_DATABASE_ID = configs.TEST_DATABASE_ID TEST_COLLECTION_ID = configs.TEST_MULTI_PARTITION_CONTAINER_ID - @classmethod - def setUpClass(cls): - # creates the database, collection, and insert all the documents - # we will gain some speed up in running the tests by creating the - # database, collection and inserting all the docs only once - - if cls.masterKey == '[YOUR_KEY_HERE]' or cls.host == '[YOUR_ENDPOINT_HERE]': - raise Exception("You must specify your Azure Cosmos account values for " - "'masterKey' and 'host' at the top of this class to run the " - "tests.") - - cls.client = cosmos_client.CosmosClient(cls.host, TestConfig.credential) - cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) - cls.created_collection = cls.created_db.create_container("TestColl", PartitionKey("/mypk"), offer_throughput=ThroughputProperties(11000)) # have a test for split, merge, different versions, compound session tokens, hpk, and for normal case for several session tokens # have tests for all the scenarios in the testing plan # should try once with 35000 session tokens - def test_session_token_merge(self): - sleep(5) - for i in range(10): - self.created_collection.create_item(body={'id': 'doc' + str(i), 'mypk': str(i)}) - session_token = self.created_collection.client_connection.last_response_headers['x-ms-session-token'] - self.created_collection.read_all_items(session_token=session_token) - session_token = self.created_collection.client_connection.last_response_headers['x-ms-session-token'] - print("\n" + "Session Token: " + session_token) - self.client.delete_database(self.TEST_DATABASE_ID) + # check if pkrangeid is being filtered out + def test_get_session_token_update(self, setup): + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + session_token = "0:1#54#3=50" + feed_ranges_and_session_tokens = [(feed_range, session_token)] + session_token = "0:1#51#3=52" + feed_ranges_and_session_tokens.append((feed_range, session_token)) + session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) + assert session_token == "0:1#54#3=52" + + def test_many_session_tokens_update_same_range(self, setup): + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_ranges_and_session_tokens = [] + for i in range(1000): + session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) + feed_ranges_and_session_tokens.append((feed_range, session_token)) + session_token = "0:1#101#3=101" + feed_ranges_and_session_tokens.append((feed_range, session_token)) + updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) + assert updated_session_token == session_token + + @pytest.mark.parametrize("split_ranges, target_feed_range, expected_session_token", create_split_ranges()) + def test_simulated_splits(self, setup, split_ranges, target_feed_range, expected_session_token): + updated_session_token = (setup["created_collection"] + .get_updated_session_token(split_ranges, target_feed_range)) + assert expected_session_token == updated_session_token From ad3ae4fd52267be6573dd1f01552a79fc84a48cc Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 4 Oct 2024 18:01:33 -0700 Subject: [PATCH 27/59] Added more tests --- .../azure/cosmos/_cosmos_client_connection.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 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 1ccb52da8cfa..bb8ff119b7e1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3401,19 +3401,29 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ session_tokens.append(subsets[k][1]) merged_indices.append(subsets[k][2]) if feed_range_cmp == merged_range: - all_global_lsns_larger = True + # if it is the bigger one remove the smaller ranges + # if it is the smaller ranges remove the bigger range + # if it is neither compound + child_lsns_larger = True + child_lsns_smaller = True for session_token in session_tokens: tokens = session_token.split(":") vector_session_token = VectorSessionToken.create(tokens[1]) if vector_session_token.global_lsn < vector_session_token_cmp.global_lsn: - all_global_lsns_larger = False - break + child_lsns_smaller = False + else: + child_lsns_larger = False feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] for feed_range_to_remove in feed_ranges_to_remove: overlapping_ranges.remove(feed_range_to_remove) - if all_global_lsns_larger: + if child_lsns_larger: + session_tokens.remove(session_token_cmp) overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) overlapping_ranges.remove(overlapping_ranges[0]) + elif child_lsns_smaller: + overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) + overlapping_ranges.remove(overlapping_ranges[0]) + not_found = False break j += 1 From 5249d0a6ebefe0b515ec9e89285e237db756e76c Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Sun, 6 Oct 2024 13:11:58 -0700 Subject: [PATCH 28/59] Changed tests to use new public feed range and more test coverage for request context --- .../azure/cosmos/_change_feed/feed_range.py | 124 ------------------ .../azure/cosmos/_cosmos_client_connection.py | 3 + .../azure-cosmos/azure/cosmos/container.py | 17 ++- .../azure-cosmos/test/test_feed_range.py | 8 +- .../azure-cosmos/test/test_request_context.py | 64 ++++++--- .../test/test_session_token_helpers.py | 10 +- 6 files changed, 69 insertions(+), 157 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py deleted file mode 100644 index 293178436cdf..000000000000 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py +++ /dev/null @@ -1,124 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) 2014 Microsoft Corporation - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -"""Internal class for feed range implementation in the Azure Cosmos -database service. -""" -from abc import ABC, abstractmethod -from typing import Union, List, Dict, Any - -from azure.cosmos._routing.routing_range import Range, PartitionKeyRange -from azure.cosmos.partition_key import _Undefined, _Empty - - -class FeedRange(ABC): - - @abstractmethod - def get_normalized_range(self) -> Range: - pass - - @abstractmethod - def to_dict(self) -> Dict[str, Any]: - pass - - - - @staticmethod - def _get_range(pk_ranges) -> Range: - max_range = pk_ranges[0][PartitionKeyRange.MaxExclusive] - min_range = pk_ranges[0][PartitionKeyRange.MinInclusive] - for i in range(1, len(pk_ranges)): - pk_ranges_min = pk_ranges[i][PartitionKeyRange.MinInclusive] - pk_ranges_max = pk_ranges[i][PartitionKeyRange.MaxExclusive] - if pk_ranges_min < min_range: - min_range = pk_ranges_min - if pk_ranges_max > max_range: - max_range = pk_ranges_max - return Range(min_range, max_range, True, False) - -class FeedRangePartitionKey(FeedRange): - type_property_name = "PK" - - def __init__( - self, - pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], - feed_range: Range) -> None: # pylint: disable=line-too-long - - if pk_value is None: - raise ValueError("PartitionKey cannot be None") - if feed_range is None: - raise ValueError("Feed range cannot be None") - - self._pk_value = pk_value - self._feed_range = feed_range - - def get_normalized_range(self) -> Range: - return self._feed_range.to_normalized_range() - - def to_dict(self) -> Dict[str, Any]: - if isinstance(self._pk_value, _Undefined): - return { self.type_property_name: [{}] } - if isinstance(self._pk_value, _Empty): - return { self.type_property_name: [] } - if isinstance(self._pk_value, list): - return { self.type_property_name: list(self._pk_value) } - - return { self.type_property_name: self._pk_value } - - @classmethod - def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartitionKey': - if data.get(cls.type_property_name): - pk_value = data.get(cls.type_property_name) - if not pk_value: - return cls(_Empty(), feed_range) - if pk_value == [{}]: - return cls(_Undefined(), feed_range) - if isinstance(pk_value, list): - return cls(list(pk_value), feed_range) - return cls(data[cls.type_property_name], feed_range) - - raise ValueError(f"Can not parse FeedRangePartitionKey from the json," - f" there is no property {cls.type_property_name}") - - -class FeedRangeEpk(FeedRange): - type_property_name = "Range" - - def __init__(self, feed_range: Range) -> None: - if feed_range is None: - raise ValueError("feed_range cannot be None") - - self._range = feed_range - - def get_normalized_range(self) -> Range: - return self._range.to_normalized_range() - - def to_dict(self) -> Dict[str, Any]: - return { - self.type_property_name: self._range.to_dict() - } - - @classmethod - def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeEpk': - if data.get(cls.type_property_name): - feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) - return cls(feed_range) - raise ValueError(f"Can not parse FeedRangeEPK from the json, there is no property {cls.type_property_name}") 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 17c0e5551de7..a5fcc046a906 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -1299,6 +1299,7 @@ def UpsertItem( self, database_or_container_link: str, document: Dict[str, Any], + request_context: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -1325,6 +1326,7 @@ def UpsertItem( # link in case of client side partitioning if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) + request_context["partitionKey"] = options["partitionKey"] collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options @@ -2010,6 +2012,7 @@ def ReplaceItem( # Extract the document collection link and add the partition key to options collection_link = base.GetItemContainerLink(document_link) options = self._AddPartitionKey(collection_link, new_document, options) + request_context["partitionKey"] = options["partitionKey"] return self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index a69f505f4197..eaed99d681c2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -259,8 +259,10 @@ def read_item( # pylint:disable=docstring-missing-param request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - self._add_request_context(request_options) - return self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) + request_context = {"partitionKey": request_options["partitionKey"]} + items = self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) + self._add_request_context(request_context) + return items @distributed_trace def read_all_items( # pylint:disable=docstring-missing-param @@ -317,7 +319,6 @@ def read_all_items( # pylint:disable=docstring-missing-param items = self.client_connection.ReadItems( collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs) - # Change to be full range self._add_request_context({}) if response_hook: response_hook(self.client_connection.last_response_headers, items) @@ -643,6 +644,7 @@ def query_items( # pylint:disable=docstring-missing-param feed_options["populateQueryMetrics"] = populate_query_metrics if populate_index_metrics is not None: feed_options["populateIndexMetrics"] = populate_index_metrics + request_context = {} if partition_key is not None: partition_key_value = self._set_partition_key(partition_key) if self.__is_prefix_partitionkey(partition_key): @@ -652,6 +654,7 @@ def query_items( # pylint:disable=docstring-missing-param kwargs["partitionKeyDefinition"]["partition_key"] = partition_key_value else: feed_options["partitionKey"] = partition_key_value + request_context["partitionKey"] = feed_options["partitionKey"] if enable_scan_in_query is not None: feed_options["enableScanInQuery"] = enable_scan_in_query if max_integrated_cache_staleness_in_ms: @@ -673,7 +676,6 @@ def query_items( # pylint:disable=docstring-missing-param response_hook=response_hook, **kwargs ) - request_context = {"partitionKey": feed_options["partitionKey"]} self._add_request_context(request_context) if response_hook: response_hook(self.client_connection.last_response_headers, items) @@ -1447,10 +1449,13 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - return child_feed_range.get_normalized_range().is_subset(parent_feed_range.get_normalized_range()) + return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_range()) def _add_request_context(self, request_context): request_context['session_token'] = self.client_connection.last_response_headers['x-ms-session-token'] if 'partitionKey' in request_context: - request_context["feed_range"] = self._get_epk_range_for_partition_key(request_context['partitionKey']) + request_context["feed_range"] = FeedRangeEpk(self._get_epk_range_for_partition_key(request_context['partitionKey'])) + else: + # default to full range + request_context["feed_range"] = FeedRangeEpk(Range("", "FF", True, False)) self.client_connection.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 139e266f1d15..42bf9bc0de67 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -9,7 +9,7 @@ import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.partition_key as partition_key import test_config -from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey +from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range from test.test_config import TestConfig @@ -23,7 +23,7 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.credential), + test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.masterKey), created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID) return { "created_db": created_db, @@ -49,7 +49,7 @@ def test_partition_key_to_feed_range(self, setup): partition_key=partition_key.PartitionKey(path="/id") ) feed_range = created_container.feed_range_from_partition_key("1") - assert feed_range.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", + assert feed_range._feed_range_internal.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) setup["created_db"].delete_container(created_container) @@ -85,8 +85,6 @@ def test_feed_range_is_subset_from_pk(self, setup): epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - child_feed_range = FeedRangePartitionKey("1", Range("", "CC", True, False)) - assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, child_feed_range) if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index ad46108c131c..bd7cbab2b310 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -5,12 +5,9 @@ import uuid import pytest -from requests import session import azure.cosmos.cosmos_client as cosmos_client -import azure.cosmos.partition_key as partition_key import test_config -from azure.cosmos._change_feed.feed_range import FeedRangeEpk, FeedRangePartitionKey from azure.cosmos._routing.routing_range import Range from test.test_config import TestConfig @@ -24,13 +21,33 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestRequestContext.host, TestConfig.credential), + test_client = cosmos_client.CosmosClient(TestRequestContext.host, TestConfig.masterKey), created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) return { "created_db": created_db, "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) } +def validate_request_context(setup, with_partition_key=True): + request_context = setup["created_collection"].client_connection.last_response_headers["request_context"] + keys_expected = ["session_token", "feed_range"] + if with_partition_key: + keys_expected.append("partitionKey") + assert request_context is not None + for key in keys_expected: + assert request_context[key] is not None + if not with_partition_key: + # when operation is not scoped with partition key that operation is done on full feed range + assert request_context["feed_range"]._feed_range_internal.get_normalized_range() == Range("", "FF", True, False) + +def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'): + item = { + 'id': id, + 'name': name, + 'pk': pk + } + return item + @pytest.mark.cosmosEmulator @pytest.mark.unittest @pytest.mark.usefixtures("setup") @@ -43,23 +60,36 @@ class TestRequestContext: TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID - test_operations = [] - - #@pytest.mark.parametrize("operation", test_operations) # test out all operations and using response hook def test_crud_request_context(self, setup): - keys_expected = ["session_token", "feed_range", "partitionKey"] - item = { - 'id': 'item' + str(uuid.uuid4()), - 'name': 'sample', - 'key': 'A' - } + item = createItem() setup["created_collection"].create_item(item) - request_context = setup["created_collection"].client_connection.last_response_headers["request_context"] - for key in keys_expected: - assert request_context is not None - assert request_context[key] is not None + validate_request_context(setup) + + setup["created_collection"].read_item(item['id'], item['pk']) + validate_request_context(setup) + + new_item = createItem(name='sample_replaced') + setup["created_collection"].replace_item(item['id'], new_item) + validate_request_context(setup) + + setup["created_collection"].upsert_item(createItem()) + validate_request_context(setup) + + # with partition key + setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])], partition_key=item['pk']) + validate_request_context(setup) + + # without partition key + setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])]) + validate_request_context(setup, False) + + + setup["created_collection"].read_all_items() + validate_request_context(setup, False) + setup["created_collection"].delete_item(item['id'], item['pk']) + validate_request_context(setup) if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index fd87c05ba043..25ca568943cf 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -8,7 +8,7 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos import DatabaseProxy -from azure.cosmos._change_feed.feed_range import FeedRangeEpk +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._routing.routing_range import Range from test.test_config import TestConfig @@ -20,7 +20,7 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, TestConfig.credential), + test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, TestConfig.masterKey), created_db = test_client[0].get_database_client(TestSessionTokenHelpers.TEST_DATABASE_ID) return { "created_db": created_db, @@ -38,9 +38,9 @@ def create_split_ranges(): for test_param in test_params: split_ranges = [] for feed_range, session_token in test_param[0]: - split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], + split_ranges.append((FeedRangeInternalEpk(Range(feed_range[0], feed_range[1], True, False)), session_token)) - target_feed_range = FeedRangeEpk(Range(test_param[1][0], test_param[1][1], True, False)) + target_feed_range = FeedRangeInternalEpk(Range(test_param[1][0], test_param[1][1], True, False)) actual_test_params.append((split_ranges, target_feed_range, test_param[2])) return actual_test_params @@ -74,7 +74,7 @@ def test_get_session_token_update(self, setup): assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeInternalEpk(Range("AA", "BB", True, False)) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) From 40523f5a5edc51474588bc7caab90b71c93ca81b Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Sun, 6 Oct 2024 21:34:25 -0700 Subject: [PATCH 29/59] Added more tests --- .../azure/cosmos/_cosmos_client_connection.py | 113 ++++++++---------- .../azure/cosmos/_routing/routing_range.py | 5 +- .../azure/cosmos/_session_token_helpers.py | 35 +++++- .../azure-cosmos/azure/cosmos/container.py | 4 +- .../azure-cosmos/test/test_feed_range.py | 79 +++++++----- .../azure-cosmos/test/test_request_context.py | 7 +- .../test/test_session_token_helpers.py | 45 ++++--- 7 files changed, 169 insertions(+), 119 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 a5fcc046a906..35efabbbeedf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -30,7 +30,8 @@ from urllib3.util.retry import Retry from azure.core import PipelineClient -from ._session_token_helpers import is_compound_session_token, merge_session_tokens +from ._session_token_helpers import is_compound_session_token, merge_session_tokens, split_compound_session_tokens, \ + merge_session_tokens_with_same_pkrangeid from ._vector_session_token import VectorSessionToken from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged @@ -3337,15 +3338,17 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s return partition_key_definition def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_range): - target_feed_range_normalized = target_feed_range.get_normalized_range() + target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges - overlapping_ranges = [(feed_range[0].get_normalized_range(), feed_range[1]) for feed_range in feed_ranges_to_session_tokens if - Range.overlaps(target_feed_range_normalized, feed_range[0].get_normalized_range())] + overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1]) + for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( + target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples # with compound session tokens? - if overlapping_ranges == 0: + if len(overlapping_ranges) == 0: raise ValueError('There were no overlapping feed ranges with the target.') + # merge any session tokens that are the same exact feed range i = 0 j = 1 while i < len(overlapping_ranges) and j < len(overlapping_ranges): @@ -3368,7 +3371,6 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ updated_session_token = "" - remaining_session_tokens = [] done_overlapping_ranges = [] while len(overlapping_ranges) != 0: feed_range_cmp, session_token_cmp = overlapping_ranges[0] @@ -3393,75 +3395,56 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ merged_range = subsets[j][0] session_tokens = [subsets[j][1]] merged_indices = [subsets[j][2]] - for k in range(len(subsets)): - if j == k: - continue - if merged_range.can_merge(subsets[k][0]): - merged_range = merged_range.merge(subsets[k][0]) - session_tokens.append(subsets[k][1]) - merged_indices.append(subsets[k][2]) - if feed_range_cmp == merged_range: - # if it is the bigger one remove the smaller ranges - # if it is the smaller ranges remove the bigger range - # if it is neither compound - child_lsns_larger = True - child_lsns_smaller = True - for session_token in session_tokens: - tokens = session_token.split(":") - vector_session_token = VectorSessionToken.create(tokens[1]) - if vector_session_token.global_lsn < vector_session_token_cmp.global_lsn: - child_lsns_smaller = False - else: - child_lsns_larger = False - feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] - for feed_range_to_remove in feed_ranges_to_remove: - overlapping_ranges.remove(feed_range_to_remove) - if child_lsns_larger: - session_tokens.remove(session_token_cmp) - overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) - overlapping_ranges.remove(overlapping_ranges[0]) - elif child_lsns_smaller: - overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) - overlapping_ranges.remove(overlapping_ranges[0]) - - not_found = False - break + if len(subsets) == 1: + tokens = session_tokens[0].split(":") + vector_session_token = VectorSessionToken.create(tokens[1]) + if vector_session_token_cmp.is_greater(vector_session_token): + overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) + else: + for k in range(len(subsets)): + if j == k: + continue + if merged_range.can_merge(subsets[k][0]): + merged_range = merged_range.merge(subsets[k][0]) + session_tokens.append(subsets[k][1]) + merged_indices.append(subsets[k][2]) + if feed_range_cmp == merged_range: + # if it is the bigger one remove the smaller ranges + # if it is the smaller ranges remove the bigger range + # if it is neither compound + child_lsns_larger = True + for session_token in session_tokens: + tokens = session_token.split(":") + vector_session_token = VectorSessionToken.create(tokens[1]) + if vector_session_token_cmp.is_greater(vector_session_token): + child_lsns_larger = False + feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + if child_lsns_larger: + overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) + overlapping_ranges.remove(overlapping_ranges[0]) + not_found = False + break + j += 1 done_overlapping_ranges.append(overlapping_ranges[0]) overlapping_ranges.remove(overlapping_ranges[0]) - for _, session_token in done_overlapping_ranges: - # here break up session tokens that are compound - if is_compound_session_token(session_token): - tokens = session_token.split(",") - for token in tokens: - remaining_session_tokens.append(token) - else: - remaining_session_tokens.append(session_token) + # break up session tokens that are compound + remaining_session_tokens = split_compound_session_tokens(done_overlapping_ranges) if len(remaining_session_tokens) == 1: return remaining_session_tokens[0] - new_session_tokens = [] # merging any session tokens with same pkrangeid + remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens) + + # compound the remaining session tokens for i in range(len(remaining_session_tokens)): - for j in range(i + 1, len(remaining_session_tokens)): - tokens1 = remaining_session_tokens[i].split(":") - tokens2 = remaining_session_tokens[j].split(":") - pk_range_id1 = tokens1[0] - pk_range_id2 = tokens2[0] - if pk_range_id1 == pk_range_id2: - vector_session_token1 = VectorSessionToken.create(tokens1[1]) - vector_session_token2 = VectorSessionToken.create(tokens2[1]) - vector_session_token = vector_session_token1.merge(vector_session_token2) - new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) - remaining_session_tokens.remove(remaining_session_tokens[i]) - remaining_session_tokens.remove(remaining_session_tokens[j]) - new_session_tokens.extend(remaining_session_tokens) - for i in range(len(new_session_tokens)): - if i == len(new_session_tokens) - 1: - updated_session_token += new_session_tokens[i] + if i == len(remaining_session_tokens) - 1: + updated_session_token += remaining_session_tokens[i] else: - updated_session_token += new_session_tokens[i] + "," + updated_session_token += remaining_session_tokens[i] + "," return updated_session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index f6d8e3d0d7e3..fd0c806b3dc7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -186,7 +186,6 @@ def _compare_helper(a, b): @staticmethod def overlaps(range1, range2): - if range1 is None or range2 is None: return False if range1.isEmpty() or range2.isEmpty(): @@ -195,7 +194,7 @@ def overlaps(range1, range2): cmp1 = Range._compare_helper(range1.min, range2.max) cmp2 = Range._compare_helper(range2.min, range1.max) - if cmp1 <= 0 or cmp2 <= 0: + if cmp1 <= 0 and cmp2 <= 0: if (cmp1 == 0 and not (range1.isMinInclusive and range2.isMaxInclusive)) or ( cmp2 == 0 and not (range2.isMinInclusive and range1.isMaxInclusive) ): @@ -225,4 +224,4 @@ def is_subset(self, parent_range) -> bool: normalized_child_range = self.to_normalized_range() return normalized_parent_range.contains(normalized_child_range.min) and \ (normalized_parent_range.contains(normalized_child_range.max) - or normalized_parent_range.max == normalized_child_range.max) \ No newline at end of file + or normalized_parent_range.max == normalized_child_range.max) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 5447d26d38ee..dbb07a4a066b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -36,12 +36,45 @@ def merge_session_tokens(session_token1, session_token2): pk_range_id = pk_range_id1 if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ - if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 + if vector_session_token1.is_greater(vector_session_token2) else pk_range_id2 vector_session_token = vector_session_token1.merge(vector_session_token2) return pk_range_id + ":" + vector_session_token.session_token def is_compound_session_token(session_token): return "," in session_token +def split_compound_session_tokens(compound_session_tokens): + session_tokens = [] + for _, session_token in compound_session_tokens: + if is_compound_session_token(session_token): + tokens = session_token.split(",") + for token in tokens: + session_tokens.append(token) + else: + session_tokens.append(session_token) + return session_tokens +def merge_session_tokens_with_same_pkrangeid(session_tokens): + new_session_tokens = [] + i = 0 + while i < len(session_tokens): + j = i + 1 + while j < len(session_tokens): + tokens1 = session_tokens[i].split(":") + tokens2 = session_tokens[j].split(":") + pk_range_id1 = tokens1[0] + pk_range_id2 = tokens2[0] + if pk_range_id1 == pk_range_id2: + vector_session_token1 = VectorSessionToken.create(tokens1[1]) + vector_session_token2 = VectorSessionToken.create(tokens2[1]) + vector_session_token = vector_session_token1.merge(vector_session_token2) + new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) + remove_session_tokens = [session_tokens[i], session_tokens[j]] + for token in remove_session_tokens: + session_tokens.remove(token) + i = -1 + j += 1 + i += 1 + new_session_tokens.extend(session_tokens) + return new_session_tokens diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index eaed99d681c2..eecd05f2b3c1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -36,9 +36,7 @@ _deserialize_throughput, _replace_throughput, GenerateGuidId, - _set_properties_cache, - ParsePaths, - TrimBeginningAndEndingSlashes + _set_properties_cache ) from ._cosmos_client_connection import CosmosClientConnection from ._feed_range import FeedRange, FeedRangeEpk diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 42bf9bc0de67..1efc9f13c3c3 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -11,9 +11,6 @@ import test_config from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range -from test.test_config import TestConfig - - @pytest.fixture(scope="class") def setup(): @@ -23,13 +20,58 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.masterKey), + test_client = cosmos_client.CosmosClient(TestFeedRange.host, test_config.TestConfig.masterKey), created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID) return { "created_db": created_db, "created_collection": created_db.get_container_client(TestFeedRange.TEST_CONTAINER_ID) } +test_subset_ranges = [(Range("", "FF", True, False), + Range("3F", "7F", True, False), + True), + (Range("3F", "7F", True, False), + Range("", "FF", True, False), + False), + (Range("3F", "7F", True, False), + Range("", "5F", True, False), + False), + (Range("3F", "7F", True, True), + Range("3F", "7F", True, True), + True), + (Range("3F", "7F", False, True), + Range("3F", "7F", True, True), + False), + (Range("3F", "7F", True, False), + Range("3F", "7F", True, True), + False), + (Range("3F", "7F", True, False), + Range("", "2F", True, False), + False)] + + +test_overlaps_ranges = [(Range("", "FF", True, False), + Range("3F", "7F", True, False), + True), + (Range("3F", "7F", True, False), + Range("", "FF", True, False), + True), + (Range("3F", "7F", True, False), + Range("", "5F", True, False), + True), + (Range("3F", "7F", True, False), + Range("3F", "7F", True, False), + True), + (Range("3F", "7F", True, False), + Range("", "2F", True, False), + False), + (Range("3F", "7F", True, False), + Range("6F", "FF", True, False), + True), + (Range("AA", "BB", True, False), + Range("CC", "FF", True, False), + False)] + @pytest.mark.cosmosEmulator @pytest.mark.unittest @pytest.mark.usefixtures("setup") @@ -53,29 +95,7 @@ def test_partition_key_to_feed_range(self, setup): "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) setup["created_db"].delete_container(created_container) - test_ranges = [(Range("", "FF", True, False), - Range("3F", "7F", True, False), - True), - (Range("3F", "7F", True, False), - Range("", "FF", True, False), - False), - (Range("3F", "7F", True, False), - Range("", "5F", True, False), - False), - (Range("3F", "7F", True, True), - Range("3F", "7F", True, True), - True), - (Range("3F", "7F", False, True), - Range("3F", "7F", True, True), - False), - (Range("3F", "7F", True, False), - Range("3F", "7F", True, True), - False), - (Range("3F", "7F", True, False), - Range("", "2F", True, False), - False)] - - @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_ranges) + @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_subset_ranges) def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset): epk_parent_feed_range = FeedRangeEpk(parent_feed_range) epk_child_feed_range = FeedRangeEpk(child_feed_range) @@ -86,5 +106,10 @@ def test_feed_range_is_subset_from_pk(self, setup): epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + @pytest.mark.parametrize("range1, range2, overlaps", test_overlaps_ranges) + def test_overlaps(self, setup, range1, range2, overlaps): + assert Range.overlaps(range1, range2) == overlaps + + if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index bd7cbab2b310..19d89bef55a0 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -9,9 +9,6 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos._routing.routing_range import Range -from test.test_config import TestConfig - - @pytest.fixture(scope="class") def setup(): @@ -21,7 +18,7 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestRequestContext.host, TestConfig.masterKey), + test_client = cosmos_client.CosmosClient(TestRequestContext.host, test_config.TestConfig.masterKey), created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) return { "created_db": created_db, @@ -52,7 +49,7 @@ def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'): @pytest.mark.unittest @pytest.mark.usefixtures("setup") class TestRequestContext: - """Tests to verify methods for operations on feed ranges + """Tests to verify request context gets populated correctly """ host = test_config.TestConfig.host diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 25ca568943cf..21db40f96f22 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -1,5 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. + import random import unittest @@ -8,9 +9,8 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos import DatabaseProxy -from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk +from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range -from test.test_config import TestConfig @pytest.fixture(scope="class") def setup(): @@ -20,7 +20,7 @@ def setup(): "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") - test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, TestConfig.masterKey), + test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, test_config.TestConfig.masterKey), created_db = test_client[0].get_database_client(TestSessionTokenHelpers.TEST_DATABASE_ID) return { "created_db": created_db, @@ -28,19 +28,28 @@ def setup(): } def create_split_ranges(): + # add one with several ranges being equal to one test_params = [([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], - ("AA", "DD"), "0:1#55#3=52")] + ("AA", "DD"), "0:1#55#3=52"), + ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52"), (("BB", "DD"),"2:1#54#3=52")], + ("AA", "DD"), "0:1#55#3=52"), + ([(("AA", "DD"), "2:1#54#3=52,1:1#55#3=52"), (("AA", "BB"),"0:1#51#3=52")], + ("AA", "BB"), "2:1#54#3=52,1:1#55#3=52,0:1#51#3=52"), + ([(("AA", "DD"), "2:1#57#3=52,1:1#57#3=52"), (("AA", "DD"),"2:1#56#3=52,1:1#58#3=52")], + ("AA", "DD"), "2:1#57#3=52,1:1#58#3=52"), + ([(("AA", "CC"), "0:1#54#3=52"), (("BB", "FF"),"2:1#51#3=52")], + ("AA", "EE"), "0:1#54#3=52,2:1#51#3=52")] actual_test_params = [] for test_param in test_params: split_ranges = [] for feed_range, session_token in test_param[0]: - split_ranges.append((FeedRangeInternalEpk(Range(feed_range[0], feed_range[1], + split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], True, False)), session_token)) - target_feed_range = FeedRangeInternalEpk(Range(test_param[1][0], test_param[1][1], True, False)) + target_feed_range = FeedRangeEpk(Range(test_param[1][0], test_param[1][1], True, False)) actual_test_params.append((split_ranges, target_feed_range, test_param[2])) return actual_test_params @@ -48,7 +57,7 @@ def create_split_ranges(): @pytest.mark.unittest @pytest.mark.usefixtures("setup") class TestSessionTokenHelpers: - """Test to ensure escaping of non-ascii characters from partition key""" + """Test for session token helpers""" created_db: DatabaseProxy = None client: cosmos_client.CosmosClient = None @@ -74,26 +83,32 @@ def test_get_session_token_update(self, setup): assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): - feed_range = FeedRangeInternalEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) feed_ranges_and_session_tokens.append((feed_range, session_token)) session_token = "0:1#101#3=101" feed_ranges_and_session_tokens.append((feed_range, session_token)) - updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) + updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, + feed_range) assert updated_session_token == session_token @pytest.mark.parametrize("split_ranges, target_feed_range, expected_session_token", create_split_ranges()) - def test_simulated_splits(self, setup, split_ranges, target_feed_range, expected_session_token): - updated_session_token = (setup["created_collection"] - .get_updated_session_token(split_ranges, target_feed_range)) - assert expected_session_token == updated_session_token - + def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, expected_session_token): + updated_session_token = setup["created_collection"].get_updated_session_token(split_ranges, target_feed_range) + assert updated_session_token == expected_session_token + def test_invalid_feed_range(self, setup): + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + session_token = "0:1#54#3=50" + feed_ranges_and_session_tokens = [(feed_range, session_token)] + with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): + setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, + FeedRangeEpk(Range("CC", "FF", True, False))) if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main() From 9f88b4ed163fd93928a3170dc06ff2bd4d8aa71c Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 7 Oct 2024 12:30:55 -0700 Subject: [PATCH 30/59] Fix tests and add changelog --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure-cosmos/azure/cosmos/container.py | 12 +++---- .../azure-cosmos/test/test_request_context.py | 35 ++++++++++--------- 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index df20a43c36bf..0f7e216ec86b 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,6 +7,7 @@ * Added option to disable write payload on writes. See [PR 37365](https://github.com/Azure/azure-sdk-for-python/pull/37365) * Added get feed ranges API. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) * Added feed range support in `query_items_change_feed`. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) +* Added helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index eecd05f2b3c1..9de67dd61c15 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -523,18 +523,18 @@ def query_items_change_feed( feed_options["maxItemCount"] = kwargs.pop('max_item_count') except KeyError: feed_options["maxItemCount"] = args[3] - + request_context = {} if kwargs.get("partition_key") is not None: change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) - request_context = {"partitionKey": change_feed_state_context["partitionKey"]} + request_context["partitionKey"] = change_feed_state_context["partitionKey"] change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: - change_feed_state_context["feedRange"] = kwargs.pop('feed_range') feed_range: FeedRangeEpk = kwargs.pop('feed_range') change_feed_state_context["feedRange"] = feed_range._feed_range_internal + request_context["feed_range"] = feed_range container_properties = self._get_properties() feed_options["changeFeedStateContext"] = change_feed_state_context @@ -1450,10 +1450,8 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_range()) def _add_request_context(self, request_context): - request_context['session_token'] = self.client_connection.last_response_headers['x-ms-session-token'] + if 'x-ms-session-token' in self.client_connection.last_response_headers: + request_context['session_token'] = self.client_connection.last_response_headers['x-ms-session-token'] if 'partitionKey' in request_context: request_context["feed_range"] = FeedRangeEpk(self._get_epk_range_for_partition_key(request_context['partitionKey'])) - else: - # default to full range - request_context["feed_range"] = FeedRangeEpk(Range("", "FF", True, False)) self.client_connection.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index 19d89bef55a0..0d7df2590f65 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -8,7 +8,6 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config -from azure.cosmos._routing.routing_range import Range @pytest.fixture(scope="class") def setup(): @@ -25,17 +24,19 @@ def setup(): "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) } -def validate_request_context(setup, with_partition_key=True): - request_context = setup["created_collection"].client_connection.last_response_headers["request_context"] - keys_expected = ["session_token", "feed_range"] +def validate_request_context(collection, with_partition_key=True): + request_context = collection.client_connection.last_response_headers["request_context"] + keys_expected = ["session_token"] if with_partition_key: keys_expected.append("partitionKey") + keys_expected.append("feed_range") assert request_context is not None for key in keys_expected: assert request_context[key] is not None if not with_partition_key: - # when operation is not scoped with partition key that operation is done on full feed range - assert request_context["feed_range"]._feed_range_internal.get_normalized_range() == Range("", "FF", True, False) + # when operation is not scoped with partition key feed range will not be present + with pytest.raises(KeyError): + request_context["feed_range"] def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'): item = { @@ -57,36 +58,38 @@ class TestRequestContext: TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID - # test out all operations and using response hook + # test out all operations def test_crud_request_context(self, setup): item = createItem() setup["created_collection"].create_item(item) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) setup["created_collection"].read_item(item['id'], item['pk']) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) new_item = createItem(name='sample_replaced') setup["created_collection"].replace_item(item['id'], new_item) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) setup["created_collection"].upsert_item(createItem()) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) + + setup["created_collection"].query_items_change_feed() + validate_request_context(setup["created_collection"], False) # with partition key setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])], partition_key=item['pk']) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) # without partition key setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])]) - validate_request_context(setup, False) - + validate_request_context(setup["created_collection"], False) setup["created_collection"].read_all_items() - validate_request_context(setup, False) + validate_request_context(setup["created_collection"], False) setup["created_collection"].delete_item(item['id'], item['pk']) - validate_request_context(setup) + validate_request_context(setup["created_collection"]) if __name__ == '__main__': unittest.main() From 7c23e87c87fe52ef783ec3d87d0c5b7833f027b6 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 7 Oct 2024 13:01:26 -0700 Subject: [PATCH 31/59] fix spell checks --- .../azure/cosmos/_cosmos_client_connection.py | 8 ++++---- .../azure-cosmos/test/test_session_token_helpers.py | 13 ++++++++++++- 2 files changed, 16 insertions(+), 5 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 35efabbbeedf..47ca281d40bf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3412,16 +3412,16 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ # if it is the bigger one remove the smaller ranges # if it is the smaller ranges remove the bigger range # if it is neither compound - child_lsns_larger = True + children_more_updated = True for session_token in session_tokens: tokens = session_token.split(":") vector_session_token = VectorSessionToken.create(tokens[1]) if vector_session_token_cmp.is_greater(vector_session_token): - child_lsns_larger = False + children_more_updated = False feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] for feed_range_to_remove in feed_ranges_to_remove: overlapping_ranges.remove(feed_range_to_remove) - if child_lsns_larger: + if children_more_updated: overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) overlapping_ranges.remove(overlapping_ranges[0]) not_found = False @@ -3437,7 +3437,7 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_ if len(remaining_session_tokens) == 1: return remaining_session_tokens[0] - # merging any session tokens with same pkrangeid + # merging any session tokens with same physical partition key range id remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens) # compound the remaining session tokens diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 21db40f96f22..b300bfb06f32 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -29,18 +29,29 @@ def setup(): def create_split_ranges(): # add one with several ranges being equal to one - test_params = [([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], + test_params = [ # split with two children + ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), + # several ranges being equal to one range ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), + # split with one child + ([(("AA", "DD"), "0:1#42#3=52"), (("AA", "BB"), "1:1#51#3=52"), + (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], + ("AA", "DD"), "1:1#55#3=52"), + # merge with one child ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], ("AA", "DD"), "0:1#55#3=52"), + # merge with two children ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "0:1#55#3=52"), + # compound session token ([(("AA", "DD"), "2:1#54#3=52,1:1#55#3=52"), (("AA", "BB"),"0:1#51#3=52")], ("AA", "BB"), "2:1#54#3=52,1:1#55#3=52,0:1#51#3=52"), + # several compound session token with one range ([(("AA", "DD"), "2:1#57#3=52,1:1#57#3=52"), (("AA", "DD"),"2:1#56#3=52,1:1#58#3=52")], ("AA", "DD"), "2:1#57#3=52,1:1#58#3=52"), + # Overlapping ranges ([(("AA", "CC"), "0:1#54#3=52"), (("BB", "FF"),"2:1#51#3=52")], ("AA", "EE"), "0:1#54#3=52,2:1#51#3=52")] actual_test_params = [] From d7c598edaa8717a3a26680c5b7b7c46f13f6bf6f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 7 Oct 2024 23:33:11 -0700 Subject: [PATCH 32/59] Added tests and pushed request context to client level --- .../_change_feed/feed_range_internal.py | 12 +- .../azure/cosmos/_cosmos_client_connection.py | 250 +++++++----------- .../azure/cosmos/_session_token_helpers.py | 143 ++++++++-- .../azure-cosmos/azure/cosmos/container.py | 69 +++-- .../samples/merge_session_tokens.py | 121 --------- .../azure-cosmos/test/test_request_context.py | 16 +- .../test/test_session_token_helpers.py | 156 ++++++++++- 7 files changed, 411 insertions(+), 356 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index c04fda0952f9..ec55a7cc2fab 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -25,11 +25,12 @@ import base64 import json from abc import ABC, abstractmethod -from typing import Union, List, Dict, Any, Optional +from typing import Union, List, Dict, Any, Optional, Sequence, Type from azure.cosmos._routing.routing_range import Range -from azure.cosmos.partition_key import _Undefined, _Empty +from azure.cosmos.partition_key import _Undefined, _Empty, PartitionKey, NonePartitionKeyValue +PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long class FeedRangeInternal(ABC): @@ -130,3 +131,10 @@ def __str__(self) -> str: self._base64_encoded_string = self._to_base64_encoded_string() return self._base64_encoded_string + + @staticmethod + def get_epk_range_for_partition_key(container_properties, partition_key_value: PartitionKeyType) -> Range: + partition_key_definition = container_properties["partitionKey"] + partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) + + return partition_key._get_epk_range_for_partition_key(partition_key_value) \ No newline at end of file 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 47ca281d40bf..e4f0b655cf5a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -30,9 +30,8 @@ from urllib3.util.retry import Retry from azure.core import PipelineClient -from ._session_token_helpers import is_compound_session_token, merge_session_tokens, split_compound_session_tokens, \ - merge_session_tokens_with_same_pkrangeid -from ._vector_session_token import VectorSessionToken +from ._change_feed.feed_range_internal import FeedRangeInternalEpk +from ._feed_range import FeedRangeEpk from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged from azure.core.pipeline.policies import ( @@ -372,7 +371,7 @@ def CreateDatabase( options = {} base._validate_resource(database) path = "/dbs" - return self.Create(database, path, "dbs", None, None, options, **kwargs) + return self.Create(database, path, "dbs", None, None, None, options, **kwargs) def ReadDatabase( self, @@ -397,7 +396,7 @@ def ReadDatabase( path = base.GetPathFromLink(database_link) database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Read(path, "dbs", database_id, None, options, **kwargs) + return self.Read(path, "dbs", database_id, None, None, options, **kwargs) def ReadDatabases( self, @@ -533,7 +532,7 @@ def CreateContainer( base._validate_resource(collection) path = base.GetPathFromLink(database_link, "colls") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Create(collection, path, "colls", database_id, None, options, **kwargs) + return self.Create(collection, path, "colls", database_id, None, None, options, **kwargs) def ReplaceContainer( self, @@ -563,7 +562,7 @@ def ReplaceContainer( base._validate_resource(collection) path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) + return self.Replace(collection, path, "colls", collection_id, None, None, options, **kwargs) def ReadContainer( self, @@ -589,7 +588,7 @@ def ReadContainer( path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Read(path, "colls", collection_id, None, options, **kwargs) + return self.Read(path, "colls", collection_id, None, None, options, **kwargs) def CreateUser( self, @@ -617,7 +616,7 @@ def CreateUser( options = {} database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Create(user, path, "users", database_id, None, options, **kwargs) + return self.Create(user, path, "users", database_id, None, None, options, **kwargs) def UpsertUser( self, @@ -643,7 +642,7 @@ def UpsertUser( options = {} database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Upsert(user, path, "users", database_id, None, options, **kwargs) + return self.Upsert(user, path, "users", database_id, None, None, options, **kwargs) def _GetDatabaseIdWithPathForUser(self, database_link: str, user: Mapping[str, Any]) -> Tuple[Optional[str], str]: base._validate_resource(user) @@ -675,7 +674,7 @@ def ReadUser( path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Read(path, "users", user_id, None, options, **kwargs) + return self.Read(path, "users", user_id, None, None, options, **kwargs) def ReadUsers( self, @@ -760,7 +759,7 @@ def DeleteDatabase( path = base.GetPathFromLink(database_link) database_id = base.GetResourceIdOrFullNameFromLink(database_link) - self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) + self.DeleteResource(path, "dbs", database_id, None, None, options, **kwargs) def CreatePermission( self, @@ -788,7 +787,7 @@ def CreatePermission( options = {} path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Create(permission, path, "permissions", user_id, None, options, **kwargs) + return self.Create(permission, path, "permissions", user_id, None, None, options, **kwargs) def UpsertPermission( self, @@ -816,7 +815,7 @@ def UpsertPermission( options = {} path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) + return self.Upsert(permission, path, "permissions", user_id, None, None, options, **kwargs) def _GetUserIdWithPathForPermission( self, @@ -852,7 +851,7 @@ def ReadPermission( path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Read(path, "permissions", permission_id, None, options, **kwargs) + return self.Read(path, "permissions", permission_id, None, None, options, **kwargs) def ReadPermissions( self, @@ -941,7 +940,7 @@ def ReplaceUser( base._validate_resource(user) path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Replace(user, path, "users", user_id, None, options, **kwargs) + return self.Replace(user, path, "users", user_id, None, None, options, **kwargs) def DeleteUser( self, @@ -967,7 +966,7 @@ def DeleteUser( path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - self.DeleteResource(path, "users", user_id, None, options, **kwargs) + self.DeleteResource(path, "users", user_id, None, None, options, **kwargs) def ReplacePermission( self, @@ -996,7 +995,7 @@ def ReplacePermission( base._validate_resource(permission) path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) + return self.Replace(permission, path, "permissions", permission_id, None, None, options, **kwargs) def DeletePermission( self, @@ -1022,11 +1021,12 @@ def DeletePermission( path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) + self.DeleteResource(path, "permissions", permission_id, None, None, options, **kwargs) def ReadItems( self, collection_link: str, + request_context: Optional[Dict[str, Any]], feed_options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, **kwargs: Any @@ -1042,12 +1042,13 @@ def ReadItems( if feed_options is None: feed_options = {} - return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) + return self.QueryItems(collection_link, None, request_context, feed_options, response_hook=response_hook, **kwargs) def QueryItems( self, database_or_container_link: str, query: Optional[Union[str, Dict[str, Any]]], + request_context: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, partition_key: Optional[PartitionKeyType] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, @@ -1100,7 +1101,7 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str response_hook=response_hook, **kwargs) - return ItemPaged( + items = ItemPaged( self, query, options, @@ -1108,10 +1109,13 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str collection_link=database_or_container_link, page_iterator_class=query_iterable.QueryIterable ) + self._add_request_context(request_context) + return items def QueryItemsChangeFeed( self, collection_link: str, + request_context: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, **kwargs: Any @@ -1134,9 +1138,11 @@ def QueryItemsChangeFeed( if options is not None and "partitionKeyRangeId" in options: partition_key_range_id = options["partitionKeyRangeId"] - return self._QueryChangeFeed( + result = self._QueryChangeFeed( collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs ) + self._add_request_context(request_context) + return result def _QueryChangeFeed( self, @@ -1293,8 +1299,7 @@ def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) request_context["partitionKey"] = options["partitionKey"] - result = self.Create(document, path, "docs", collection_id, None, options, **kwargs) - return result + return self.Create(document, path, "docs", collection_id, None, request_context, options, **kwargs) def UpsertItem( self, @@ -1309,6 +1314,7 @@ def UpsertItem( :param str database_or_container_link: The link to the database when using partitioning, otherwise link to the document collection. :param dict document: The Azure Cosmos document to upsert. + :param dict request_context: The request context for the request. :param dict options: The request options for the request. :return: The upserted Document. :rtype: dict @@ -1332,7 +1338,7 @@ def UpsertItem( collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options ) - return self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) + return self.Upsert(document, path, "docs", collection_id, None, request_context, options, **kwargs) PartitionResolverErrorMessage = ( "Couldn't find any partition resolvers for the database link provided. " @@ -1377,6 +1383,7 @@ def _GetContainerIdWithPathForItem( def ReadItem( self, document_link: str, + request_context: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs ) -> Dict[str, Any]: @@ -1398,7 +1405,7 @@ def ReadItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return self.Read(path, "docs", document_id, None, options, **kwargs) + return self.Read(path, "docs", document_id, None, request_context, options, **kwargs) def ReadTriggers( self, @@ -1485,7 +1492,7 @@ def CreateTrigger( options = {} collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) + return self.Create(trigger, path, "triggers", collection_id, None, None, options, **kwargs) def _GetContainerIdWithPathForTrigger( self, @@ -1527,7 +1534,7 @@ def ReadTrigger( path = base.GetPathFromLink(trigger_link) trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - return self.Read(path, "triggers", trigger_id, None, options, **kwargs) + return self.Read(path, "triggers", trigger_id, None, None, options, **kwargs) def UpsertTrigger( self, @@ -1551,7 +1558,7 @@ def UpsertTrigger( options = {} collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Upsert(trigger, path, "triggers", collection_id, None, options, **kwargs) + return self.Upsert(trigger, path, "triggers", collection_id, None, None, options, **kwargs) def ReadUserDefinedFunctions( self, @@ -1638,7 +1645,7 @@ def CreateUserDefinedFunction( options = {} collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) + return self.Create(udf, path, "udfs", collection_id, None, None, options, **kwargs) def UpsertUserDefinedFunction( self, @@ -1665,7 +1672,7 @@ def UpsertUserDefinedFunction( options = {} collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Upsert(udf, path, "udfs", collection_id, None, options, **kwargs) + return self.Upsert(udf, path, "udfs", collection_id, None, None, options, **kwargs) def _GetContainerIdWithPathForUDF( self, @@ -1707,7 +1714,7 @@ def ReadUserDefinedFunction( path = base.GetPathFromLink(udf_link) udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - return self.Read(path, "udfs", udf_id, None, options, **kwargs) + return self.Read(path, "udfs", udf_id, None, None, options, **kwargs) def ReadStoredProcedures( self, @@ -1794,7 +1801,7 @@ def CreateStoredProcedure( options = {} collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) + return self.Create(sproc, path, "sprocs", collection_id, None, None, options, **kwargs) def UpsertStoredProcedure( self, @@ -1821,7 +1828,7 @@ def UpsertStoredProcedure( options = {} collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Upsert(sproc, path, "sprocs", collection_id, None, options, **kwargs) + return self.Upsert(sproc, path, "sprocs", collection_id, None, None, options, **kwargs) def _GetContainerIdWithPathForSproc( self, @@ -1862,7 +1869,7 @@ def ReadStoredProcedure( path = base.GetPathFromLink(sproc_link) sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - return self.Read(path, "sprocs", sproc_id, None, options, **kwargs) + return self.Read(path, "sprocs", sproc_id, None, None, options, **kwargs) def ReadConflicts( self, @@ -1946,7 +1953,7 @@ def ReadConflict( path = base.GetPathFromLink(conflict_link) conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - return self.Read(path, "conflicts", conflict_id, None, options, **kwargs) + return self.Read(path, "conflicts", conflict_id, None, None, options, **kwargs) def DeleteContainer( self, @@ -1972,13 +1979,13 @@ def DeleteContainer( path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) + self.DeleteResource(path, "colls", collection_id, None, None, options, **kwargs) def ReplaceItem( self, document_link: str, new_document: Dict[str, Any], - request_context: Dict[str, Any], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -1987,6 +1994,7 @@ def ReplaceItem( :param str document_link: The link to the document. :param dict new_document: + :param dict request_context: :param dict options: The request options for the request. @@ -2015,12 +2023,13 @@ def ReplaceItem( options = self._AddPartitionKey(collection_link, new_document, options) request_context["partitionKey"] = options["partitionKey"] - return self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) + return self.Replace(new_document, path, "docs", document_id, None, request_context, options, **kwargs) def PatchItem( self, document_link: str, operations: List[Dict[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2028,6 +2037,7 @@ def PatchItem( :param str document_link: The link to the document. :param list operations: The operations for the patch request. + :param dict request_context: The request context for the request. :param dict options: The request options for the request. :return: @@ -2055,6 +2065,8 @@ def PatchItem( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -2063,6 +2075,7 @@ def Batch( self, collection_link: str, batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> List[Dict[str, Any]]: @@ -2070,6 +2083,7 @@ def Batch( :param str collection_link: The link to the collection :param list batch_operations: The batch of operations for the batch request. + :param dict request_context: The request context of the operation. :param dict options: The request options for the request. :return: @@ -2094,6 +2108,8 @@ def Batch( **kwargs ) self.last_response_headers = last_response_headers + if request_context is not None: + self._add_request_context(request_context) final_responses = [] is_error = False error_status = 0 @@ -2141,6 +2157,7 @@ def _Batch( def DeleteItem( self, document_link: str, + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> None: @@ -2148,6 +2165,7 @@ def DeleteItem( :param str document_link: The link to the document. + :param dict request_context: :param dict options: The request options for the request. @@ -2162,7 +2180,7 @@ def DeleteItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - self.DeleteResource(path, "docs", document_id, None, options, **kwargs) + self.DeleteResource(path, "docs", document_id, None, request_context, options, **kwargs) def DeleteAllItemsByPartitionKey( self, @@ -2262,7 +2280,7 @@ def DeleteTrigger( path = base.GetPathFromLink(trigger_link) trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) + self.DeleteResource(path, "triggers", trigger_id, None, None, options, **kwargs) def ReplaceUserDefinedFunction( self, @@ -2323,7 +2341,7 @@ def DeleteUserDefinedFunction( path = base.GetPathFromLink(udf_link) udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) + self.DeleteResource(path, "udfs", udf_id, None, None, options, **kwargs) def ExecuteStoredProcedure( self, @@ -2424,7 +2442,7 @@ def DeleteStoredProcedure( path = base.GetPathFromLink(sproc_link) sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) + self.DeleteResource(path, "sprocs", sproc_id, None, None, options, **kwargs) def DeleteConflict( self, @@ -2450,7 +2468,7 @@ def DeleteConflict( path = base.GetPathFromLink(conflict_link) conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) + self.DeleteResource(path, "conflicts", conflict_id, None, None, options, **kwargs) def ReplaceOffer( self, @@ -2490,7 +2508,7 @@ def ReadOffer( """ path = base.GetPathFromLink(offer_link) offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Read(path, "offers", offer_id, None, {}, **kwargs) + return self.Read(path, "offers", offer_id, None, None, {}, **kwargs) def ReadOffers( self, @@ -2596,6 +2614,7 @@ def Create( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2606,6 +2625,7 @@ def Create( :param str typ: :param str id: :param dict initial_headers: + :param dict request_context: :param dict options: The request options for the request. @@ -2629,6 +2649,8 @@ def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -2640,6 +2662,7 @@ def Upsert( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2650,6 +2673,7 @@ def Upsert( :param str typ: :param str id: :param dict initial_headers: + :param dict request_context: :param dict options: The request options for the request. @@ -2673,6 +2697,8 @@ def Upsert( self.last_response_headers = last_response_headers # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -2684,6 +2710,7 @@ def Replace( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2694,6 +2721,7 @@ def Replace( :param str typ: :param str id: :param dict initial_headers: + :param dict request_context: :param dict options: The request options for the request. @@ -2716,6 +2744,8 @@ def Replace( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -2726,6 +2756,7 @@ def Read( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: @@ -2735,6 +2766,7 @@ def Read( :param str typ: :param str id: :param dict initial_headers: + :param dict request_context: :param dict options: The request options for the request. @@ -2754,6 +2786,8 @@ def Read( request_params = RequestObject(typ, documents._OperationType.Read) result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -2764,6 +2798,7 @@ def DeleteResource( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], + request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> None: @@ -2773,6 +2808,7 @@ def DeleteResource( :param str typ: :param str id: :param dict initial_headers: + :param dict request_context: :param dict options: The request options for the request. @@ -2795,6 +2831,8 @@ def DeleteResource( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, None) @@ -3337,114 +3375,14 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s self.__container_properties_cache[collection_link] = _set_properties_cache(container) return partition_key_definition - def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_range): - target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() - # filter out tuples that overlap with target_feed_range and normalizes all the ranges - overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1]) - for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( - target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] - # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples - # with compound session tokens? - if len(overlapping_ranges) == 0: - raise ValueError('There were no overlapping feed ranges with the target.') - - # merge any session tokens that are the same exact feed range - i = 0 - j = 1 - while i < len(overlapping_ranges) and j < len(overlapping_ranges): - cur_feed_range = overlapping_ranges[i][0] - session_token = overlapping_ranges[i][1] - session_token_1 = overlapping_ranges[j][1] - if (not is_compound_session_token(session_token) and - not is_compound_session_token(overlapping_ranges[j][1]) and - cur_feed_range == overlapping_ranges[j][0]): - session_token = merge_session_tokens(session_token, session_token_1) - feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] - for feed_range_to_remove in feed_ranges_to_remove: - overlapping_ranges.remove(feed_range_to_remove) - overlapping_ranges.append((cur_feed_range, session_token)) - else: - j += 1 - if j == len(overlapping_ranges): - i += 1 - j = i + 1 - - - updated_session_token = "" - done_overlapping_ranges = [] - while len(overlapping_ranges) != 0: - feed_range_cmp, session_token_cmp = overlapping_ranges[0] - if is_compound_session_token(session_token_cmp): - done_overlapping_ranges.append(overlapping_ranges[0]) - overlapping_ranges.remove(overlapping_ranges[0]) - continue - tokens_cmp = session_token_cmp.split(":") - vector_session_token_cmp = VectorSessionToken.create(tokens_cmp[1]) - subsets = [] - # creating subsets of feed ranges that are subsets of the current feed range - for j in range(1, len(overlapping_ranges)): - feed_range = overlapping_ranges[j][0] - if not is_compound_session_token(overlapping_ranges[j][1]) and \ - feed_range.is_subset(feed_range_cmp): - subsets.append(overlapping_ranges[j] + (j,)) - - # go through subsets to see if can create current feed range from the subsets - not_found = True - j = 0 - while not_found and j < len(subsets): - merged_range = subsets[j][0] - session_tokens = [subsets[j][1]] - merged_indices = [subsets[j][2]] - if len(subsets) == 1: - tokens = session_tokens[0].split(":") - vector_session_token = VectorSessionToken.create(tokens[1]) - if vector_session_token_cmp.is_greater(vector_session_token): - overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) - else: - for k in range(len(subsets)): - if j == k: - continue - if merged_range.can_merge(subsets[k][0]): - merged_range = merged_range.merge(subsets[k][0]) - session_tokens.append(subsets[k][1]) - merged_indices.append(subsets[k][2]) - if feed_range_cmp == merged_range: - # if it is the bigger one remove the smaller ranges - # if it is the smaller ranges remove the bigger range - # if it is neither compound - children_more_updated = True - for session_token in session_tokens: - tokens = session_token.split(":") - vector_session_token = VectorSessionToken.create(tokens[1]) - if vector_session_token_cmp.is_greater(vector_session_token): - children_more_updated = False - feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] - for feed_range_to_remove in feed_ranges_to_remove: - overlapping_ranges.remove(feed_range_to_remove) - if children_more_updated: - overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) - overlapping_ranges.remove(overlapping_ranges[0]) - not_found = False - break - - j += 1 - - done_overlapping_ranges.append(overlapping_ranges[0]) - overlapping_ranges.remove(overlapping_ranges[0]) - - # break up session tokens that are compound - remaining_session_tokens = split_compound_session_tokens(done_overlapping_ranges) - - if len(remaining_session_tokens) == 1: - return remaining_session_tokens[0] - # merging any session tokens with same physical partition key range id - remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens) - - # compound the remaining session tokens - for i in range(len(remaining_session_tokens)): - if i == len(remaining_session_tokens) - 1: - updated_session_token += remaining_session_tokens[i] - else: - updated_session_token += remaining_session_tokens[i] + "," - - return updated_session_token + def _add_request_context(self, request_context): + if http_constants.HttpHeaders.SessionToken in self.last_response_headers: + request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] + if 'partitionKey' in request_context: + request_context["feed_range"] = FeedRangeEpk(FeedRangeInternalEpk.get_epk_range_for_partition_key( + request_context['container_properties'], + request_context['partitionKey'])) + request_context.pop('partitionKey') + if 'container_properties' in request_context: + request_context.pop('container_properties') + self.last_response_headers["request_context"] = request_context \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index dbb07a4a066b..dd32019607bc 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -21,18 +21,13 @@ """Internal Helper functions for manipulating session tokens. """ - +from azure.cosmos._routing.routing_range import Range from azure.cosmos._vector_session_token import VectorSessionToken -def merge_session_tokens(session_token1, session_token2): - # TODO method for splitting the session token into pk_range_id and vector_session_token - token_pairs1 = session_token1.split(":") - pk_range_id1 = token_pairs1[0] - vector_session_token1 = VectorSessionToken.create(token_pairs1[1]) - token_pairs2 = session_token2.split(":") - pk_range_id2 = token_pairs2[0] - vector_session_token2 = VectorSessionToken.create(token_pairs2[1]) +def merge_session_tokens_with_same_range(session_token1, session_token2): + pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrangeid(session_token1) + pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrangeid(session_token2) pk_range_id = pk_range_id1 if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ @@ -43,6 +38,10 @@ def merge_session_tokens(session_token1, session_token2): def is_compound_session_token(session_token): return "," in session_token +def create_vector_session_token_and_pkrangeid(session_token): + tokens = session_token.split(":") + return tokens[0], VectorSessionToken.create(tokens[1]) + def split_compound_session_tokens(compound_session_tokens): session_tokens = [] for _, session_token in compound_session_tokens: @@ -55,20 +54,15 @@ def split_compound_session_tokens(compound_session_tokens): return session_tokens def merge_session_tokens_with_same_pkrangeid(session_tokens): - new_session_tokens = [] i = 0 while i < len(session_tokens): j = i + 1 while j < len(session_tokens): - tokens1 = session_tokens[i].split(":") - tokens2 = session_tokens[j].split(":") - pk_range_id1 = tokens1[0] - pk_range_id2 = tokens2[0] + pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrangeid(session_tokens[i]) + pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrangeid(session_tokens[j]) if pk_range_id1 == pk_range_id2: - vector_session_token1 = VectorSessionToken.create(tokens1[1]) - vector_session_token2 = VectorSessionToken.create(tokens2[1]) vector_session_token = vector_session_token1.merge(vector_session_token2) - new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) + session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) remove_session_tokens = [session_tokens[i], session_tokens[j]] for token in remove_session_tokens: session_tokens.remove(token) @@ -76,5 +70,116 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): j += 1 i += 1 - new_session_tokens.extend(session_tokens) - return new_session_tokens + return session_tokens + +def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): + target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() + # filter out tuples that overlap with target_feed_range and normalizes all the ranges + overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1]) + for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( + target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] + # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples + # with compound session tokens? + if len(overlapping_ranges) == 0: + raise ValueError('There were no overlapping feed ranges with the target.') + + # merge any session tokens that are the same exact feed range + i = 0 + j = 1 + while i < len(overlapping_ranges) and j < len(overlapping_ranges): + cur_feed_range = overlapping_ranges[i][0] + session_token = overlapping_ranges[i][1] + session_token_1 = overlapping_ranges[j][1] + if (not is_compound_session_token(session_token) and + not is_compound_session_token(overlapping_ranges[j][1]) and + cur_feed_range == overlapping_ranges[j][0]): + session_token = merge_session_tokens_with_same_range(session_token, session_token_1) + feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + overlapping_ranges.append((cur_feed_range, session_token)) + else: + j += 1 + if j == len(overlapping_ranges): + i += 1 + j = i + 1 + + + done_overlapping_ranges = [] + # checking for merging of feed ranges that can be created from other feed ranges + while len(overlapping_ranges) != 0: + feed_range_cmp, session_token_cmp = overlapping_ranges[0] + # compound session tokens are not considered for merging + if is_compound_session_token(session_token_cmp): + done_overlapping_ranges.append(overlapping_ranges[0]) + overlapping_ranges.remove(overlapping_ranges[0]) + continue + tokens_cmp = session_token_cmp.split(":") + vector_session_token_cmp = VectorSessionToken.create(tokens_cmp[1]) + subsets = [] + # finding the subset feed ranges of the current feed range + for j in range(1, len(overlapping_ranges)): + feed_range = overlapping_ranges[j][0] + if not is_compound_session_token(overlapping_ranges[j][1]) and \ + feed_range.is_subset(feed_range_cmp): + subsets.append(overlapping_ranges[j] + (j,)) + + # go through subsets to see if can create current feed range from the subsets + not_found = True + j = 0 + while not_found and j < len(subsets): + merged_range = subsets[j][0] + session_tokens = [subsets[j][1]] + merged_indices = [subsets[j][2]] + if len(subsets) == 1: + _, vector_session_token = create_vector_session_token_and_pkrangeid(session_tokens[0]) + if vector_session_token_cmp.is_greater(vector_session_token): + overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) + else: + for k in range(len(subsets)): + if j == k: + continue + if merged_range.can_merge(subsets[k][0]): + merged_range = merged_range.merge(subsets[k][0]) + session_tokens.append(subsets[k][1]) + merged_indices.append(subsets[k][2]) + if feed_range_cmp == merged_range: + # if feed range can be created from the subsets + # take the subsets if their global lsn is larger + # else take the current feed range + children_more_updated = True + for session_token in session_tokens: + _, vector_session_token = create_vector_session_token_and_pkrangeid(session_token) + if vector_session_token_cmp.is_greater(vector_session_token): + children_more_updated = False + feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + if children_more_updated: + overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) + overlapping_ranges.remove(overlapping_ranges[0]) + not_found = False + break + + j += 1 + + done_overlapping_ranges.append(overlapping_ranges[0]) + overlapping_ranges.remove(overlapping_ranges[0]) + + # break up session tokens that are compound + remaining_session_tokens = split_compound_session_tokens(done_overlapping_ranges) + + if len(remaining_session_tokens) == 1: + return remaining_session_tokens[0] + # merging any session tokens with same physical partition key range id + remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens) + + updated_session_token = "" + # compound the remaining session tokens + for i in range(len(remaining_session_tokens)): + if i == len(remaining_session_tokens) - 1: + updated_session_token += remaining_session_tokens[i] + else: + updated_session_token += remaining_session_tokens[i] + "," + + return updated_session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 9de67dd61c15..ae34e34e0a7f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -38,9 +38,11 @@ GenerateGuidId, _set_properties_cache ) +from ._change_feed.feed_range_internal import FeedRangeInternalEpk from ._cosmos_client_connection import CosmosClientConnection from ._feed_range import FeedRange, FeedRangeEpk from ._routing.routing_range import Range +from ._session_token_helpers import get_updated_session_token from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -131,12 +133,7 @@ def _set_partition_key( 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) - def _get_epk_range_for_partition_key( self, partition_key_value: PartitionKeyType) -> Range: - container_properties = self._get_properties() - partition_key_definition = container_properties["partitionKey"] - partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) - return partition_key._get_epk_range_for_partition_key(partition_key_value) def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: return self.client_connection._container_properties_cache @@ -257,10 +254,10 @@ def read_item( # pylint:disable=docstring-missing-param request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {"partitionKey": request_options["partitionKey"]} - items = self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) - self._add_request_context(request_context) - return items + request_context = {"partitionKey": request_options["partitionKey"], + "container_properties": self._get_properties()} + return self.client_connection.ReadItem(document_link=doc_link, options=request_options, + request_context=request_context, **kwargs) @distributed_trace def read_all_items( # pylint:disable=docstring-missing-param @@ -315,9 +312,10 @@ def read_all_items( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + request_context = {"container_properties": self._get_properties()} items = self.client_connection.ReadItems( - collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs) - self._add_request_context({}) + collection_link=self.container_link, request_context=request_context, feed_options=feed_options, + response_hook=response_hook, **kwargs) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -527,9 +525,9 @@ def query_items_change_feed( if kwargs.get("partition_key") is not None: change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) - request_context["partitionKey"] = change_feed_state_context["partitionKey"] change_feed_state_context["partitionKeyFeedRange"] =\ - self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) + FeedRangeInternalEpk.get_epk_range_for_partition_key(self._get_properties(), kwargs.pop('partition_key')) + request_context["feed_range"] = FeedRangeEpk(change_feed_state_context["partitionKeyFeedRange"]) if kwargs.get("feed_range") is not None: feed_range: FeedRangeEpk = kwargs.pop('feed_range') @@ -545,9 +543,9 @@ def query_items_change_feed( response_hook.clear() result = self.client_connection.QueryItemsChangeFeed( - self.container_link, options=feed_options, response_hook=response_hook, **kwargs + self.container_link, request_context=request_context, + options=feed_options, response_hook=response_hook, **kwargs ) - self._add_request_context(request_context) if response_hook: response_hook(self.client_connection.last_response_headers, result) return result @@ -666,15 +664,16 @@ def query_items( # pylint:disable=docstring-missing-param 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"] + request_context["container_properties"] = self._get_properties() items = self.client_connection.QueryItems( database_or_container_link=self.container_link, query=query if parameters is None else {"query": query, "parameters": parameters}, + request_context=request_context, options=feed_options, partition_key=partition_key, response_hook=response_hook, **kwargs ) - self._add_request_context(request_context) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -759,11 +758,10 @@ def replace_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {} + request_context = {'container_properties': self._get_properties()} result = self.client_connection.ReplaceItem( document_link=item_link, new_document=body, request_context=request_context, options=request_options, **kwargs ) - self._add_request_context(request_context) return result or {} @distributed_trace @@ -833,7 +831,7 @@ def upsert_item( # pylint:disable=docstring-missing-param request_options["populateQueryMetrics"] = populate_query_metrics if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {} + request_context = {'container_properties': self._get_properties()} result = self.client_connection.UpsertItem( database_or_container_link=self.container_link, document=body, @@ -841,7 +839,6 @@ def upsert_item( # pylint:disable=docstring-missing-param options=request_options, **kwargs ) - self._add_request_context(request_context) return result or {} @distributed_trace @@ -919,10 +916,9 @@ def create_item( # pylint:disable=docstring-missing-param request_options["indexingDirective"] = indexing_directive if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {} + request_context = {'container_properties': self._get_properties()} result = self.client_connection.CreateItem( database_or_container_link=self.container_link, document=body, request_context=request_context, options=request_options, **kwargs) - self._add_request_context(request_context) return result or {} @distributed_trace @@ -996,10 +992,10 @@ def patch_item( if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] item_link = self._get_document_link(item) - request_context = {} + request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} result = self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, request_context=request_context, options=request_options, **kwargs) - self._add_request_context(request_context) + document_link=item_link, operations=patch_operations, request_context=request_context, + options=request_options, **kwargs) return result or {} @distributed_trace @@ -1051,11 +1047,11 @@ def execute_item_batch( kwargs['priority'] = priority request_options = build_options(kwargs) request_options["partitionKey"] = self._set_partition_key(partition_key) - request_context = {"partitionKey": request_options["partitionKey"]} + request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} request_options["disableAutomaticIdGeneration"] = True result = self.client_connection.Batch( - collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) - self._add_request_context(request_context) + collection_link=self.container_link, batch_operations=batch_operations, request_context=request_context, + options=request_options, **kwargs) return result @distributed_trace @@ -1123,9 +1119,9 @@ def delete_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) - request_context = {"partitionKey": request_options["partitionKey"]} - self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) - self._add_request_context(request_context) + request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} + self.client_connection.DeleteItem(document_link=document_link, request_context=request_context, + options=request_options, **kwargs) @distributed_trace def read_offer(self, **kwargs: Any) -> Offer: @@ -1427,7 +1423,7 @@ def get_updated_session_token(self, :returns: a session token :rtype: str """ - return self.client_connection._get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. @@ -1436,7 +1432,7 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Feed :returns: a feed range :rtype: Range """ - return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) + return FeedRangeEpk(FeedRangeInternalEpk.get_epk_range_for_partition_key(self._get_properties(), partition_key)) def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. @@ -1449,9 +1445,4 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F """ return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_range()) - def _add_request_context(self, request_context): - if 'x-ms-session-token' in self.client_connection.last_response_headers: - request_context['session_token'] = self.client_connection.last_response_headers['x-ms-session-token'] - if 'partitionKey' in request_context: - request_context["feed_range"] = FeedRangeEpk(self._get_epk_range_for_partition_key(request_context['partitionKey'])) - self.client_connection.last_response_headers["request_context"] = request_context + diff --git a/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py b/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py deleted file mode 100644 index 726cbbf8f3f4..000000000000 --- a/sdk/cosmos/azure-cosmos/samples/merge_session_tokens.py +++ /dev/null @@ -1,121 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See LICENSE.txt in the project root for -# license information. -# ------------------------------------------------------------------------- -import os - -import azure.cosmos.cosmos_client as cosmos_client -import azure.cosmos.exceptions as exceptions -from azure.cosmos.aio import CosmosClient -from azure.cosmos.http_constants import StatusCodes -from azure.cosmos.partition_key import PartitionKey -import datetime - -import config - -# ---------------------------------------------------------------------------------------------------------- -# Prerequisites - -# -# 1. An Azure Cosmos account - -# https:#azure.microsoft.com/en-us/documentation/articles/documentdb-create-account/ -# -# 2. Microsoft Azure Cosmos PyPi package - -# https://pypi.python.org/pypi/azure-cosmos/ -# ---------------------------------------------------------------------------------------------------------- -# Sample - demonstrates how to merge session tokens using different strategies for storing the session token -# ---------------------------------------------------------------------------------------------------------- - - - -# url = os.environ["ACCOUNT_URI"] -# key = os.environ["ACCOUNT_KEY"] -# client = CosmosClient(url, key) -# # Create a database in the account using the CosmosClient, -# # specifying that the operation shouldn't throw an exception -# # if a database with the given ID already exists. -# # [START create_database] -# database_name = "testDatabase" -# try: -# database = client.create_database(id=database_name) -# except exceptions.CosmosResourceExistsError: -# database = client.get_database_client(database=database_name) -# # [END create_database] -# -# # Create a container, handling the exception if a container with the -# # same ID (name) already exists in the database. -# # [START create_container] -# container_name = "products" -# try: -# container = database.create_container( -# id=container_name, partition_key=PartitionKey(path="/productName") -# ) -# except exceptions.CosmosResourceExistsError: -# container = database.get_container_client(container_name) -# # [END create_container] -# -# # This would be happening through different clients -# # Using physical partition model for read operations -# cache = {} -# session_token = "" -# feed_range = container.feed_range_for_logical_partition(logical_pk) -# for stored_feed_range, stored_session_token in cache: -# if is_feed_range_subset(stored_feed_range, feed_range): -# session_token = stored_session_token -# read_item = container.read_item(doc_to_read, logical_pk, session_token) -# -# # the feed range returned in the request context will correspond to the logical partition key -# logical_pk_feed_range = container.client_connection.last_response_headers["request-context"]["feed-range"] -# session_token = container.client_connection.last_response_headers["request-context"]["session-token"] -# feed_ranges_and_session_tokens = [] -# -# # Get feed ranges for physical partitions -# container_feed_ranges = container.read_feed_ranges() -# target_feed_range = "" -# -# # which feed range maps to the logical pk from the operation -# for feed_range in container_feed_ranges: -# if is_feed_range_subset(feed_range, logical_pk_feed_range): -# target_feed_range = feed_range -# break -# for cached_feed_range, cached_session_token in cache: -# feed_ranges_and_session_tokens.append((cached_feed_range, cached_session_token)) -# # Add the target feed range and session token from the operation -# feed_ranges_and_session_tokens.append((target_feed_range, session_token)) -# cache[feed_range] = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) -# -# -# -# # Different ways of storing the session token and how to get most updated session token -# -# # ---------------------1. using logical partition key --------------------------------------------------- -# # could also use the one stored from the responses headers -# target_feed_range = container.feed_range_for_logical_partition(logical_pk) -# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) -# # ---------------------2. using artificial feed range ---------------------------------------------------- -# # Get four artificial feed ranges -# container_feed_ranges = container.read_feed_ranges(4) -# -# pk_feed_range = container.feed_range_for_logical_partition(logical_pk) -# target_feed_range = "" -# # which feed range maps to the logical pk from the operation -# for feed_range in container_feed_ranges: -# if is_feed_range_subset(feed_range, pk_feed_range): -# target_feed_range = feed_range -# break -# -# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) -# # ---------------------3. using physical partitions ----------------------------------------------------- -# # Get feed ranges for physical partitions -# container_feed_ranges = container.read_feed_ranges() -# -# pk_feed_range = container.feed_range_for_logical_partition(logical_pk) -# target_feed_range = "" -# # which feed range maps to the logical pk from the operation -# for feed_range in container_feed_ranges: -# if is_feed_range_subset(feed_range, pk_feed_range): -# target_feed_range = feed_range -# break -# -# updated_session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) -# # ------------------------------------------------------------------------------------------------------ \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index 0d7df2590f65..f62342671dd8 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -28,15 +28,10 @@ def validate_request_context(collection, with_partition_key=True): request_context = collection.client_connection.last_response_headers["request_context"] keys_expected = ["session_token"] if with_partition_key: - keys_expected.append("partitionKey") keys_expected.append("feed_range") assert request_context is not None for key in keys_expected: assert request_context[key] is not None - if not with_partition_key: - # when operation is not scoped with partition key feed range will not be present - with pytest.raises(KeyError): - request_context["feed_range"] def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'): item = { @@ -67,9 +62,18 @@ def test_crud_request_context(self, setup): setup["created_collection"].read_item(item['id'], item['pk']) validate_request_context(setup["created_collection"]) - new_item = createItem(name='sample_replaced') + new_item = createItem(item['id'], name='sample_replaced') setup["created_collection"].replace_item(item['id'], new_item) validate_request_context(setup["created_collection"]) + operations = [ + {"op": "add", "path": "/favorite_color", "value": "red"}, + {"op": "replace", "path": "/name", "value": 14}, + ] + setup["created_collection"].patch_item(item['id'], item['pk'], operations) + validate_request_context(setup["created_collection"], False) + + setup["created_collection"].read_all_items() + validate_request_context(setup["created_collection"], False) setup["created_collection"].upsert_item(createItem()) validate_request_context(setup["created_collection"]) diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index b300bfb06f32..a043665f21ab 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -2,7 +2,9 @@ # Copyright (c) Microsoft Corporation. All rights reserved. import random +import time import unittest +import uuid import pytest @@ -12,6 +14,8 @@ from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range +COLLECTION = "created_collection" +DATABASE = "created_db" @pytest.fixture(scope="class") def setup(): if (TestSessionTokenHelpers.masterKey == '[YOUR_KEY_HERE]' or @@ -23,8 +27,8 @@ def setup(): test_client = cosmos_client.CosmosClient(TestSessionTokenHelpers.host, test_config.TestConfig.masterKey), created_db = test_client[0].get_database_client(TestSessionTokenHelpers.TEST_DATABASE_ID) return { - "created_db": created_db, - "created_collection": created_db.get_container_client(TestSessionTokenHelpers.TEST_COLLECTION_ID) + DATABASE: created_db, + COLLECTION: created_db.get_container_client(TestSessionTokenHelpers.TEST_COLLECTION_ID) } def create_split_ranges(): @@ -32,13 +36,17 @@ def create_split_ranges(): test_params = [ # split with two children ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), - # several ranges being equal to one range + # split with one child ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), - # split with one child + # several ranges being equal to one range ([(("AA", "DD"), "0:1#42#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "1:1#55#3=52"), + # several ranges being equal to one range + ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), + (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], + ("AA", "DD"), "0:1#60#3=52"), # merge with one child ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], ("AA", "DD"), "0:1#55#3=52"), @@ -51,9 +59,13 @@ def create_split_ranges(): # several compound session token with one range ([(("AA", "DD"), "2:1#57#3=52,1:1#57#3=52"), (("AA", "DD"),"2:1#56#3=52,1:1#58#3=52")], ("AA", "DD"), "2:1#57#3=52,1:1#58#3=52"), - # Overlapping ranges + # overlapping ranges ([(("AA", "CC"), "0:1#54#3=52"), (("BB", "FF"),"2:1#51#3=52")], - ("AA", "EE"), "0:1#54#3=52,2:1#51#3=52")] + ("AA", "EE"), "0:1#54#3=52,2:1#51#3=52"), + # different version numbers + ([(("AA", "BB"), "0:1#54#3=52"), (("AA", "BB"),"0:2#57#3=53")], + ("AA", "BB"), "0:2#57#3=53") + ] actual_test_params = [] for test_param in test_params: split_ranges = [] @@ -64,6 +76,67 @@ def create_split_ranges(): actual_test_params.append((split_ranges, target_feed_range, test_param[2])) return actual_test_params +def trigger_split(setup): + print("Triggering a split in session token helpers") + setup[COLLECTION].replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = setup[COLLECTION].get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in session token helpers has completed") + +def create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens): + target_session_token = "" + for i in range(100): + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + setup[COLLECTION].create_item(item, session_token=previous_session_token) + request_context = setup[COLLECTION].client_connection.last_response_headers["request_context"] + if item['pk'] == target_pk: + target_session_token = request_context["session_token"] + previous_session_token = request_context["session_token"] + feed_ranges_and_session_tokens.append((request_context["feed_range"], request_context["session_token"])) + return target_session_token + +def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): + target_session_token = "" + container_feed_ranges = setup[COLLECTION].read_feed_ranges() + target_feed_range = None + for feed_range in container_feed_ranges: + if setup[COLLECTION].is_feed_range_subset(feed_range, pk_feed_range): + target_feed_range = feed_range + break + + for i in range(100): + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + setup[COLLECTION].create_item(item, session_token=previous_session_token) + request_context = setup[COLLECTION].client_connection.last_response_headers["request_context"] + if setup[COLLECTION].is_feed_range_subset(target_feed_range, request_context["feed_range"]): + target_session_token = request_context["session_token"] + previous_session_token = request_context["session_token"] + feed_ranges_and_session_tokens.append((request_context["feed_range"], request_context["session_token"])) + + return target_session_token, target_feed_range + @pytest.mark.cosmosEmulator @pytest.mark.unittest @pytest.mark.usefixtures("setup") @@ -79,18 +152,13 @@ class TestSessionTokenHelpers: TEST_DATABASE_ID = configs.TEST_DATABASE_ID TEST_COLLECTION_ID = configs.TEST_MULTI_PARTITION_CONTAINER_ID - - # have a test for split, merge, different versions, compound session tokens, hpk, and for normal case for several session tokens - # have tests for all the scenarios in the testing plan - # should try once with 35000 session tokens - # check if pkrangeid is being filtered out def test_get_session_token_update(self, setup): feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] session_token = "0:1#51#3=52" feed_ranges_and_session_tokens.append((feed_range, session_token)) - session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) + session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): @@ -105,9 +173,31 @@ def test_many_session_tokens_update_same_range(self, setup): feed_range) assert updated_session_token == session_token + def test_many_session_tokens_update(self, setup): + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_ranges_and_session_tokens = [] + for i in range(1000): + session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) + feed_ranges_and_session_tokens.append((feed_range, session_token)) + + # adding irrelevant feed ranges + feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False)) + feed_range2 = FeedRangeEpk(Range("00", "55", True, False)) + for i in range(1000): + session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) + if i % 2 == 0: + feed_ranges_and_session_tokens.append((feed_range1, session_token)) + else: + feed_ranges_and_session_tokens.append((feed_range2, session_token)) + session_token = "0:1#101#3=101" + feed_ranges_and_session_tokens.append((feed_range, session_token)) + updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, + feed_range) + assert updated_session_token == session_token + @pytest.mark.parametrize("split_ranges, target_feed_range, expected_session_token", create_split_ranges()) def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, expected_session_token): - updated_session_token = setup["created_collection"].get_updated_session_token(split_ranges, target_feed_range) + updated_session_token = setup[COLLECTION].get_updated_session_token(split_ranges, target_feed_range) assert updated_session_token == expected_session_token def test_invalid_feed_range(self, setup): @@ -118,6 +208,46 @@ def test_invalid_feed_range(self, setup): setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, FeedRangeEpk(Range("CC", "FF", True, False))) + def test_updated_session_token_from_logical_pk(self, setup): + feed_ranges_and_session_tokens = [] + previous_session_token = "" + target_pk = 'A1' + target_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) + target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) + session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert target_session_token == session_token + + trigger_split(setup) + + target_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) + target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) + session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert target_session_token == session_token + + + def test_updated_session_token_from_physical_pk(self, setup): + feed_ranges_and_session_tokens = [] + previous_session_token = "" + pk_feed_range = setup[COLLECTION].feed_range_from_partition_key('A1') + target_session_token, target_feed_range = create_items_physical_pk(setup, pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens) + + session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert target_session_token == session_token + + trigger_split(setup) + + target_session_token, target_feed_range = create_items_physical_pk(setup, pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens) + + session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert target_session_token == session_token + + From 8698098ca698bd9d26fdd74a2533262e19f5920e Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 8 Oct 2024 11:20:11 -0700 Subject: [PATCH 33/59] Added async methods and removed feed range from request context --- .../_change_feed/feed_range_internal.py | 8 - .../azure/cosmos/_cosmos_client_connection.py | 137 +++++++++--------- .../azure/cosmos/_session_token_helpers.py | 2 +- .../azure/cosmos/aio/_container.py | 66 ++++++++- .../aio/_cosmos_client_connection_async.py | 77 +++++++++- .../azure-cosmos/azure/cosmos/container.py | 69 +++++---- .../azure-cosmos/test/test_request_context.py | 38 +++-- .../test/test_request_context_async.py | 99 +++++++++++++ .../test/test_session_token_helpers.py | 13 +- 9 files changed, 356 insertions(+), 153 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/test/test_request_context_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index ec55a7cc2fab..4884da3d1ea3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -30,7 +30,6 @@ from azure.cosmos._routing.routing_range import Range from azure.cosmos.partition_key import _Undefined, _Empty, PartitionKey, NonePartitionKeyValue -PartitionKeyType = Union[str, int, float, bool, Sequence[Union[str, int, float, bool, None]], Type[NonePartitionKeyValue]] # pylint: disable=line-too-long class FeedRangeInternal(ABC): @@ -131,10 +130,3 @@ def __str__(self) -> str: self._base64_encoded_string = self._to_base64_encoded_string() return self._base64_encoded_string - - @staticmethod - def get_epk_range_for_partition_key(container_properties, partition_key_value: PartitionKeyType) -> Range: - partition_key_definition = container_properties["partitionKey"] - partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) - - return partition_key._get_epk_range_for_partition_key(partition_key_value) \ No newline at end of file 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 e4f0b655cf5a..047892a6c38e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -30,8 +30,6 @@ from urllib3.util.retry import Retry from azure.core import PipelineClient -from ._change_feed.feed_range_internal import FeedRangeInternalEpk -from ._feed_range import FeedRangeEpk from azure.core.credentials import TokenCredential from azure.core.paging import ItemPaged from azure.core.pipeline.policies import ( @@ -67,7 +65,6 @@ from ._retry_utility import ConnectionRetryPolicy from ._routing import routing_map_provider, routing_range from .documents import ConnectionPolicy, DatabaseAccount -from azure.cosmos._routing.routing_range import Range from .partition_key import ( _Undefined, _Empty, @@ -371,7 +368,7 @@ def CreateDatabase( options = {} base._validate_resource(database) path = "/dbs" - return self.Create(database, path, "dbs", None, None, None, options, **kwargs) + return self.Create(database, path, "dbs", None, None, options, **kwargs) def ReadDatabase( self, @@ -396,7 +393,7 @@ def ReadDatabase( path = base.GetPathFromLink(database_link) database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Read(path, "dbs", database_id, None, None, options, **kwargs) + return self.Read(path, "dbs", database_id, None, options, **kwargs) def ReadDatabases( self, @@ -532,7 +529,7 @@ def CreateContainer( base._validate_resource(collection) path = base.GetPathFromLink(database_link, "colls") database_id = base.GetResourceIdOrFullNameFromLink(database_link) - return self.Create(collection, path, "colls", database_id, None, None, options, **kwargs) + return self.Create(collection, path, "colls", database_id, None, options, **kwargs) def ReplaceContainer( self, @@ -562,7 +559,7 @@ def ReplaceContainer( base._validate_resource(collection) path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Replace(collection, path, "colls", collection_id, None, None, options, **kwargs) + return self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) def ReadContainer( self, @@ -588,7 +585,7 @@ def ReadContainer( path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - return self.Read(path, "colls", collection_id, None, None, options, **kwargs) + return self.Read(path, "colls", collection_id, None, options, **kwargs) def CreateUser( self, @@ -616,7 +613,7 @@ def CreateUser( options = {} database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Create(user, path, "users", database_id, None, None, options, **kwargs) + return self.Create(user, path, "users", database_id, None, options, **kwargs) def UpsertUser( self, @@ -642,7 +639,7 @@ def UpsertUser( options = {} database_id, path = self._GetDatabaseIdWithPathForUser(database_link, user) - return self.Upsert(user, path, "users", database_id, None, None, options, **kwargs) + return self.Upsert(user, path, "users", database_id, None, options, **kwargs) def _GetDatabaseIdWithPathForUser(self, database_link: str, user: Mapping[str, Any]) -> Tuple[Optional[str], str]: base._validate_resource(user) @@ -674,7 +671,7 @@ def ReadUser( path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Read(path, "users", user_id, None, None, options, **kwargs) + return self.Read(path, "users", user_id, None, options, **kwargs) def ReadUsers( self, @@ -759,7 +756,7 @@ def DeleteDatabase( path = base.GetPathFromLink(database_link) database_id = base.GetResourceIdOrFullNameFromLink(database_link) - self.DeleteResource(path, "dbs", database_id, None, None, options, **kwargs) + self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) def CreatePermission( self, @@ -787,7 +784,7 @@ def CreatePermission( options = {} path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Create(permission, path, "permissions", user_id, None, None, options, **kwargs) + return self.Create(permission, path, "permissions", user_id, None, options, **kwargs) def UpsertPermission( self, @@ -815,7 +812,7 @@ def UpsertPermission( options = {} path, user_id = self._GetUserIdWithPathForPermission(permission, user_link) - return self.Upsert(permission, path, "permissions", user_id, None, None, options, **kwargs) + return self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) def _GetUserIdWithPathForPermission( self, @@ -940,7 +937,7 @@ def ReplaceUser( base._validate_resource(user) path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - return self.Replace(user, path, "users", user_id, None, None, options, **kwargs) + return self.Replace(user, path, "users", user_id, None, options, **kwargs) def DeleteUser( self, @@ -966,7 +963,7 @@ def DeleteUser( path = base.GetPathFromLink(user_link) user_id = base.GetResourceIdOrFullNameFromLink(user_link) - self.DeleteResource(path, "users", user_id, None, None, options, **kwargs) + self.DeleteResource(path, "users", user_id, None, options, **kwargs) def ReplacePermission( self, @@ -995,7 +992,7 @@ def ReplacePermission( base._validate_resource(permission) path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Replace(permission, path, "permissions", permission_id, None, None, options, **kwargs) + return self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) def DeletePermission( self, @@ -1021,20 +1018,21 @@ def DeletePermission( path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - self.DeleteResource(path, "permissions", permission_id, None, None, options, **kwargs) + self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) def ReadItems( self, collection_link: str, - request_context: Optional[Dict[str, Any]], feed_options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all documents in a collection. :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. + :param request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -1042,16 +1040,17 @@ def ReadItems( if feed_options is None: feed_options = {} - return self.QueryItems(collection_link, None, request_context, feed_options, response_hook=response_hook, **kwargs) + return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, + request_context=request_context, **kwargs) def QueryItems( self, database_or_container_link: str, query: Optional[Union[str, Dict[str, Any]]], - request_context: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, partition_key: Optional[PartitionKeyType] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -1063,6 +1062,7 @@ def QueryItems( :param partition_key: Partition key for the query(default value None) :type: partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :param response_hook: A callable invoked with the response metadata. + :param request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1099,9 +1099,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str query, options, response_hook=response_hook, + request_context=request_context, **kwargs) - items = ItemPaged( + return ItemPaged( self, query, options, @@ -1109,15 +1110,13 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str collection_link=database_or_container_link, page_iterator_class=query_iterable.QueryIterable ) - self._add_request_context(request_context) - return items def QueryItemsChangeFeed( self, collection_link: str, - request_context: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. @@ -1125,6 +1124,7 @@ def QueryItemsChangeFeed( :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. + :param request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1138,11 +1138,10 @@ def QueryItemsChangeFeed( if options is not None and "partitionKeyRangeId" in options: partition_key_range_id = options["partitionKeyRangeId"] - result = self._QueryChangeFeed( - collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs + return self._QueryChangeFeed( + collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, + request_context=request_context, **kwargs ) - self._add_request_context(request_context) - return result def _QueryChangeFeed( self, @@ -1151,6 +1150,7 @@ def _QueryChangeFeed( options: Optional[Mapping[str, Any]] = None, partition_key_range_id: Optional[str] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries change feed of a resource in a collection. @@ -1160,6 +1160,7 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata + :param request_context: The request context for the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1198,6 +1199,7 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str options, partition_key_range_id, response_hook=response_hook, + request_context=request_context, **kwargs) return ItemPaged( @@ -1267,8 +1269,8 @@ def CreateItem( self, database_or_container_link: str, document: Dict[str, Any], - request_context: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, + request_context: Dict[str, Any] = None, **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -1277,6 +1279,7 @@ def CreateItem( The link to the database when using partitioning, otherwise link to the document collection. :param dict document: The Azure Cosmos document to create. :param dict options: The request options for the request. + :param dict request_context: The request context for the request. :return: The created Document. :rtype: dict """ @@ -1298,8 +1301,7 @@ def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) - request_context["partitionKey"] = options["partitionKey"] - return self.Create(document, path, "docs", collection_id, None, request_context, options, **kwargs) + return self.Create(document, path, "docs", collection_id, None, options, request_context, **kwargs) def UpsertItem( self, @@ -1333,12 +1335,11 @@ def UpsertItem( # link in case of client side partitioning if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) - request_context["partitionKey"] = options["partitionKey"] collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options ) - return self.Upsert(document, path, "docs", collection_id, None, request_context, options, **kwargs) + return self.Upsert(document, path, "docs", collection_id, None, options, request_context, **kwargs) PartitionResolverErrorMessage = ( "Couldn't find any partition resolvers for the database link provided. " @@ -1391,6 +1392,7 @@ def ReadItem( :param str document_link: The link to the document. + :param dict request_context: :param dict options: The request options for the request. @@ -1405,7 +1407,7 @@ def ReadItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return self.Read(path, "docs", document_id, None, request_context, options, **kwargs) + return self.Read(path, "docs", document_id, None, options, request_context, **kwargs) def ReadTriggers( self, @@ -1492,7 +1494,7 @@ def CreateTrigger( options = {} collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Create(trigger, path, "triggers", collection_id, None, None, options, **kwargs) + return self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForTrigger( self, @@ -1534,7 +1536,7 @@ def ReadTrigger( path = base.GetPathFromLink(trigger_link) trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - return self.Read(path, "triggers", trigger_id, None, None, options, **kwargs) + return self.Read(path, "triggers", trigger_id, None, options, **kwargs) def UpsertTrigger( self, @@ -1558,7 +1560,7 @@ def UpsertTrigger( options = {} collection_id, path, trigger = self._GetContainerIdWithPathForTrigger(collection_link, trigger) - return self.Upsert(trigger, path, "triggers", collection_id, None, None, options, **kwargs) + return self.Upsert(trigger, path, "triggers", collection_id, None, options, **kwargs) def ReadUserDefinedFunctions( self, @@ -1645,7 +1647,7 @@ def CreateUserDefinedFunction( options = {} collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Create(udf, path, "udfs", collection_id, None, None, options, **kwargs) + return self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) def UpsertUserDefinedFunction( self, @@ -1672,7 +1674,7 @@ def UpsertUserDefinedFunction( options = {} collection_id, path, udf = self._GetContainerIdWithPathForUDF(collection_link, udf) - return self.Upsert(udf, path, "udfs", collection_id, None, None, options, **kwargs) + return self.Upsert(udf, path, "udfs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForUDF( self, @@ -1714,7 +1716,7 @@ def ReadUserDefinedFunction( path = base.GetPathFromLink(udf_link) udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - return self.Read(path, "udfs", udf_id, None, None, options, **kwargs) + return self.Read(path, "udfs", udf_id, None, options, **kwargs) def ReadStoredProcedures( self, @@ -1801,7 +1803,7 @@ def CreateStoredProcedure( options = {} collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Create(sproc, path, "sprocs", collection_id, None, None, options, **kwargs) + return self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) def UpsertStoredProcedure( self, @@ -1828,7 +1830,7 @@ def UpsertStoredProcedure( options = {} collection_id, path, sproc = self._GetContainerIdWithPathForSproc(collection_link, sproc) - return self.Upsert(sproc, path, "sprocs", collection_id, None, None, options, **kwargs) + return self.Upsert(sproc, path, "sprocs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForSproc( self, @@ -1869,7 +1871,7 @@ def ReadStoredProcedure( path = base.GetPathFromLink(sproc_link) sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - return self.Read(path, "sprocs", sproc_id, None, None, options, **kwargs) + return self.Read(path, "sprocs", sproc_id, None, options, **kwargs) def ReadConflicts( self, @@ -1953,7 +1955,7 @@ def ReadConflict( path = base.GetPathFromLink(conflict_link) conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - return self.Read(path, "conflicts", conflict_id, None, None, options, **kwargs) + return self.Read(path, "conflicts", conflict_id, None, options, **kwargs) def DeleteContainer( self, @@ -1979,7 +1981,7 @@ def DeleteContainer( path = base.GetPathFromLink(collection_link) collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) - self.DeleteResource(path, "colls", collection_id, None, None, options, **kwargs) + self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) def ReplaceItem( self, @@ -2021,9 +2023,8 @@ def ReplaceItem( # Extract the document collection link and add the partition key to options collection_link = base.GetItemContainerLink(document_link) options = self._AddPartitionKey(collection_link, new_document, options) - request_context["partitionKey"] = options["partitionKey"] - return self.Replace(new_document, path, "docs", document_id, None, request_context, options, **kwargs) + return self.Replace(new_document, path, "docs", document_id, None, options, request_context, **kwargs) def PatchItem( self, @@ -2180,7 +2181,7 @@ def DeleteItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - self.DeleteResource(path, "docs", document_id, None, request_context, options, **kwargs) + self.DeleteResource(path, "docs", document_id, None, options, request_context, **kwargs) def DeleteAllItemsByPartitionKey( self, @@ -2280,7 +2281,7 @@ def DeleteTrigger( path = base.GetPathFromLink(trigger_link) trigger_id = base.GetResourceIdOrFullNameFromLink(trigger_link) - self.DeleteResource(path, "triggers", trigger_id, None, None, options, **kwargs) + self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) def ReplaceUserDefinedFunction( self, @@ -2341,7 +2342,7 @@ def DeleteUserDefinedFunction( path = base.GetPathFromLink(udf_link) udf_id = base.GetResourceIdOrFullNameFromLink(udf_link) - self.DeleteResource(path, "udfs", udf_id, None, None, options, **kwargs) + self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) def ExecuteStoredProcedure( self, @@ -2442,7 +2443,7 @@ def DeleteStoredProcedure( path = base.GetPathFromLink(sproc_link) sproc_id = base.GetResourceIdOrFullNameFromLink(sproc_link) - self.DeleteResource(path, "sprocs", sproc_id, None, None, options, **kwargs) + self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) def DeleteConflict( self, @@ -2468,7 +2469,7 @@ def DeleteConflict( path = base.GetPathFromLink(conflict_link) conflict_id = base.GetResourceIdOrFullNameFromLink(conflict_link) - self.DeleteResource(path, "conflicts", conflict_id, None, None, options, **kwargs) + self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) def ReplaceOffer( self, @@ -2491,7 +2492,7 @@ def ReplaceOffer( base._validate_resource(offer) path = base.GetPathFromLink(offer_link) offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) + return self.Replace(offer, path, "offers", offer_id, None, **kwargs) def ReadOffer( self, @@ -2508,7 +2509,7 @@ def ReadOffer( """ path = base.GetPathFromLink(offer_link) offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Read(path, "offers", offer_id, None, None, {}, **kwargs) + return self.Read(path, "offers", offer_id, None,{}, **kwargs) def ReadOffers( self, @@ -2614,8 +2615,8 @@ def Create( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], - request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -2662,8 +2663,8 @@ def Upsert( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], - request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -2710,8 +2711,8 @@ def Replace( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], - request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -2756,8 +2757,8 @@ def Read( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], - request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -2798,8 +2799,8 @@ def DeleteResource( typ: str, id: Optional[str], initial_headers: Optional[Mapping[str, Any]], - request_context: Optional[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -3019,6 +3020,7 @@ def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements partition_key_range_id: Optional[str] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, is_query_plan: bool = False, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -3079,6 +3081,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return __GetBodiesFromQueryResult(result), last_response_headers @@ -3161,6 +3165,9 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: results["Documents"].extend(partial_result["Documents"]) else: results = partial_result + + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, partial_result) # if the prefix partition query has results lets return it @@ -3173,7 +3180,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: index_metrics_raw = last_response_headers[INDEX_METRICS_HEADER] last_response_headers[INDEX_METRICS_HEADER] = _utils.get_index_metrics_info(index_metrics_raw) self.last_response_headers = last_response_headers - + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) @@ -3378,11 +3386,4 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s def _add_request_context(self, request_context): if http_constants.HttpHeaders.SessionToken in self.last_response_headers: request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] - if 'partitionKey' in request_context: - request_context["feed_range"] = FeedRangeEpk(FeedRangeInternalEpk.get_epk_range_for_partition_key( - request_context['container_properties'], - request_context['partitionKey'])) - request_context.pop('partitionKey') - if 'container_properties' in request_context: - request_context.pop('container_properties') self.last_response_headers["request_context"] = request_context \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index dd32019607bc..e3331088fde9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -38,7 +38,7 @@ def merge_session_tokens_with_same_range(session_token1, session_token2): def is_compound_session_token(session_token): return "," in session_token -def create_vector_session_token_and_pkrangeid(session_token): +def create_vector_session_token_and_pkrangeid(session_token): # cspell:disable-line tokens = session_token.split(":") return tokens[0], VectorSessionToken.create(tokens[1]) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 3f879ded3187..dcfb228c1bc7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -43,6 +43,7 @@ ) from .._feed_range import FeedRange, FeedRangeEpk from .._routing.routing_range import Range +from .._session_token_helpers import get_updated_session_token from ..offer import ThroughputProperties from ..partition_key import ( NonePartitionKeyValue, @@ -259,7 +260,8 @@ async def create_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = await self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, options=request_options, **kwargs + database_or_container_link=self.container_link, document=body, options=request_options, + request_context={}, **kwargs ) return result or {} @@ -325,7 +327,8 @@ async def read_item( if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - return await self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) + return await self.client_connection.ReadItem(document_link=doc_link, options=request_options, + request_context={}, **kwargs) @distributed_trace def read_all_items( @@ -373,7 +376,8 @@ def read_all_items( feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] items = self.client_connection.ReadItems( - collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs + collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, + request_context={}, **kwargs ) if response_hook: response_hook(self.client_connection.last_response_headers, items) @@ -492,6 +496,7 @@ def query_items( options=feed_options, partition_key=partition_key, response_hook=response_hook, + request_context= {}, **kwargs ) if response_hook: @@ -690,9 +695,11 @@ def query_items_change_feed( # pylint: disable=unused-argument change_feed_state_context["partitionKeyFeedRange"] = \ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) + request_context = {} if kwargs.get("feed_range") is not None: feed_range: FeedRangeEpk = kwargs.pop('feed_range') change_feed_state_context["feedRange"] = feed_range._feed_range_internal + request_context["feedRange"] = feed_range feed_options["containerProperties"] = self._get_properties() feed_options["changeFeedStateContext"] = change_feed_state_context @@ -705,7 +712,8 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.QueryItemsChangeFeed( - self.container_link, options=feed_options, response_hook=response_hook, **kwargs + self.container_link, options=feed_options, response_hook=response_hook, + request_context=request_context, **kwargs ) if response_hook: @@ -779,6 +787,7 @@ async def upsert_item( database_or_container_link=self.container_link, document=body, options=request_options, + request_context={}, **kwargs ) return result or {} @@ -851,7 +860,8 @@ async def replace_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = await self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, options=request_options, **kwargs + document_link=item_link, new_document=body, options=request_options, + request_context={}, **kwargs ) return result or {} @@ -927,7 +937,8 @@ async def patch_item( item_link = self._get_document_link(item) result = await self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, options=request_options, **kwargs) + document_link=item_link, operations=patch_operations, options=request_options, + request_context={}, **kwargs) return result or {} @distributed_trace_async @@ -990,7 +1001,8 @@ async def delete_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) - await self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) + await self.client_connection.DeleteItem(document_link=document_link, options=request_options, + request_context={}, **kwargs) @distributed_trace_async async def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @@ -1291,7 +1303,8 @@ async def execute_item_batch( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] return await self.client_connection.Batch( - collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) + collection_link=self.container_link, batch_operations=batch_operations, options=request_options, + request_context={}, **kwargs) async def read_feed_ranges( self, @@ -1319,3 +1332,40 @@ async def read_feed_ranges( return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) for partitionKeyRange in partition_key_ranges] + + async def get_updated_session_token(self, + feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], + target_feed_range: FeedRange + ) -> "Session Token": + """Gets the the most up to date session token from the list of session token and feed range tuples + for a specific target feed range. The feed range can be obtained from a response from any crud operation. + This should only be used if maintaining own session token or else the sdk will keep track of + session token. + :param feed_ranges_to_session_tokens: List of partition key and session token tuples. + :type feed_ranges_to_session_tokens: List[Tuple(str, FeedRange)] + :param target_feed_range: feed range to get most up to date session token. + :type target_feed_range: FeedRange + :returns: a session token + :rtype: str + """ + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + + async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: + """Gets the feed range for a given partition key. + :param partition_key: partition key to get feed range. + :type partition_key: PartitionKey + :returns: a feed range + :rtype: FeedRange + """ + return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key)) + + async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + """Checks if child feed range is a subset of parent feed range. + :param parent_feed_range: left feed range + :type parent_feed_range: FeedRange + :param child_feed_range: right feed range + :type child_feed_range: FeedRange + :returns: a boolean indicating if child feed range is a subset of parent feed range + :rtype: bool + """ + return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_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 eeb67225660a..2a9741eff831 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 @@ -24,6 +24,7 @@ """Document client class for the Azure Cosmos database service. """ import os +from optparse import Option from urllib.parse import urlparse from typing import ( Callable, Dict, Any, Iterable, Mapping, Optional, List, @@ -515,6 +516,7 @@ async def CreateItem( database_or_container_link: str, document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -550,7 +552,8 @@ async def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = await self._AddPartitionKey(database_or_container_link, document, options) - return await self.Create(document, path, "docs", collection_id, None, options, **kwargs) + return await self.Create(document, path, "docs", collection_id, None, options, + request_context, **kwargs) async def CreatePermission( self, @@ -704,6 +707,7 @@ async def Create( id: Optional[str], initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -715,6 +719,7 @@ async def Create( :param dict initial_headers: :param dict options: The request options for the request. + :param dict request_context: :return: The created Azure Cosmos resource. :rtype: @@ -735,6 +740,8 @@ async def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -796,6 +803,7 @@ async def UpsertItem( database_or_container_link: str, document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Upserts a document in a collection. @@ -830,7 +838,8 @@ async def UpsertItem( collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options ) - return await self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) + return await self.Upsert(document, path, "docs", collection_id, None, options, + request_context, **kwargs) async def Upsert( self, @@ -840,6 +849,7 @@ async def Upsert( id: Optional[str], initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -851,6 +861,7 @@ async def Upsert( :param dict initial_headers: :param dict options: The request options for the request. + :param dict request_context: :return: The upserted Azure Cosmos resource. :rtype: @@ -872,6 +883,8 @@ async def Upsert( self.last_response_headers = last_response_headers # update session for write request self._UpdateSessionIfRequired(headers, result, self.last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -959,6 +972,7 @@ async def ReadItem( self, document_link: str, options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Reads a document. @@ -967,6 +981,7 @@ async def ReadItem( The link to the document. :param dict options: The request options for the request. + :param dict request_context: :return: The read Document. @@ -979,7 +994,7 @@ async def ReadItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return await self.Read(path, "docs", document_id, None, options, **kwargs) + return await self.Read(path, "docs", document_id, None, options, request_context, **kwargs) async def ReadUser( self, @@ -1142,6 +1157,7 @@ async def Read( id: Optional[str], initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -1169,6 +1185,8 @@ async def Read( request_params = _request_object.RequestObject(typ, documents._OperationType.Read) result, last_response_headers = await self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -1358,6 +1376,7 @@ async def ReplaceItem( document_link: str, new_document: Dict[str, Any], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Replaces a document and returns it. @@ -1367,6 +1386,7 @@ async def ReplaceItem( :param dict new_document: :param dict options: The request options for the request. + :param request_context: :return: The new Document. :rtype: @@ -1391,13 +1411,15 @@ async def ReplaceItem( collection_link = base.GetItemContainerLink(document_link) options = await self._AddPartitionKey(collection_link, new_document, options) - return await self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) + return await self.Replace(new_document, path, "docs", document_id, None, options, + request_context, **kwargs) async def PatchItem( self, document_link: str, operations: List[Dict[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Patches a document and returns it. @@ -1405,6 +1427,7 @@ async def PatchItem( :param str document_link: The link to the document. :param list operations: The operations for the patch request. :param dict options: The request options for the request. + :param dict request_context: The request context for the operation :return: The new Document. :rtype: @@ -1432,6 +1455,8 @@ async def PatchItem( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -1500,6 +1525,7 @@ async def Replace( id: Optional[str], initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -1511,6 +1537,7 @@ async def Replace( :param dict initial_headers: :param dict options: The request options for the request. + :param dict request_context: :return: The new Azure Cosmos resource. :rtype: @@ -1530,6 +1557,8 @@ async def Replace( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -1695,6 +1724,7 @@ async def DeleteItem( self, document_link: str, options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> None: """Deletes a document. @@ -1703,6 +1733,7 @@ async def DeleteItem( The link to the document. :param dict options: The request options for the request. + :param dict request_context: :return: The deleted Document. :rtype: @@ -1714,7 +1745,8 @@ async def DeleteItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - await self.DeleteResource(path, "docs", document_id, None, options, **kwargs) + await self.DeleteResource(path, "docs", document_id, None, options, + request_context, **kwargs) async def DeleteUserDefinedFunction( self, @@ -1822,6 +1854,7 @@ async def DeleteResource( id: Optional[str], initial_headers: Optional[Mapping[str, Any]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -1832,6 +1865,7 @@ async def DeleteResource( :param dict initial_headers: :param dict options: The request options for the request. + :param dict request_context: :return: The deleted Azure Cosmos resource. :rtype: @@ -1851,6 +1885,8 @@ async def DeleteResource( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, None) @@ -1886,6 +1922,7 @@ async def Batch( collection_link: str, batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], options: Optional[Mapping[str, Any]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> List[Dict[str, Any]]: """Executes the given operations in transactional batch. @@ -1941,6 +1978,8 @@ async def Batch( ), operation_responses=final_responses ) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, final_responses) return final_responses @@ -2138,6 +2177,7 @@ def ReadItems( collection_link: str, feed_options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all documents in a collection. @@ -2145,6 +2185,7 @@ def ReadItems( :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. + :param request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -2153,7 +2194,8 @@ def ReadItems( if feed_options is None: feed_options = {} - return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) + return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, + request_context=request_context, **kwargs) def QueryItems( self, @@ -2162,6 +2204,7 @@ def QueryItems( options: Optional[Mapping[str, Any]] = None, partition_key: Optional[PartitionKeyType] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -2172,6 +2215,7 @@ def QueryItems( :param dict options: The request options for the request. :param str partition_key: Partition key for the query(default value None) :param response_hook: A callable invoked with the response metadata. + :param request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2208,6 +2252,7 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di query, options, response_hook=response_hook, + request_context=request_context, **kwargs ), self.last_response_headers, @@ -2227,6 +2272,7 @@ def QueryItemsChangeFeed( collection_link: str, options: Optional[Mapping[str, Any]] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. @@ -2234,6 +2280,7 @@ def QueryItemsChangeFeed( :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. + :param request_context: The request_context of the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2247,7 +2294,8 @@ def QueryItemsChangeFeed( partition_key_range_id = options["partitionKeyRangeId"] return self._QueryChangeFeed( - collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs + collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, + request_context=request_context, **kwargs ) def _QueryChangeFeed( @@ -2257,6 +2305,7 @@ def _QueryChangeFeed( options: Optional[Mapping[str, Any]] = None, partition_key_range_id: Optional[str] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries change feed of a resource in a collection. @@ -2266,6 +2315,7 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata + :param request_context: The request context of the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2304,6 +2354,7 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di options, partition_key_range_id, response_hook=response_hook, + request_context=request_context, **kwargs ), self.last_response_headers, @@ -2768,6 +2819,7 @@ async def __QueryFeed( # pylint: disable=too-many-branches,too-many-statements, partition_key_range_id: Optional[str] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, is_query_plan: bool = False, + request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> List[Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -2818,6 +2870,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: await change_feed_state.populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, result) return __GetBodiesFromQueryResult(result) @@ -2902,6 +2956,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: results["Documents"].extend(partial_result["Documents"]) else: results = partial_result + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, partial_result) # if the prefix partition query has results lets return it @@ -2913,6 +2969,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: INDEX_METRICS_HEADER = http_constants.HttpHeaders.IndexUtilization index_metrics_raw = self.last_response_headers[INDEX_METRICS_HEADER] self.last_response_headers[INDEX_METRICS_HEADER] = _utils.get_index_metrics_info(index_metrics_raw) + if request_context is not None: + self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, result) @@ -3222,3 +3280,8 @@ async def _get_partition_key_definition(self, collection_link: str) -> Optional[ partition_key_definition = container.get("partitionKey") self.__container_properties_cache[collection_link] = _set_properties_cache(container) return partition_key_definition + + def _add_request_context(self, request_context): + if http_constants.HttpHeaders.SessionToken in self.last_response_headers: + request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] + self.last_response_headers["request_context"] = request_context \ 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 ae34e34e0a7f..1ec70a244956 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -133,7 +133,12 @@ def _set_partition_key( 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) + def _get_epk_range_for_partition_key( self, partition_key_value: PartitionKeyType) -> Range: + container_properties = self._get_properties() + partition_key_definition = container_properties["partitionKey"] + partition_key = PartitionKey(path=partition_key_definition["paths"], kind=partition_key_definition["kind"]) + return partition_key._get_epk_range_for_partition_key(partition_key_value) def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: return self.client_connection._container_properties_cache @@ -254,10 +259,8 @@ def read_item( # pylint:disable=docstring-missing-param request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {"partitionKey": request_options["partitionKey"], - "container_properties": self._get_properties()} - return self.client_connection.ReadItem(document_link=doc_link, options=request_options, - request_context=request_context, **kwargs) + return self.client_connection.ReadItem(document_link=doc_link, request_context={}, + options=request_options, **kwargs) @distributed_trace def read_all_items( # pylint:disable=docstring-missing-param @@ -312,10 +315,9 @@ def read_all_items( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {"container_properties": self._get_properties()} items = self.client_connection.ReadItems( - collection_link=self.container_link, request_context=request_context, feed_options=feed_options, - response_hook=response_hook, **kwargs) + collection_link=self.container_link, feed_options=feed_options, + response_hook=response_hook, request_context={}, **kwargs) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -526,8 +528,7 @@ def query_items_change_feed( change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) change_feed_state_context["partitionKeyFeedRange"] =\ - FeedRangeInternalEpk.get_epk_range_for_partition_key(self._get_properties(), kwargs.pop('partition_key')) - request_context["feed_range"] = FeedRangeEpk(change_feed_state_context["partitionKeyFeedRange"]) + self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: feed_range: FeedRangeEpk = kwargs.pop('feed_range') @@ -543,8 +544,11 @@ def query_items_change_feed( response_hook.clear() result = self.client_connection.QueryItemsChangeFeed( - self.container_link, request_context=request_context, - options=feed_options, response_hook=response_hook, **kwargs + self.container_link, + options=feed_options, + response_hook=response_hook, + request_context=request_context, + **kwargs ) if response_hook: response_hook(self.client_connection.last_response_headers, result) @@ -640,7 +644,6 @@ def query_items( # pylint:disable=docstring-missing-param feed_options["populateQueryMetrics"] = populate_query_metrics if populate_index_metrics is not None: feed_options["populateIndexMetrics"] = populate_index_metrics - request_context = {} if partition_key is not None: partition_key_value = self._set_partition_key(partition_key) if self.__is_prefix_partitionkey(partition_key): @@ -650,7 +653,6 @@ def query_items( # pylint:disable=docstring-missing-param kwargs["partitionKeyDefinition"]["partition_key"] = partition_key_value else: feed_options["partitionKey"] = partition_key_value - request_context["partitionKey"] = feed_options["partitionKey"] if enable_scan_in_query is not None: feed_options["enableScanInQuery"] = enable_scan_in_query if max_integrated_cache_staleness_in_ms: @@ -664,14 +666,13 @@ def query_items( # pylint:disable=docstring-missing-param 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"] - request_context["container_properties"] = self._get_properties() items = self.client_connection.QueryItems( database_or_container_link=self.container_link, query=query if parameters is None else {"query": query, "parameters": parameters}, - request_context=request_context, options=feed_options, partition_key=partition_key, response_hook=response_hook, + request_context={}, **kwargs ) if response_hook: @@ -758,9 +759,8 @@ def replace_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {'container_properties': self._get_properties()} result = self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, request_context=request_context, options=request_options, **kwargs + document_link=item_link, new_document=body, request_context={}, options=request_options, **kwargs ) return result or {} @@ -831,11 +831,10 @@ def upsert_item( # pylint:disable=docstring-missing-param request_options["populateQueryMetrics"] = populate_query_metrics if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {'container_properties': self._get_properties()} result = self.client_connection.UpsertItem( database_or_container_link=self.container_link, document=body, - request_context=request_context, + request_context={}, options=request_options, **kwargs ) @@ -916,9 +915,9 @@ def create_item( # pylint:disable=docstring-missing-param request_options["indexingDirective"] = indexing_directive if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - request_context = {'container_properties': self._get_properties()} result = self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, request_context=request_context, options=request_options, **kwargs) + database_or_container_link=self.container_link, document=body, + options=request_options, request_context={}, **kwargs) return result or {} @distributed_trace @@ -992,10 +991,9 @@ def patch_item( if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] item_link = self._get_document_link(item) - request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} result = self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, request_context=request_context, - options=request_options, **kwargs) + document_link=item_link, operations=patch_operations, + request_context={}, options=request_options, **kwargs) return result or {} @distributed_trace @@ -1047,10 +1045,10 @@ def execute_item_batch( kwargs['priority'] = priority request_options = build_options(kwargs) request_options["partitionKey"] = self._set_partition_key(partition_key) - request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} request_options["disableAutomaticIdGeneration"] = True result = self.client_connection.Batch( - collection_link=self.container_link, batch_operations=batch_operations, request_context=request_context, + collection_link=self.container_link, batch_operations=batch_operations, + request_context={}, options=request_options, **kwargs) return result @@ -1119,8 +1117,8 @@ def delete_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) - request_context = {"partitionKey": request_options["partitionKey"], 'container_properties': self._get_properties()} - self.client_connection.DeleteItem(document_link=document_link, request_context=request_context, + self.client_connection.DeleteItem(document_link=document_link, + request_context={}, options=request_options, **kwargs) @distributed_trace @@ -1408,12 +1406,11 @@ def read_feed_ranges( for partitionKeyRange in partition_key_ranges] def get_updated_session_token(self, - feed_ranges_to_session_tokens: List, + feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], target_feed_range: FeedRange ) -> "Session Token": - """Gets the best session token from the list of session token and feed range tuples - to figure out which is the most up to date for a specific - feed range. The feed range can be obtained from a response from any crud operation. + """Gets the the most up to date session token from the list of session token and feed range tuples + for a specific target feed range. The feed range can be obtained from a response from any crud operation. This should only be used if maintaining own session token or else the sdk will keep track of session token. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. @@ -1430,16 +1427,16 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Feed :param partition_key: partition key to get feed range. :type partition_key: PartitionKey :returns: a feed range - :rtype: Range + :rtype: FeedRange """ - return FeedRangeEpk(FeedRangeInternalEpk.get_epk_range_for_partition_key(self._get_properties(), partition_key)) + return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. :param parent_feed_range: left feed range - :type parent_feed_range: Range + :type parent_feed_range: FeedRange :param child_feed_range: right feed range - :type child_feed_range: Range + :type child_feed_range: FeedRange :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index f62342671dd8..5aced6f7d2c1 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -1,6 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import time import unittest import uuid @@ -24,16 +24,16 @@ def setup(): "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) } -def validate_request_context(collection, with_partition_key=True): +def validate_request_context(collection): request_context = collection.client_connection.last_response_headers["request_context"] keys_expected = ["session_token"] - if with_partition_key: - keys_expected.append("feed_range") assert request_context is not None for key in keys_expected: assert request_context[key] is not None -def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'): +def createItem(id = None, pk='A', name='sample'): + if id is None: + id = 'item' + str(uuid.uuid4()) item = { 'id': id, 'name': name, @@ -53,7 +53,6 @@ class TestRequestContext: TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID - # test out all operations def test_crud_request_context(self, setup): item = createItem() setup["created_collection"].create_item(item) @@ -70,27 +69,26 @@ def test_crud_request_context(self, setup): {"op": "replace", "path": "/name", "value": 14}, ] setup["created_collection"].patch_item(item['id'], item['pk'], operations) - validate_request_context(setup["created_collection"], False) + validate_request_context(setup["created_collection"]) - setup["created_collection"].read_all_items() - validate_request_context(setup["created_collection"], False) + items = list(setup["created_collection"].read_all_items()) + assert len(items) == 1 + validate_request_context(setup["created_collection"]) setup["created_collection"].upsert_item(createItem()) validate_request_context(setup["created_collection"]) - setup["created_collection"].query_items_change_feed() - validate_request_context(setup["created_collection"], False) - - # with partition key - setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])], partition_key=item['pk']) + for i in range(100): + setup["created_collection"].create_item(createItem()) + items = list(setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)) + assert len(items) == 102 validate_request_context(setup["created_collection"]) - # without partition key - setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", parameters=[dict(name="@id", value=item['id'])]) - validate_request_context(setup["created_collection"], False) - - setup["created_collection"].read_all_items() - validate_request_context(setup["created_collection"], False) + items = list(setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", + parameters=[dict(name="@id", value=item['id'])], + partition_key=item['pk'])) + assert len(items) == 1 + validate_request_context(setup["created_collection"]) setup["created_collection"].delete_item(item['id'], item['pk']) validate_request_context(setup["created_collection"]) diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py new file mode 100644 index 000000000000..268f115a099a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py @@ -0,0 +1,99 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest +import pytest_asyncio + +import test_config +from azure.cosmos.aio import CosmosClient + +@pytest_asyncio.fixture() +async def setup(): + if (TestRequestContext.masterKey == '[YOUR_KEY_HERE]' or + TestRequestContext.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = CosmosClient(TestRequestContext.host, test_config.TestConfig.masterKey), + created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) + return { + "created_db": created_db, + "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) + } + +def validate_request_context(collection): + request_context = collection.client_connection.last_response_headers["request_context"] + keys_expected = ["session_token"] + assert request_context is not None + for key in keys_expected: + assert request_context[key] is not None + +def createItem(id = None, pk='A', name='sample'): + if id is None: + id = 'item' + str(uuid.uuid4()) + item = { + 'id': id, + 'name': name, + 'pk': pk + } + return item + + +@pytest.mark.cosmosEmulator +@pytest.mark.asyncio +@pytest.mark.usefixtures("setup") +class TestRequestContext: + """Tests to verify request context gets populated correctly + """ + + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID + TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + + async def test_crud_request_context(self, setup): + item = createItem() + await setup["created_collection"].create_item(item) + validate_request_context(setup["created_collection"]) + + await setup["created_collection"].read_item(item['id'], item['pk']) + validate_request_context(setup["created_collection"]) + + new_item = createItem(item['id'], name='sample_replaced') + await setup["created_collection"].replace_item(item['id'], new_item) + validate_request_context(setup["created_collection"]) + operations = [ + {"op": "add", "path": "/favorite_color", "value": "red"}, + {"op": "replace", "path": "/name", "value": 14}, + ] + await setup["created_collection"].patch_item(item['id'], item['pk'], operations) + validate_request_context(setup["created_collection"]) + + items = [item async for item in setup["created_collection"].read_all_items()] + assert len(items) == 1 + validate_request_context(setup["created_collection"]) + + await setup["created_collection"].upsert_item(createItem()) + validate_request_context(setup["created_collection"]) + + for i in range(100): + await setup["created_collection"].create_item(createItem()) + items = [item async for item in setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)] + assert len(items) == 102 + validate_request_context(setup["created_collection"]) + + items = [item async for item in setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", + parameters=[dict(name="@id", value=item['id'])], + partition_key=item['pk'])] + assert len(items) == 1 + validate_request_context(setup["created_collection"]) + + await setup["created_collection"].delete_item(item['id'], item['pk']) + validate_request_context(setup["created_collection"]) + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index a043665f21ab..bf6e84e8bb8d 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -110,7 +110,8 @@ def create_items_logical_pk(setup, target_pk, previous_session_token, feed_range if item['pk'] == target_pk: target_session_token = request_context["session_token"] previous_session_token = request_context["session_token"] - feed_ranges_and_session_tokens.append((request_context["feed_range"], request_context["session_token"])) + feed_ranges_and_session_tokens.append((setup[COLLECTION].feed_range_from_partition_key(item['pk']), + request_context["session_token"])) return target_session_token def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): @@ -122,7 +123,7 @@ def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ target_feed_range = feed_range break - for i in range(100): + for i in range(10): item = { 'id': 'item' + str(uuid.uuid4()), 'name': 'sample', @@ -130,10 +131,12 @@ def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ } setup[COLLECTION].create_item(item, session_token=previous_session_token) request_context = setup[COLLECTION].client_connection.last_response_headers["request_context"] - if setup[COLLECTION].is_feed_range_subset(target_feed_range, request_context["feed_range"]): + curr_feed_range = setup[COLLECTION].feed_range_from_partition_key(item['pk']) + if setup[COLLECTION].is_feed_range_subset(target_feed_range, curr_feed_range): target_session_token = request_context["session_token"] previous_session_token = request_context["session_token"] - feed_ranges_and_session_tokens.append((request_context["feed_range"], request_context["session_token"])) + feed_ranges_and_session_tokens.append((curr_feed_range, + request_context["session_token"])) return target_session_token, target_feed_range @@ -150,7 +153,7 @@ class TestSessionTokenHelpers: connectionPolicy = test_config.TestConfig.connectionPolicy configs = test_config.TestConfig TEST_DATABASE_ID = configs.TEST_DATABASE_ID - TEST_COLLECTION_ID = configs.TEST_MULTI_PARTITION_CONTAINER_ID + TEST_COLLECTION_ID = configs.TEST_SINGLE_PARTITION_CONTAINER_ID def test_get_session_token_update(self, setup): feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) From c252d8840130c6c2e4b8faf33d61532bb9e52942 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 8 Oct 2024 18:07:25 -0700 Subject: [PATCH 34/59] fix tests --- .../azure/cosmos/_session_token_helpers.py | 19 ++++++----- .../aio/_cosmos_client_connection_async.py | 1 - .../azure-cosmos/test/test_feed_range.py | 25 +++++++------- .../azure-cosmos/test/test_request_context.py | 4 +-- .../test/test_request_context_async.py | 4 +-- .../test/test_session_token_helpers.py | 34 ++++++++++++------- 6 files changed, 49 insertions(+), 38 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index e3331088fde9..c7bbaad98606 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -26,8 +26,8 @@ def merge_session_tokens_with_same_range(session_token1, session_token2): - pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrangeid(session_token1) - pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrangeid(session_token2) + pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrange_id(session_token1) + pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrange_id(session_token2) pk_range_id = pk_range_id1 if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ @@ -38,7 +38,7 @@ def merge_session_tokens_with_same_range(session_token1, session_token2): def is_compound_session_token(session_token): return "," in session_token -def create_vector_session_token_and_pkrangeid(session_token): # cspell:disable-line +def create_vector_session_token_and_pkrange_id(session_token): tokens = session_token.split(":") return tokens[0], VectorSessionToken.create(tokens[1]) @@ -58,8 +58,8 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): while i < len(session_tokens): j = i + 1 while j < len(session_tokens): - pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrangeid(session_tokens[i]) - pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrangeid(session_tokens[j]) + pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrange_id(session_tokens[i]) + pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrange_id(session_tokens[j]) if pk_range_id1 == pk_range_id2: vector_session_token = vector_session_token1.merge(vector_session_token2) session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) @@ -67,6 +67,7 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): for token in remove_session_tokens: session_tokens.remove(token) i = -1 + break j += 1 i += 1 @@ -98,6 +99,7 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): for feed_range_to_remove in feed_ranges_to_remove: overlapping_ranges.remove(feed_range_to_remove) overlapping_ranges.append((cur_feed_range, session_token)) + i, j = 0, 1 else: j += 1 if j == len(overlapping_ranges): @@ -114,8 +116,7 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): done_overlapping_ranges.append(overlapping_ranges[0]) overlapping_ranges.remove(overlapping_ranges[0]) continue - tokens_cmp = session_token_cmp.split(":") - vector_session_token_cmp = VectorSessionToken.create(tokens_cmp[1]) + _, vector_session_token_cmp = create_vector_session_token_and_pkrange_id(session_token_cmp) subsets = [] # finding the subset feed ranges of the current feed range for j in range(1, len(overlapping_ranges)): @@ -132,7 +133,7 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): session_tokens = [subsets[j][1]] merged_indices = [subsets[j][2]] if len(subsets) == 1: - _, vector_session_token = create_vector_session_token_and_pkrangeid(session_tokens[0]) + _, vector_session_token = create_vector_session_token_and_pkrange_id(session_tokens[0]) if vector_session_token_cmp.is_greater(vector_session_token): overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) else: @@ -149,7 +150,7 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): # else take the current feed range children_more_updated = True for session_token in session_tokens: - _, vector_session_token = create_vector_session_token_and_pkrangeid(session_token) + _, vector_session_token = create_vector_session_token_and_pkrange_id(session_token) if vector_session_token_cmp.is_greater(vector_session_token): children_more_updated = False feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] 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 2a9741eff831..547b16407d73 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 @@ -24,7 +24,6 @@ """Document client class for the Azure Cosmos database service. """ import os -from optparse import Option from urllib.parse import urlparse from typing import ( Callable, Dict, Any, Iterable, Mapping, Optional, List, diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 1efc9f13c3c3..d8d44b6ceb6e 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -31,23 +31,24 @@ def setup(): Range("3F", "7F", True, False), True), (Range("3F", "7F", True, False), - Range("", "FF", True, False), - False), + Range("", "FF", True, False), + False), (Range("3F", "7F", True, False), - Range("", "5F", True, False), - False), + Range("", "5F", True, False), + False), (Range("3F", "7F", True, True), - Range("3F", "7F", True, True), - True), + Range("3F", "7F", True, True), + True), (Range("3F", "7F", False, True), - Range("3F", "7F", True, True), - False), + Range("3F", "7F", True, True), + False), (Range("3F", "7F", True, False), - Range("3F", "7F", True, True), - False), + Range("3F", "7F", True, True), + False), (Range("3F", "7F", True, False), - Range("", "2F", True, False), - False)] + Range("", "2F", True, False), + False) + ] test_overlaps_ranges = [(Range("", "FF", True, False), diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py index 5aced6f7d2c1..1c46152dc0c3 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context.py @@ -72,7 +72,7 @@ def test_crud_request_context(self, setup): validate_request_context(setup["created_collection"]) items = list(setup["created_collection"].read_all_items()) - assert len(items) == 1 + assert len(items) > 0 validate_request_context(setup["created_collection"]) setup["created_collection"].upsert_item(createItem()) @@ -81,7 +81,7 @@ def test_crud_request_context(self, setup): for i in range(100): setup["created_collection"].create_item(createItem()) items = list(setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)) - assert len(items) == 102 + assert len(items) > 100 validate_request_context(setup["created_collection"]) items = list(setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py index 268f115a099a..2de921aa3494 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py @@ -74,7 +74,7 @@ async def test_crud_request_context(self, setup): validate_request_context(setup["created_collection"]) items = [item async for item in setup["created_collection"].read_all_items()] - assert len(items) == 1 + assert len(items) > 0 validate_request_context(setup["created_collection"]) await setup["created_collection"].upsert_item(createItem()) @@ -83,7 +83,7 @@ async def test_crud_request_context(self, setup): for i in range(100): await setup["created_collection"].create_item(createItem()) items = [item async for item in setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)] - assert len(items) == 102 + assert len(items) > 1000 validate_request_context(setup["created_collection"]) items = [item async for item in setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index bf6e84e8bb8d..0d354be7a797 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -13,6 +13,7 @@ from azure.cosmos import DatabaseProxy from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range +from azure.cosmos._session_token_helpers import is_compound_session_token, create_vector_session_token_and_pkrange_id COLLECTION = "created_collection" DATABASE = "created_db" @@ -112,7 +113,7 @@ def create_items_logical_pk(setup, target_pk, previous_session_token, feed_range previous_session_token = request_context["session_token"] feed_ranges_and_session_tokens.append((setup[COLLECTION].feed_range_from_partition_key(item['pk']), request_context["session_token"])) - return target_session_token + return target_session_token, previous_session_token def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): target_session_token = "" @@ -123,7 +124,7 @@ def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ target_feed_range = feed_range break - for i in range(10): + for i in range(100): item = { 'id': 'item' + str(uuid.uuid4()), 'name': 'sample', @@ -138,7 +139,7 @@ def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ feed_ranges_and_session_tokens.append((curr_feed_range, request_context["session_token"])) - return target_session_token, target_feed_range + return target_session_token, target_feed_range, previous_session_token @pytest.mark.cosmosEmulator @pytest.mark.unittest @@ -215,40 +216,49 @@ def test_updated_session_token_from_logical_pk(self, setup): feed_ranges_and_session_tokens = [] previous_session_token = "" target_pk = 'A1' - target_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) + target_session_token, previous_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert target_session_token == session_token + assert session_token == target_session_token trigger_split(setup) - target_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) + target_session_token, _ = create_items_logical_pk(setup, target_pk, session_token, feed_ranges_and_session_tokens) target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert target_session_token == session_token + assert session_token == target_session_token def test_updated_session_token_from_physical_pk(self, setup): feed_ranges_and_session_tokens = [] previous_session_token = "" pk_feed_range = setup[COLLECTION].feed_range_from_partition_key('A1') - target_session_token, target_feed_range = create_items_physical_pk(setup, pk_feed_range, + target_session_token, target_feed_range, previous_session_token = create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens) session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert target_session_token == session_token + assert session_token == target_session_token trigger_split(setup) - target_session_token, target_feed_range = create_items_physical_pk(setup, pk_feed_range, - previous_session_token, + _, target_feed_range, previous_session_token = create_items_physical_pk(setup, pk_feed_range, + session_token, feed_ranges_and_session_tokens) session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert target_session_token == session_token + assert is_compound_session_token(session_token) + session_tokens = session_token.split(",") + assert len(session_tokens) == 2 + pk_range_id1, session_token1 = create_vector_session_token_and_pkrange_id(session_tokens[0]) + pk_range_id2, session_token2 = create_vector_session_token_and_pkrange_id(session_tokens[1]) + pk_range_ids = [pk_range_id1, pk_range_id2] + + assert 320 == (session_token1.global_lsn + session_token2.global_lsn) + assert 1 in pk_range_ids + assert 2 in pk_range_ids From 51e721b4e85f17a2baac085763fc7e779bd002d1 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 8 Oct 2024 21:12:37 -0700 Subject: [PATCH 35/59] fix tests and pylint --- .../_change_feed/feed_range_internal.py | 4 +- .../azure/cosmos/_cosmos_client_connection.py | 13 +- .../azure/cosmos/_routing/routing_range.py | 4 +- .../azure/cosmos/_session_token_helpers.py | 121 +++++++------- .../aio/_cosmos_client_connection_async.py | 20 ++- .../azure-cosmos/azure/cosmos/container.py | 6 +- .../test/test_request_context_async.py | 2 +- .../test/test_session_token_helpers.py | 118 ------------- .../test/test_updated_session_token_split.py | 158 ++++++++++++++++++ 9 files changed, 248 insertions(+), 198 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index 4884da3d1ea3..c04fda0952f9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -25,10 +25,10 @@ import base64 import json from abc import ABC, abstractmethod -from typing import Union, List, Dict, Any, Optional, Sequence, Type +from typing import Union, List, Dict, Any, Optional from azure.cosmos._routing.routing_range import Range -from azure.cosmos.partition_key import _Undefined, _Empty, PartitionKey, NonePartitionKeyValue +from azure.cosmos.partition_key import _Undefined, _Empty class FeedRangeInternal(ABC): 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 047892a6c38e..2917a4f0bc0c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -848,7 +848,7 @@ def ReadPermission( path = base.GetPathFromLink(permission_link) permission_id = base.GetResourceIdOrFullNameFromLink(permission_link) - return self.Read(path, "permissions", permission_id, None, None, options, **kwargs) + return self.Read(path, "permissions", permission_id, None, options, **kwargs) def ReadPermissions( self, @@ -1032,7 +1032,7 @@ def ReadItems( :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. - :param request_context: A dictionary representing the request context. + :param dict request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -1124,7 +1124,7 @@ def QueryItemsChangeFeed( :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. - :param request_context: A dictionary representing the request context. + :param dict request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1160,7 +1160,7 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata - :param request_context: The request context for the operation + :param dict request_context: The request context for the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -3008,7 +3008,7 @@ def QueryFeed( partition_key_range_id, **kwargs) - def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements + def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements, disable=too-many-branches self, path: str, resource_type: str, @@ -3036,6 +3036,7 @@ def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements :param str partition_key_range_id: Specifies partition key range id. :param function response_hook: + :param dict request_context: :param bool is_query_plan: Specifies if the call is to fetch query plan :returns: A list of the queried resources. @@ -3386,4 +3387,4 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s def _add_request_context(self, request_context): if http_constants.HttpHeaders.SessionToken in self.last_response_headers: request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] - self.last_response_headers["request_context"] = request_context \ No newline at end of file + self.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index fd0c806b3dc7..44f19ca83b85 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -206,7 +206,9 @@ def can_merge(self, other): if self.isSingleValue() and other.isSingleValue(): return self.min == other.min # if share the same boundary, they can merge - if (self.max == other.min and self.isMaxInclusive or other.isMinInclusive) or (other.max == self.min and other.isMaxInclusive or self.isMinInclusive): + overlap_boundary1 = self.max == other.min and self.isMaxInclusive or other.isMinInclusive + overlap_boundary2 = other.max == self.min and other.isMaxInclusive or self.isMinInclusive + if overlap_boundary1 or overlap_boundary2: return True return self.overlaps(self, other) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index c7bbaad98606..23fda970cb04 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -24,10 +24,11 @@ from azure.cosmos._routing.routing_range import Range from azure.cosmos._vector_session_token import VectorSessionToken +# pylint: disable=protected-access def merge_session_tokens_with_same_range(session_token1, session_token2): - pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrange_id(session_token1) - pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrange_id(session_token2) + pk_range_id1, vector_session_token1 = parse_session_token(session_token1) + pk_range_id2, vector_session_token2 = parse_session_token(session_token2) pk_range_id = pk_range_id1 if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ @@ -38,19 +39,19 @@ def merge_session_tokens_with_same_range(session_token1, session_token2): def is_compound_session_token(session_token): return "," in session_token -def create_vector_session_token_and_pkrange_id(session_token): +def parse_session_token(session_token): tokens = session_token.split(":") return tokens[0], VectorSessionToken.create(tokens[1]) def split_compound_session_tokens(compound_session_tokens): session_tokens = [] for _, session_token in compound_session_tokens: - if is_compound_session_token(session_token): - tokens = session_token.split(",") - for token in tokens: - session_tokens.append(token) - else: - session_tokens.append(session_token) + if is_compound_session_token(session_token): + tokens = session_token.split(",") + for token in tokens: + session_tokens.append(token) + else: + session_tokens.append(session_token) return session_tokens def merge_session_tokens_with_same_pkrangeid(session_tokens): @@ -58,8 +59,8 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): while i < len(session_tokens): j = i + 1 while j < len(session_tokens): - pk_range_id1, vector_session_token1 = create_vector_session_token_and_pkrange_id(session_tokens[i]) - pk_range_id2, vector_session_token2 = create_vector_session_token_and_pkrange_id(session_tokens[j]) + pk_range_id1, vector_session_token1 = parse_session_token(session_tokens[i]) + pk_range_id2, vector_session_token2 = parse_session_token(session_tokens[j]) if pk_range_id1 == pk_range_id2: vector_session_token = vector_session_token1.merge(vector_session_token2) session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) @@ -73,50 +74,16 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): return session_tokens -def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): - target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() - # filter out tuples that overlap with target_feed_range and normalizes all the ranges - overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1]) - for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( - target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] - # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples - # with compound session tokens? - if len(overlapping_ranges) == 0: - raise ValueError('There were no overlapping feed ranges with the target.') - - # merge any session tokens that are the same exact feed range - i = 0 - j = 1 - while i < len(overlapping_ranges) and j < len(overlapping_ranges): - cur_feed_range = overlapping_ranges[i][0] - session_token = overlapping_ranges[i][1] - session_token_1 = overlapping_ranges[j][1] - if (not is_compound_session_token(session_token) and - not is_compound_session_token(overlapping_ranges[j][1]) and - cur_feed_range == overlapping_ranges[j][0]): - session_token = merge_session_tokens_with_same_range(session_token, session_token_1) - feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] - for feed_range_to_remove in feed_ranges_to_remove: - overlapping_ranges.remove(feed_range_to_remove) - overlapping_ranges.append((cur_feed_range, session_token)) - i, j = 0, 1 - else: - j += 1 - if j == len(overlapping_ranges): - i += 1 - j = i + 1 - - - done_overlapping_ranges = [] - # checking for merging of feed ranges that can be created from other feed ranges +def merge_ranges_with_subsets(overlapping_ranges): + processed_ranges = [] while len(overlapping_ranges) != 0: feed_range_cmp, session_token_cmp = overlapping_ranges[0] # compound session tokens are not considered for merging if is_compound_session_token(session_token_cmp): - done_overlapping_ranges.append(overlapping_ranges[0]) + processed_ranges.append(overlapping_ranges[0]) overlapping_ranges.remove(overlapping_ranges[0]) continue - _, vector_session_token_cmp = create_vector_session_token_and_pkrange_id(session_token_cmp) + _, vector_session_token_cmp = parse_session_token(session_token_cmp) subsets = [] # finding the subset feed ranges of the current feed range for j in range(1, len(overlapping_ranges)): @@ -133,24 +100,24 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): session_tokens = [subsets[j][1]] merged_indices = [subsets[j][2]] if len(subsets) == 1: - _, vector_session_token = create_vector_session_token_and_pkrange_id(session_tokens[0]) + _, vector_session_token = parse_session_token(session_tokens[0]) if vector_session_token_cmp.is_greater(vector_session_token): overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) else: - for k in range(len(subsets)): + for k, subset in enumerate(subsets): if j == k: continue - if merged_range.can_merge(subsets[k][0]): - merged_range = merged_range.merge(subsets[k][0]) - session_tokens.append(subsets[k][1]) - merged_indices.append(subsets[k][2]) + if merged_range.can_merge(subset[0]): + merged_range = merged_range.merge(subset[0]) + session_tokens.append(subset[1]) + merged_indices.append(subset[2]) if feed_range_cmp == merged_range: # if feed range can be created from the subsets # take the subsets if their global lsn is larger # else take the current feed range children_more_updated = True for session_token in session_tokens: - _, vector_session_token = create_vector_session_token_and_pkrange_id(session_token) + _, vector_session_token = parse_session_token(session_token) if vector_session_token_cmp.is_greater(vector_session_token): children_more_updated = False feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] @@ -164,11 +131,49 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): j += 1 - done_overlapping_ranges.append(overlapping_ranges[0]) + processed_ranges.append(overlapping_ranges[0]) overlapping_ranges.remove(overlapping_ranges[0]) + return processed_ranges + +def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): + target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() + # filter out tuples that overlap with target_feed_range and normalizes all the ranges + overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), + feed_range_to_session_token[1]) + for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( + target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] + # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples + # with compound session tokens? + if len(overlapping_ranges) == 0: + raise ValueError('There were no overlapping feed ranges with the target.') + + # merge any session tokens that are the same exact feed range + i = 0 + j = 1 + while i < len(overlapping_ranges) and j < len(overlapping_ranges): + cur_feed_range = overlapping_ranges[i][0] + session_token = overlapping_ranges[i][1] + session_token_1 = overlapping_ranges[j][1] + if (not is_compound_session_token(session_token) and + not is_compound_session_token(overlapping_ranges[j][1]) and + cur_feed_range == overlapping_ranges[j][0]): + session_token = merge_session_tokens_with_same_range(session_token, session_token_1) + feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] + for feed_range_to_remove in feed_ranges_to_remove: + overlapping_ranges.remove(feed_range_to_remove) + overlapping_ranges.append((cur_feed_range, session_token)) + i, j = 0, 1 + else: + j += 1 + if j == len(overlapping_ranges): + i += 1 + j = i + 1 + + # checking for merging of feed ranges that can be created from other feed ranges + processed_ranges = merge_ranges_with_subsets(overlapping_ranges) # break up session tokens that are compound - remaining_session_tokens = split_compound_session_tokens(done_overlapping_ranges) + remaining_session_tokens = split_compound_session_tokens(processed_ranges) if len(remaining_session_tokens) == 1: return remaining_session_tokens[0] 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 547b16407d73..58dc4eea0dcd 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 @@ -526,6 +526,7 @@ async def CreateItem( The Azure Cosmos document to create. :param dict options: The request options for the request. + :param dict request_context: :return: The created Document. :rtype: @@ -813,6 +814,7 @@ async def UpsertItem( The Azure Cosmos document to upsert. :param dict options: The request options for the request. + :param dict request_context: :return: The upserted Document. :rtype: @@ -1167,7 +1169,7 @@ async def Read( :param dict initial_headers: :param dict options: The request options for the request. - + :param dict request_context: :return: The upserted Azure Cosmos resource. :rtype: @@ -1185,7 +1187,7 @@ async def Read( result, last_response_headers = await self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers if request_context is not None: - self._add_request_context(request_context) + self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result @@ -1385,7 +1387,7 @@ async def ReplaceItem( :param dict new_document: :param dict options: The request options for the request. - :param request_context: + :param dict request_context: :return: The new Document. :rtype: @@ -1929,6 +1931,7 @@ async def Batch( :param str collection_link: The link to the collection :param list batch_operations: The batch of operations for the batch request. :param dict options: The request options for the request. + :param dict request_context: The request context for the operation :return: The result of the batch operation. @@ -2184,7 +2187,7 @@ def ReadItems( :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. - :param request_context: The request context for the operation. + :param dict request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -2214,7 +2217,7 @@ def QueryItems( :param dict options: The request options for the request. :param str partition_key: Partition key for the query(default value None) :param response_hook: A callable invoked with the response metadata. - :param request_context: The request context for the operation. + :param dict request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2279,7 +2282,7 @@ def QueryItemsChangeFeed( :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. - :param request_context: The request_context of the operation + :param dict request_context: The request_context of the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2314,7 +2317,7 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata - :param request_context: The request context of the operation. + :param dict request_context: The request context of the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2836,6 +2839,7 @@ async def __QueryFeed( # pylint: disable=too-many-branches,too-many-statements, :param function response_hook: :param bool is_query_plan: Specifies if the call is to fetch query plan + :param dict request_context: :returns: A list of the queried resources. :rtype: list :raises SystemError: If the query compatibility mode is undefined. @@ -3283,4 +3287,4 @@ async def _get_partition_key_definition(self, collection_link: str) -> Optional[ def _add_request_context(self, request_context): if http_constants.HttpHeaders.SessionToken in self.last_response_headers: request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] - self.last_response_headers["request_context"] = request_context \ No newline at end of file + self.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 1ec70a244956..eda6d09b6565 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -38,7 +38,6 @@ GenerateGuidId, _set_properties_cache ) -from ._change_feed.feed_range_internal import FeedRangeInternalEpk from ._cosmos_client_connection import CosmosClientConnection from ._feed_range import FeedRange, FeedRangeEpk from ._routing.routing_range import Range @@ -1440,6 +1439,5 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_range()) - - + return child_feed_range._feed_range_internal.get_normalized_range().is_subset( + parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py index 2de921aa3494..5d772ed09eca 100644 --- a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py @@ -83,7 +83,7 @@ async def test_crud_request_context(self, setup): for i in range(100): await setup["created_collection"].create_item(createItem()) items = [item async for item in setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)] - assert len(items) > 1000 + assert len(items) > 100 validate_request_context(setup["created_collection"]) items = [item async for item in setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 0d354be7a797..114a60831ee5 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -13,7 +13,6 @@ from azure.cosmos import DatabaseProxy from azure.cosmos._feed_range import FeedRangeEpk from azure.cosmos._routing.routing_range import Range -from azure.cosmos._session_token_helpers import is_compound_session_token, create_vector_session_token_and_pkrange_id COLLECTION = "created_collection" DATABASE = "created_db" @@ -77,70 +76,6 @@ def create_split_ranges(): actual_test_params.append((split_ranges, target_feed_range, test_param[2])) return actual_test_params -def trigger_split(setup): - print("Triggering a split in session token helpers") - setup[COLLECTION].replace_throughput(11000) - print("changed offer to 11k") - print("--------------------------------") - print("Waiting for split to complete") - start_time = time.time() - - while True: - offer = setup[COLLECTION].get_throughput() - if offer.properties['content'].get('isOfferReplacePending', False): - if time.time() - start_time > 60 * 25: # timeout test at 25 minutes - unittest.skip("Partition split didn't complete in time.") - else: - print("Waiting for split to complete") - time.sleep(60) - else: - break - - print("Split in session token helpers has completed") - -def create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens): - target_session_token = "" - for i in range(100): - item = { - 'id': 'item' + str(uuid.uuid4()), - 'name': 'sample', - 'pk': 'A' + str(random.randint(1, 10)) - } - setup[COLLECTION].create_item(item, session_token=previous_session_token) - request_context = setup[COLLECTION].client_connection.last_response_headers["request_context"] - if item['pk'] == target_pk: - target_session_token = request_context["session_token"] - previous_session_token = request_context["session_token"] - feed_ranges_and_session_tokens.append((setup[COLLECTION].feed_range_from_partition_key(item['pk']), - request_context["session_token"])) - return target_session_token, previous_session_token - -def create_items_physical_pk(setup, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): - target_session_token = "" - container_feed_ranges = setup[COLLECTION].read_feed_ranges() - target_feed_range = None - for feed_range in container_feed_ranges: - if setup[COLLECTION].is_feed_range_subset(feed_range, pk_feed_range): - target_feed_range = feed_range - break - - for i in range(100): - item = { - 'id': 'item' + str(uuid.uuid4()), - 'name': 'sample', - 'pk': 'A' + str(random.randint(1, 10)) - } - setup[COLLECTION].create_item(item, session_token=previous_session_token) - request_context = setup[COLLECTION].client_connection.last_response_headers["request_context"] - curr_feed_range = setup[COLLECTION].feed_range_from_partition_key(item['pk']) - if setup[COLLECTION].is_feed_range_subset(target_feed_range, curr_feed_range): - target_session_token = request_context["session_token"] - previous_session_token = request_context["session_token"] - feed_ranges_and_session_tokens.append((curr_feed_range, - request_context["session_token"])) - - return target_session_token, target_feed_range, previous_session_token - @pytest.mark.cosmosEmulator @pytest.mark.unittest @pytest.mark.usefixtures("setup") @@ -151,7 +86,6 @@ class TestSessionTokenHelpers: client: cosmos_client.CosmosClient = None host = test_config.TestConfig.host masterKey = test_config.TestConfig.masterKey - connectionPolicy = test_config.TestConfig.connectionPolicy configs = test_config.TestConfig TEST_DATABASE_ID = configs.TEST_DATABASE_ID TEST_COLLECTION_ID = configs.TEST_SINGLE_PARTITION_CONTAINER_ID @@ -212,57 +146,5 @@ def test_invalid_feed_range(self, setup): setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, FeedRangeEpk(Range("CC", "FF", True, False))) - def test_updated_session_token_from_logical_pk(self, setup): - feed_ranges_and_session_tokens = [] - previous_session_token = "" - target_pk = 'A1' - target_session_token, previous_session_token = create_items_logical_pk(setup, target_pk, previous_session_token, feed_ranges_and_session_tokens) - target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) - session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - - assert session_token == target_session_token - - trigger_split(setup) - - target_session_token, _ = create_items_logical_pk(setup, target_pk, session_token, feed_ranges_and_session_tokens) - target_feed_range = setup[COLLECTION].feed_range_from_partition_key(target_pk) - session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - - assert session_token == target_session_token - - - def test_updated_session_token_from_physical_pk(self, setup): - feed_ranges_and_session_tokens = [] - previous_session_token = "" - pk_feed_range = setup[COLLECTION].feed_range_from_partition_key('A1') - target_session_token, target_feed_range, previous_session_token = create_items_physical_pk(setup, pk_feed_range, - previous_session_token, - feed_ranges_and_session_tokens) - - session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert session_token == target_session_token - - trigger_split(setup) - - _, target_feed_range, previous_session_token = create_items_physical_pk(setup, pk_feed_range, - session_token, - feed_ranges_and_session_tokens) - - session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) - assert is_compound_session_token(session_token) - session_tokens = session_token.split(",") - assert len(session_tokens) == 2 - pk_range_id1, session_token1 = create_vector_session_token_and_pkrange_id(session_tokens[0]) - pk_range_id2, session_token2 = create_vector_session_token_and_pkrange_id(session_tokens[1]) - pk_range_ids = [pk_range_id1, pk_range_id2] - - assert 320 == (session_token1.global_lsn + session_token2.global_lsn) - assert 1 in pk_range_ids - assert 2 in pk_range_ids - - - - - if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py new file mode 100644 index 000000000000..849d00be0d1c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py @@ -0,0 +1,158 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import random +import time +import unittest +import uuid + + +import azure.cosmos.cosmos_client as cosmos_client +import test_config +from azure.cosmos import DatabaseProxy, PartitionKey +from azure.cosmos._session_token_helpers import is_compound_session_token, parse_session_token + + +class TestUpdatedSessionTokenSplit(unittest.TestCase): + """Test for session token helpers""" + + created_db: DatabaseProxy = None + client: cosmos_client.CosmosClient = None + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + configs = test_config.TestConfig + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) + cls.database = cls.client.get_database_client(cls.TEST_DATABASE_ID) + + + def test_updated_session_token_from_logical_pk(self): + container = self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + target_pk = 'A1' + target_session_token, previous_session_token = self.create_items_logical_pk(container, target_pk, + previous_session_token, + feed_ranges_and_session_tokens) + target_feed_range = container.feed_range_from_partition_key(target_pk) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + + self.trigger_split(container) + + target_session_token, _ = self.create_items_logical_pk(container, target_pk, session_token, + feed_ranges_and_session_tokens) + target_feed_range = container.feed_range_from_partition_key(target_pk) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + self.database.delete_container(container.id) + + def test_updated_session_token_from_physical_pk(self): + container = self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + pk_feed_range = container.feed_range_from_partition_key('A1') + target_session_token, target_feed_range, previous_session_token = self.create_items_physical_pk(container, pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert session_token == target_session_token + + self.trigger_split(container) + + _, target_feed_range, previous_session_token = self.create_items_physical_pk(container, pk_feed_range, + session_token, + feed_ranges_and_session_tokens) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert is_compound_session_token(session_token) + session_tokens = session_token.split(",") + assert len(session_tokens) == 2 + pk_range_id1, session_token1 = parse_session_token(session_tokens[0]) + pk_range_id2, session_token2 = parse_session_token(session_tokens[1]) + pk_range_ids = [pk_range_id1, pk_range_id2] + + assert 320 == (session_token1.global_lsn + session_token2.global_lsn) + assert '1' in pk_range_ids + assert '2' in pk_range_ids + self.database.delete_container(container.id) + + + @staticmethod + def trigger_split(container): + print("Triggering a split in session token helpers") + container.replace_throughput(11000) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = container.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + + print("Split in session token helpers has completed") + + @staticmethod + def create_items_logical_pk(container, target_pk, previous_session_token, feed_ranges_and_session_tokens): + target_session_token = "" + for i in range(100): + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + container.create_item(item, session_token=previous_session_token) + request_context = container.client_connection.last_response_headers["request_context"] + if item['pk'] == target_pk: + target_session_token = request_context["session_token"] + previous_session_token = request_context["session_token"] + feed_ranges_and_session_tokens.append((container.feed_range_from_partition_key(item['pk']), + request_context["session_token"])) + return target_session_token, previous_session_token + + @staticmethod + def create_items_physical_pk(container, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): + target_session_token = "" + container_feed_ranges = container.read_feed_ranges() + target_feed_range = None + for feed_range in container_feed_ranges: + if container.is_feed_range_subset(feed_range, pk_feed_range): + target_feed_range = feed_range + break + + for i in range(100): + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + container.create_item(item, session_token=previous_session_token) + request_context = container.client_connection.last_response_headers["request_context"] + curr_feed_range = container.feed_range_from_partition_key(item['pk']) + if container.is_feed_range_subset(target_feed_range, curr_feed_range): + target_session_token = request_context["session_token"] + previous_session_token = request_context["session_token"] + feed_ranges_and_session_tokens.append((curr_feed_range, request_context["session_token"])) + + return target_session_token, target_feed_range, previous_session_token + +if __name__ == '__main__': + unittest.main() From 104e341d6f5a1efd2833f330dbff97836f9e217b Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 9 Oct 2024 23:17:56 -0700 Subject: [PATCH 36/59] Reacting to comments --- .../azure/cosmos/_cosmos_client_connection.py | 1041 ++++++++--------- .../azure-cosmos/azure/cosmos/_feed_range.py | 3 +- .../azure/cosmos/_routing/routing_range.py | 6 +- .../azure/cosmos/_session_token_helpers.py | 47 +- .../azure/cosmos/_vector_session_token.py | 5 - .../azure/cosmos/aio/_container.py | 380 +++--- .../aio/_cosmos_client_connection_async.py | 884 +++++++------- .../azure-cosmos/azure/cosmos/container.py | 413 ++++--- .../azure-cosmos/test/test_feed_range.py | 22 +- .../test/test_feed_range_async.py | 78 ++ .../azure-cosmos/test/test_request_context.py | 97 -- .../test/test_request_context_async.py | 99 -- .../test/test_session_token_helpers.py | 53 +- .../test/test_updated_session_token_split.py | 17 +- 14 files changed, 1459 insertions(+), 1686 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/test/test_feed_range_async.py delete mode 100644 sdk/cosmos/azure-cosmos/test/test_request_context.py delete mode 100644 sdk/cosmos/azure-cosmos/test/test_request_context_async.py 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 2917a4f0bc0c..cf62e92f6f2c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -108,12 +108,12 @@ class _QueryCompatibilityMode: _DefaultStringRangePrecision = -1 def __init__( - self, - url_connection: str, - auth: CredentialDict, - connection_policy: Optional[ConnectionPolicy] = None, - consistency_level: Optional[str] = None, - **kwargs: Any + self, + url_connection: str, + auth: CredentialDict, + connection_policy: Optional[ConnectionPolicy] = None, + consistency_level: Optional[str] = None, + **kwargs: Any ) -> None: """ :param str url_connection: @@ -255,9 +255,9 @@ def _set_container_properties_cache(self, container_link: str, properties: Optio self.__container_properties_cache[container_link] = {} def _set_client_consistency_level( - self, - database_account: DatabaseAccount, - consistency_level: Optional[str], + self, + database_account: DatabaseAccount, + consistency_level: Optional[str], ) -> None: """Checks if consistency level param was passed in by user and sets it to that value or to the account default. @@ -348,10 +348,10 @@ def GetPartitionResolver(self, database_link: str) -> Optional[RangePartitionRes return self.partition_resolvers.get(base.TrimBeginningAndEndingSlashes(database_link)) def CreateDatabase( - self, - database: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a database. @@ -371,10 +371,10 @@ def CreateDatabase( return self.Create(database, path, "dbs", None, None, options, **kwargs) def ReadDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a database. @@ -396,9 +396,9 @@ def ReadDatabase( return self.Read(path, "dbs", database_id, None, options, **kwargs) def ReadDatabases( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all databases. @@ -417,10 +417,10 @@ def ReadDatabases( return self.QueryDatabases(None, options, **kwargs) def QueryDatabases( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries databases. @@ -438,18 +438,18 @@ def QueryDatabases( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - "/dbs", "dbs", "", lambda r: r["Databases"], - lambda _, b: b, query, options, **kwargs) + "/dbs", "dbs", "", lambda r: r["Databases"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def ReadContainers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all collections in a database. @@ -469,11 +469,11 @@ def ReadContainers( return self.QueryContainers(database_link, None, options, **kwargs) def QueryContainers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries collections in a database. @@ -496,19 +496,19 @@ def QueryContainers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "colls", database_id, lambda r: r["DocumentCollections"], - lambda _, body: body, query, options, **kwargs) + path, "colls", database_id, lambda r: r["DocumentCollections"], + lambda _, body: body, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateContainer( - self, - database_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a collection in a database. @@ -532,11 +532,11 @@ def CreateContainer( return self.Create(collection, path, "colls", database_id, None, options, **kwargs) def ReplaceContainer( - self, - collection_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a collection and return it. @@ -562,10 +562,10 @@ def ReplaceContainer( return self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) def ReadContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a collection. @@ -588,11 +588,11 @@ def ReadContainer( return self.Read(path, "colls", collection_id, None, options, **kwargs) def CreateUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user. @@ -616,11 +616,11 @@ def CreateUser( return self.Create(user, path, "users", database_id, None, options, **kwargs) def UpsertUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user. @@ -648,10 +648,10 @@ def _GetDatabaseIdWithPathForUser(self, database_link: str, user: Mapping[str, A return database_id, path def ReadUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user. @@ -674,10 +674,10 @@ def ReadUser( return self.Read(path, "users", user_id, None, options, **kwargs) def ReadUsers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all users in a database. @@ -697,11 +697,11 @@ def ReadUsers( return self.QueryUsers(database_link, None, options, **kwargs) def QueryUsers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries users in a database. @@ -725,18 +725,18 @@ def QueryUsers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "users", database_id, lambda r: r["Users"], - lambda _, b: b, query, options, **kwargs) + path, "users", database_id, lambda r: r["Users"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def DeleteDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a database. @@ -759,11 +759,11 @@ def DeleteDatabase( self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) def CreatePermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a permission for a user. @@ -787,11 +787,11 @@ def CreatePermission( return self.Create(permission, path, "permissions", user_id, None, options, **kwargs) def UpsertPermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a permission for a user. @@ -815,9 +815,9 @@ def UpsertPermission( return self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) def _GetUserIdWithPathForPermission( - self, - permission: Mapping[str, Any], - user_link: str + self, + permission: Mapping[str, Any], + user_link: str ) -> Tuple[str, Optional[str]]: base._validate_resource(permission) path = base.GetPathFromLink(user_link, "permissions") @@ -825,10 +825,10 @@ def _GetUserIdWithPathForPermission( return path, user_id def ReadPermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a permission. @@ -851,10 +851,10 @@ def ReadPermission( return self.Read(path, "permissions", permission_id, None, options, **kwargs) def ReadPermissions( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all permissions for a user. @@ -875,11 +875,11 @@ def ReadPermissions( return self.QueryPermissions(user_link, None, options, **kwargs) def QueryPermissions( - self, - user_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries permissions for a user. @@ -903,19 +903,19 @@ def QueryPermissions( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "permissions", user_id, lambda r: r["Permissions"], - lambda _, b: b, query, options, **kwargs) + path, "permissions", user_id, lambda r: r["Permissions"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def ReplaceUser( - self, - user_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user and return it. @@ -940,10 +940,10 @@ def ReplaceUser( return self.Replace(user, path, "users", user_id, None, options, **kwargs) def DeleteUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user. @@ -966,11 +966,11 @@ def DeleteUser( self.DeleteResource(path, "users", user_id, None, options, **kwargs) def ReplacePermission( - self, - permission_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a permission and return it. @@ -995,10 +995,10 @@ def ReplacePermission( return self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) def DeletePermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a permission. @@ -1021,18 +1021,16 @@ def DeletePermission( self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) def ReadItems( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all documents in a collection. :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. - :param dict request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -1040,18 +1038,16 @@ def ReadItems( if feed_options is None: feed_options = {} - return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, - request_context=request_context, **kwargs) + return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) def QueryItems( - self, - database_or_container_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key: Optional[PartitionKeyType] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key: Optional[PartitionKeyType] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -1062,7 +1058,6 @@ def QueryItems( :param partition_key: Partition key for the query(default value None) :type: partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] :param response_hook: A callable invoked with the response metadata. - :param request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1091,16 +1086,15 @@ def QueryItems( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - response_hook=response_hook, - request_context=request_context, - **kwargs) + path, + "docs", + collection_id, + lambda r: r["Documents"], + lambda _, b: b, + query, + options, + response_hook=response_hook, + **kwargs) return ItemPaged( self, @@ -1112,19 +1106,17 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def QueryItemsChangeFeed( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. - :param dict request_context: A dictionary representing the request context. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1139,19 +1131,17 @@ def QueryItemsChangeFeed( partition_key_range_id = options["partitionKeyRangeId"] return self._QueryChangeFeed( - collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, - request_context=request_context, **kwargs + collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs ) def _QueryChangeFeed( - self, - collection_link: str, - resource_type: str, - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + resource_type: str, + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries change feed of a resource in a collection. @@ -1160,7 +1150,6 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata - :param dict request_context: The request context for the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: @@ -1199,7 +1188,6 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str options, partition_key_range_id, response_hook=response_hook, - request_context=request_context, **kwargs) return ItemPaged( @@ -1211,10 +1199,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def _ReadPartitionKeyRanges( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads Partition Key Ranges. @@ -1230,11 +1218,11 @@ def _ReadPartitionKeyRanges( return self._QueryPartitionKeyRanges(collection_link, None, feed_options, **kwargs) def _QueryPartitionKeyRanges( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries Partition Key Ranges in a collection. @@ -1258,20 +1246,19 @@ def _QueryPartitionKeyRanges( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], - lambda _, b: b, query, options, **kwargs) + path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - request_context: Dict[str, Any] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -1279,7 +1266,6 @@ def CreateItem( The link to the database when using partitioning, otherwise link to the document collection. :param dict document: The Azure Cosmos document to create. :param dict options: The request options for the request. - :param dict request_context: The request context for the request. :return: The created Document. :rtype: dict """ @@ -1301,22 +1287,20 @@ def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = self._AddPartitionKey(database_or_container_link, document, options) - return self.Create(document, path, "docs", collection_id, None, options, request_context, **kwargs) + return self.Create(document, path, "docs", collection_id, None, options, **kwargs) def UpsertItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - request_context: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a document in a collection. :param str database_or_container_link: The link to the database when using partitioning, otherwise link to the document collection. :param dict document: The Azure Cosmos document to upsert. - :param dict request_context: The request context for the request. :param dict options: The request options for the request. :return: The upserted Document. :rtype: dict @@ -1339,7 +1323,7 @@ def UpsertItem( collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options ) - return self.Upsert(document, path, "docs", collection_id, None, options, request_context, **kwargs) + return self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) PartitionResolverErrorMessage = ( "Couldn't find any partition resolvers for the database link provided. " @@ -1350,10 +1334,10 @@ def UpsertItem( # Gets the collection id and path for the document def _GetContainerIdWithPathForItem( - self, - database_or_container_link: str, - document: Mapping[str, Any], - options: Mapping[str, Any] + self, + database_or_container_link: str, + document: Mapping[str, Any], + options: Mapping[str, Any] ) -> Tuple[Optional[str], Dict[str, Any], str]: if not database_or_container_link: @@ -1382,17 +1366,15 @@ def _GetContainerIdWithPathForItem( return collection_id, document, path def ReadItem( - self, - document_link: str, - request_context: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Reads a document. :param str document_link: The link to the document. - :param dict request_context: :param dict options: The request options for the request. @@ -1407,13 +1389,13 @@ def ReadItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return self.Read(path, "docs", document_id, None, options, request_context, **kwargs) + return self.Read(path, "docs", document_id, None, options, **kwargs) def ReadTriggers( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all triggers in a collection. @@ -1434,11 +1416,11 @@ def ReadTriggers( return self.QueryTriggers(collection_link, None, options, **kwargs) def QueryTriggers( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries triggers in a collection. @@ -1470,11 +1452,11 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def CreateTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a trigger in a collection. @@ -1497,9 +1479,9 @@ def CreateTrigger( return self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForTrigger( - self, - collection_link: str, - trigger: Mapping[str, Any] + self, + collection_link: str, + trigger: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(trigger) trigger = dict(trigger) @@ -1513,10 +1495,10 @@ def _GetContainerIdWithPathForTrigger( return collection_id, path, trigger def ReadTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a trigger. @@ -1539,11 +1521,11 @@ def ReadTrigger( return self.Read(path, "triggers", trigger_id, None, options, **kwargs) def UpsertTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a trigger in a collection. :param str collection_link: @@ -1563,10 +1545,10 @@ def UpsertTrigger( return self.Upsert(trigger, path, "triggers", collection_id, None, options, **kwargs) def ReadUserDefinedFunctions( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all user-defined functions in a collection. @@ -1587,11 +1569,11 @@ def ReadUserDefinedFunctions( return self.QueryUserDefinedFunctions(collection_link, None, options, **kwargs) def QueryUserDefinedFunctions( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries user-defined functions in a collection. @@ -1615,19 +1597,19 @@ def QueryUserDefinedFunctions( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], - lambda _, b: b, query, options, **kwargs) + path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user-defined function in a collection. @@ -1650,11 +1632,11 @@ def CreateUserDefinedFunction( return self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) def UpsertUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user-defined function in a collection. @@ -1677,9 +1659,9 @@ def UpsertUserDefinedFunction( return self.Upsert(udf, path, "udfs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForUDF( - self, - collection_link: str, - udf: Mapping[str, Any] + self, + collection_link: str, + udf: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(udf) udf = dict(udf) @@ -1693,10 +1675,10 @@ def _GetContainerIdWithPathForUDF( return collection_id, path, udf def ReadUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user-defined function. @@ -1719,10 +1701,10 @@ def ReadUserDefinedFunction( return self.Read(path, "udfs", udf_id, None, options, **kwargs) def ReadStoredProcedures( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all store procedures in a collection. @@ -1743,11 +1725,11 @@ def ReadStoredProcedures( return self.QueryStoredProcedures(collection_link, None, options, **kwargs) def QueryStoredProcedures( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries stored procedures in a collection. @@ -1779,11 +1761,11 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def CreateStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a stored procedure in a collection. @@ -1806,11 +1788,11 @@ def CreateStoredProcedure( return self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) def UpsertStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a stored procedure in a collection. @@ -1833,9 +1815,9 @@ def UpsertStoredProcedure( return self.Upsert(sproc, path, "sprocs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForSproc( - self, - collection_link: str, - sproc: Mapping[str, Any] + self, + collection_link: str, + sproc: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(sproc) sproc = dict(sproc) @@ -1848,10 +1830,10 @@ def _GetContainerIdWithPathForSproc( return collection_id, path, sproc def ReadStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a stored procedure. @@ -1874,10 +1856,10 @@ def ReadStoredProcedure( return self.Read(path, "sprocs", sproc_id, None, options, **kwargs) def ReadConflicts( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads conflicts. @@ -1897,11 +1879,11 @@ def ReadConflicts( return self.QueryConflicts(collection_link, None, feed_options, **kwargs) def QueryConflicts( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries conflicts in a collection. @@ -1933,10 +1915,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def ReadConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a conflict. @@ -1958,10 +1940,10 @@ def ReadConflict( return self.Read(path, "conflicts", conflict_id, None, options, **kwargs) def DeleteContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a collection. @@ -1984,19 +1966,17 @@ def DeleteContainer( self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) def ReplaceItem( - self, - document_link: str, - new_document: Dict[str, Any], - request_context: Optional[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + new_document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a document and returns it. :param str document_link: The link to the document. :param dict new_document: - :param dict request_context: :param dict options: The request options for the request. @@ -2024,21 +2004,19 @@ def ReplaceItem( collection_link = base.GetItemContainerLink(document_link) options = self._AddPartitionKey(collection_link, new_document, options) - return self.Replace(new_document, path, "docs", document_id, None, options, request_context, **kwargs) + return self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) def PatchItem( - self, - document_link: str, - operations: List[Dict[str, Any]], - request_context: Optional[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + operations: List[Dict[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Patches a document and returns it. :param str document_link: The link to the document. :param list operations: The operations for the patch request. - :param dict request_context: The request context for the request. :param dict options: The request options for the request. :return: @@ -2066,25 +2044,21 @@ def PatchItem( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result def Batch( - self, - collection_link: str, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - request_context: Optional[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """Executes the given operations in transactional batch. :param str collection_link: The link to the collection :param list batch_operations: The batch of operations for the batch request. - :param dict request_context: The request context of the operation. :param dict options: The request options for the request. :return: @@ -2109,8 +2083,6 @@ def Batch( **kwargs ) self.last_response_headers = last_response_headers - if request_context is not None: - self._add_request_context(request_context) final_responses = [] is_error = False error_status = 0 @@ -2139,12 +2111,12 @@ def Batch( return final_responses def _Batch( - self, - batch_operations: List[Dict[str, Any]], - path: str, - collection_id: Optional[str], - options: Mapping[str, Any], - **kwargs: Any + self, + batch_operations: List[Dict[str, Any]], + path: str, + collection_id: Optional[str], + options: Mapping[str, Any], + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) @@ -2156,17 +2128,15 @@ def _Batch( ) def DeleteItem( - self, - document_link: str, - request_context: Optional[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a document. :param str document_link: The link to the document. - :param dict request_context: :param dict options: The request options for the request. @@ -2181,13 +2151,13 @@ def DeleteItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - self.DeleteResource(path, "docs", document_id, None, options, request_context, **kwargs) + self.DeleteResource(path, "docs", document_id, None, options, **kwargs) def DeleteAllItemsByPartitionKey( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Exposes an API to delete all items with a single partition key without the user having to explicitly call delete on each record in the partition key. @@ -2223,11 +2193,11 @@ def DeleteAllItemsByPartitionKey( response_hook(last_response_headers, None) def ReplaceTrigger( - self, - trigger_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a trigger and returns it. @@ -2258,10 +2228,10 @@ def ReplaceTrigger( return self.Replace(trigger, path, "triggers", trigger_id, None, options, **kwargs) def DeleteTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a trigger. @@ -2284,11 +2254,11 @@ def DeleteTrigger( self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) def ReplaceUserDefinedFunction( - self, - udf_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user-defined function and returns it. @@ -2319,10 +2289,10 @@ def ReplaceUserDefinedFunction( return self.Replace(udf, path, "udfs", udf_id, None, options, **kwargs) def DeleteUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user-defined function. @@ -2345,11 +2315,11 @@ def DeleteUserDefinedFunction( self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) def ExecuteStoredProcedure( - self, - sproc_link: str, - params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Executes a store procedure. @@ -2385,11 +2355,11 @@ def ExecuteStoredProcedure( return result def ReplaceStoredProcedure( - self, - sproc_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a stored procedure and returns it. @@ -2420,10 +2390,10 @@ def ReplaceStoredProcedure( return self.Replace(sproc, path, "sprocs", sproc_id, None, options, **kwargs) def DeleteStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a stored procedure. @@ -2446,10 +2416,10 @@ def DeleteStoredProcedure( self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) def DeleteConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a conflict. @@ -2472,10 +2442,10 @@ def DeleteConflict( self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) def ReplaceOffer( - self, - offer_link: str, - offer: Dict[str, Any], - **kwargs: Any + self, + offer_link: str, + offer: Dict[str, Any], + **kwargs: Any ) -> Dict[str, Any]: """Replaces an offer and returns it. @@ -2492,12 +2462,12 @@ def ReplaceOffer( base._validate_resource(offer) path = base.GetPathFromLink(offer_link) offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Replace(offer, path, "offers", offer_id, None, **kwargs) + return self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) def ReadOffer( - self, - offer_link: str, - **kwargs: Any + self, + offer_link: str, + **kwargs: Any ) -> Dict[str, Any]: """Reads an offer. :param str offer_link: @@ -2509,12 +2479,12 @@ def ReadOffer( """ path = base.GetPathFromLink(offer_link) offer_id = base.GetResourceIdOrFullNameFromLink(offer_link) - return self.Read(path, "offers", offer_id, None,{}, **kwargs) + return self.Read(path, "offers", offer_id, None, {}, **kwargs) def ReadOffers( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all offers. :param dict options: @@ -2530,10 +2500,10 @@ def ReadOffers( return self.QueryOffers(None, options, **kwargs) def QueryOffers( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Query for all offers. @@ -2552,17 +2522,17 @@ def QueryOffers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - "/offers", "offers", "", lambda r: r["Offers"], - lambda _, b: b, query, options, **kwargs) + "/offers", "offers", "", lambda r: r["Offers"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def GetDatabaseAccount( - self, - url_connection: Optional[str] = None, - **kwargs: Any + self, + url_connection: Optional[str] = None, + **kwargs: Any ) -> DatabaseAccount: """Gets database account info. @@ -2609,15 +2579,14 @@ def GetDatabaseAccount( return database_account def Create( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -2626,7 +2595,6 @@ def Create( :param str typ: :param str id: :param dict initial_headers: - :param dict request_context: :param dict options: The request options for the request. @@ -2650,22 +2618,19 @@ def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result def Upsert( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -2674,7 +2639,6 @@ def Upsert( :param str typ: :param str id: :param dict initial_headers: - :param dict request_context: :param dict options: The request options for the request. @@ -2698,22 +2662,19 @@ def Upsert( self.last_response_headers = last_response_headers # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result def Replace( - self, - resource: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + resource: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -2722,7 +2683,6 @@ def Replace( :param str typ: :param str id: :param dict initial_headers: - :param dict request_context: :param dict options: The request options for the request. @@ -2745,21 +2705,18 @@ def Replace( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result def Read( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -2767,7 +2724,6 @@ def Read( :param str typ: :param str id: :param dict initial_headers: - :param dict request_context: :param dict options: The request options for the request. @@ -2787,21 +2743,18 @@ def Read( request_params = RequestObject(typ, documents._OperationType.Read) result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result def DeleteResource( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -2809,7 +2762,6 @@ def DeleteResource( :param str typ: :param str id: :param dict initial_headers: - :param dict request_context: :param dict options: The request options for the request. @@ -2832,17 +2784,15 @@ def DeleteResource( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, None) def __Get( - self, - path: str, - request_params: RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'GET' http request. @@ -2865,12 +2815,12 @@ def __Get( ) def __Post( - self, - path: str, - request_params: RequestObject, - body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'POST' http request. @@ -2894,12 +2844,12 @@ def __Post( ) def __Put( - self, - path: str, - request_params: RequestObject, - body: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + body: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PUT' http request. @@ -2923,12 +2873,12 @@ def __Put( ) def __Patch( - self, - path: str, - request_params: RequestObject, - request_data: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + request_data: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PATCH' http request. @@ -2952,11 +2902,11 @@ def __Patch( ) def __Delete( - self, - path: str, - request_params: RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[None, Dict[str, Any]]: """Azure Cosmos 'DELETE' http request. @@ -2979,13 +2929,13 @@ def __Delete( ) def QueryFeed( - self, - path: str, - collection_id: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Mapping[str, Any], - partition_key_range_id: Optional[str] = None, - **kwargs: Any + self, + path: str, + collection_id: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Mapping[str, Any], + partition_key_range_id: Optional[str] = None, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query Feed for Document Collection resource. @@ -2998,30 +2948,29 @@ def QueryFeed( :rtype: tuple of (dict, dict) """ return self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - partition_key_range_id, - **kwargs) + path, + "docs", + collection_id, + lambda r: r["Documents"], + lambda _, b: b, + query, + options, + partition_key_range_id, + **kwargs) - def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements, disable=too-many-branches - self, - path: str, - resource_type: str, - resource_id: Optional[str], - result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], - create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - is_query_plan: bool = False, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements + self, + path: str, + resource_type: str, + resource_id: Optional[str], + result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], + create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + is_query_plan: bool = False, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -3036,7 +2985,6 @@ def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements, disabl :param str partition_key_range_id: Specifies partition key range id. :param function response_hook: - :param dict request_context: :param bool is_query_plan: Specifies if the call is to fetch query plan :returns: A list of the queried resources. @@ -3082,8 +3030,6 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: result, last_response_headers = self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return __GetBodiesFromQueryResult(result), last_response_headers @@ -3166,9 +3112,6 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: results["Documents"].extend(partial_result["Documents"]) else: results = partial_result - - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, partial_result) # if the prefix partition query has results lets return it @@ -3181,8 +3124,7 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: index_metrics_raw = last_response_headers[INDEX_METRICS_HEADER] last_response_headers[INDEX_METRICS_HEADER] = _utils.get_index_metrics_info(index_metrics_raw) self.last_response_headers = last_response_headers - if request_context is not None: - self._add_request_context(request_context) + if response_hook: response_hook(last_response_headers, result) @@ -3263,10 +3205,10 @@ def __CheckAndUnifyQueryFormat(self, query_body: Union[str, Dict[str, Any]]) -> # Adds the partition key to options def _AddPartitionKey( - self, - collection_link: str, - document: Mapping[str, Any], - options: Mapping[str, Any] + self, + collection_link: str, + document: Mapping[str, Any], + options: Mapping[str, Any] ) -> Dict[str, Any]: collection_link = base.TrimBeginningAndEndingSlashes(collection_link) partitionKeyDefinition = self._get_partition_key_definition(collection_link) @@ -3281,9 +3223,9 @@ def _AddPartitionKey( # Extracts the partition key from the document using the partitionKey definition def _ExtractPartitionKey( - self, - partitionKeyDefinition: Mapping[str, Any], - document: Mapping[str, Any] + self, + partitionKeyDefinition: Mapping[str, Any], + document: Mapping[str, Any] ) -> Union[List[Optional[Union[str, float, bool]]], str, float, bool, _Empty, _Undefined]: if partitionKeyDefinition["kind"] == "MultiHash": ret: List[Optional[Union[str, float, bool]]] = [] @@ -3310,10 +3252,10 @@ def _ExtractPartitionKey( # Navigates the document to retrieve the partitionKey specified in the partition key parts def _retrieve_partition_key( - self, - partition_key_parts: List[str], - document: Mapping[str, Any], - is_system_key: bool + self, + partition_key_parts: List[str], + document: Mapping[str, Any], + is_system_key: bool ) -> Union[str, float, bool, _Empty, _Undefined]: expected_matchCount = len(partition_key_parts) matchCount = 0 @@ -3345,10 +3287,10 @@ def _refresh_container_properties_cache(self, container_link: str): self._set_container_properties_cache(container_link, _set_properties_cache(container)) def _UpdateSessionIfRequired( - self, - request_headers: Mapping[str, Any], - response_result: Optional[Mapping[str, Any]], - response_headers: Optional[Mapping[str, Any]] + self, + request_headers: Mapping[str, Any], + response_result: Optional[Mapping[str, Any]], + response_headers: Optional[Mapping[str, Any]] ) -> None: """ Updates session if necessary. @@ -3383,8 +3325,3 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s partition_key_definition = container.get("partitionKey") self.__container_properties_cache[collection_link] = _set_properties_cache(container) return partition_key_definition - - def _add_request_context(self, request_context): - if http_constants.HttpHeaders.SessionToken in self.last_response_headers: - request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] - self.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 2bda669b6bc0..f4b8dc576cb5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -51,11 +51,12 @@ def from_string(json_str: str) -> 'FeedRange': class FeedRangeEpk(FeedRange): type_property_name = "Range" - def __init__(self, feed_range: Range) -> None: + def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") self._feed_range_internal = FeedRangeInternalEpk(feed_range) + self._container_link = container_link def __str__(self) -> str: """Get a json representation of the feed range. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index 44f19ca83b85..ba1a88b2d4f3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -202,7 +202,7 @@ def overlaps(range1, range2): return True return False - def can_merge(self, other): + def can_merge(self, other: 'Range') -> bool: if self.isSingleValue() and other.isSingleValue(): return self.min == other.min # if share the same boundary, they can merge @@ -212,7 +212,7 @@ def can_merge(self, other): return True return self.overlaps(self, other) - def merge(self, other): + def merge(self, other: 'Range') -> 'Range': if not self.can_merge(other): raise ValueError("Ranges do not overlap") min_val = self.min if self.min < other.min else other.min @@ -221,7 +221,7 @@ def merge(self, other): is_max_inclusive = self.isMaxInclusive if self.max > other.max else other.isMaxInclusive return Range(min_val, max_val, is_min_inclusive, is_max_inclusive) - def is_subset(self, parent_range) -> bool: + def is_subset(self, parent_range: 'Range') -> bool: normalized_parent_range = parent_range.to_normalized_range() normalized_child_range = self.to_normalized_range() return normalized_parent_range.contains(normalized_child_range.min) and \ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 23fda970cb04..5c94a9f3e094 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -22,28 +22,29 @@ """Internal Helper functions for manipulating session tokens. """ from azure.cosmos._routing.routing_range import Range +from ._feed_range import FeedRange from azure.cosmos._vector_session_token import VectorSessionToken # pylint: disable=protected-access -def merge_session_tokens_with_same_range(session_token1, session_token2): +def merge_session_tokens_with_same_range(session_token1: str, session_token2: str) -> str: pk_range_id1, vector_session_token1 = parse_session_token(session_token1) pk_range_id2, vector_session_token2 = parse_session_token(session_token2) pk_range_id = pk_range_id1 if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ - if vector_session_token1.is_greater(vector_session_token2) else pk_range_id2 + if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 vector_session_token = vector_session_token1.merge(vector_session_token2) return pk_range_id + ":" + vector_session_token.session_token -def is_compound_session_token(session_token): +def is_compound_session_token(session_token: str) -> bool: return "," in session_token -def parse_session_token(session_token): +def parse_session_token(session_token: str) -> (str, VectorSessionToken): tokens = session_token.split(":") return tokens[0], VectorSessionToken.create(tokens[1]) -def split_compound_session_tokens(compound_session_tokens): +def split_compound_session_tokens(compound_session_tokens: [(Range, str)]) -> [str]: session_tokens = [] for _, session_token in compound_session_tokens: if is_compound_session_token(session_token): @@ -54,7 +55,7 @@ def split_compound_session_tokens(compound_session_tokens): session_tokens.append(session_token) return session_tokens -def merge_session_tokens_with_same_pkrangeid(session_tokens): +def merge_session_tokens_for_same_physical_pk(session_tokens: [str]) -> [str]: i = 0 while i < len(session_tokens): j = i + 1 @@ -74,7 +75,7 @@ def merge_session_tokens_with_same_pkrangeid(session_tokens): return session_tokens -def merge_ranges_with_subsets(overlapping_ranges): +def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: processed_ranges = [] while len(overlapping_ranges) != 0: feed_range_cmp, session_token_cmp = overlapping_ranges[0] @@ -101,7 +102,7 @@ def merge_ranges_with_subsets(overlapping_ranges): merged_indices = [subsets[j][2]] if len(subsets) == 1: _, vector_session_token = parse_session_token(session_tokens[0]) - if vector_session_token_cmp.is_greater(vector_session_token): + if vector_session_token_cmp.global_lsn > vector_session_token.global_lsn: overlapping_ranges.remove(overlapping_ranges[merged_indices[0]]) else: for k, subset in enumerate(subsets): @@ -116,16 +117,22 @@ def merge_ranges_with_subsets(overlapping_ranges): # take the subsets if their global lsn is larger # else take the current feed range children_more_updated = True + parent_more_updated = True for session_token in session_tokens: _, vector_session_token = parse_session_token(session_token) - if vector_session_token_cmp.is_greater(vector_session_token): + if vector_session_token_cmp.global_lsn > vector_session_token.global_lsn: children_more_updated = False + else: + parent_more_updated = False feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices] for feed_range_to_remove in feed_ranges_to_remove: overlapping_ranges.remove(feed_range_to_remove) if children_more_updated: overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) overlapping_ranges.remove(overlapping_ranges[0]) + elif not parent_more_updated and not children_more_updated: + session_tokens.append(session_token_cmp) + overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens)))) not_found = False break @@ -135,15 +142,21 @@ def merge_ranges_with_subsets(overlapping_ranges): overlapping_ranges.remove(overlapping_ranges[0]) return processed_ranges -def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): +def get_updated_session_token(feed_ranges_to_session_tokens: [(FeedRange, str)], target_feed_range: FeedRange, + container_link: str): + if target_feed_range._container_link != container_link: + raise ValueError('The target feed range does not belong to the container.') target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges - overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), - feed_range_to_session_token[1]) - for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps( - target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())] - # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples - # with compound session tokens? + overlapping_ranges = [] + for feed_range_to_session_token in feed_ranges_to_session_tokens: + if feed_range_to_session_token[0]._container_link != container_link: + raise ValueError('The feed range does not belong to the container.') + if Range.overlaps(target_feed_range_normalized, + feed_range_to_session_token[0]._feed_range_internal.get_normalized_range()): + overlapping_ranges.append((feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), + feed_range_to_session_token[1])) + if len(overlapping_ranges) == 0: raise ValueError('There were no overlapping feed ranges with the target.') @@ -178,7 +191,7 @@ def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range): if len(remaining_session_tokens) == 1: return remaining_session_tokens[0] # merging any session tokens with same physical partition key range id - remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens) + remaining_session_tokens = merge_session_tokens_for_same_physical_pk(remaining_session_tokens) updated_session_token = "" # compound the remaining session tokens diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py index a5be01799447..f8285f000b67 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py @@ -113,11 +113,6 @@ def equals(self, other): and self.are_region_progress_equal(other.local_lsn_by_region) ) - def is_greater(self, other): - if self.global_lsn > other.global_lsn: - return True - return False - def merge(self, other): if other is None: raise ValueError("Invalid Session Token (should not be None)") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index dcfb228c1bc7..c973ebd9ae95 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -77,11 +77,11 @@ class ContainerProxy: """ def __init__( - self, - client_connection: CosmosClientConnection, - database_link: str, - id: str, - properties: Optional[Dict[str, Any]] = None + self, + client_connection: CosmosClientConnection, + database_link: str, + id: str, + properties: Optional[Dict[str, Any]] = None ) -> None: self.client_connection = client_connection self.id = id @@ -130,8 +130,8 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> return conflict_or_link["_self"] async def _set_partition_key( - self, - partition_key: PartitionKeyType + self, + partition_key: PartitionKeyType ) -> Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(await self.is_system_key) @@ -147,14 +147,14 @@ async def _get_epk_range_for_partition_key(self, partition_key_value: PartitionK @distributed_trace_async async def read( - self, - *, - populate_partition_key_range_statistics: Optional[bool] = None, - populate_quota_info: Optional[bool] = None, - session_token: Optional[str] = None, - priority: Optional[Literal["High", "Low"]] = None, - initial_headers: Optional[Dict[str, str]] = None, - **kwargs: Any + self, + *, + populate_partition_key_range_statistics: Optional[bool] = None, + populate_quota_info: Optional[bool] = None, + session_token: Optional[str] = None, + priority: Optional[Literal["High", "Low"]] = None, + initial_headers: Optional[Dict[str, str]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Read the container properties. @@ -191,20 +191,20 @@ async def read( @distributed_trace_async async def create_item( - self, - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - indexing_directive: Optional[int] = None, - enable_automatic_id_generation: bool = False, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + indexing_directive: Optional[int] = None, + enable_automatic_id_generation: bool = False, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Create an item in the container. @@ -231,8 +231,8 @@ async def create_item( before high priority requests start getting throttled. Feature must first be enabled at the account level. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -260,23 +260,22 @@ async def create_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = await self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, options=request_options, - request_context={}, **kwargs + database_or_container_link=self.container_link, document=body, options=request_options, **kwargs ) return result or {} @distributed_trace_async async def read_item( - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - *, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + *, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Get the item identified by `item`. @@ -327,19 +326,18 @@ async def read_item( if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - return await self.client_connection.ReadItem(document_link=doc_link, options=request_options, - request_context={}, **kwargs) + return await self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) @distributed_trace def read_all_items( - self, - *, - max_item_count: Optional[int] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + *, + max_item_count: Optional[int] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """List all the items in the container. @@ -376,8 +374,7 @@ def read_all_items( feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] items = self.client_connection.ReadItems( - collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, - request_context={}, **kwargs + collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs ) if response_hook: response_hook(self.client_connection.last_response_headers, items) @@ -385,21 +382,21 @@ def read_all_items( @distributed_trace def query_items( - self, - query: str, - *, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: 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, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, - **kwargs: Any + self, + query: str, + *, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: 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, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + continuation_token_limit: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Return all results matching the given `query`. @@ -496,7 +493,6 @@ def query_items( options=feed_options, partition_key=partition_key, response_hook=response_hook, - request_context= {}, **kwargs ) if response_hook: @@ -690,16 +686,14 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["maxItemCount"] = kwargs.pop('max_item_count') if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] =\ + change_feed_state_context["partitionKey"] = \ self._set_partition_key(cast(PartitionKeyType, kwargs.get("partition_key"))) change_feed_state_context["partitionKeyFeedRange"] = \ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) - request_context = {} if kwargs.get("feed_range") is not None: feed_range: FeedRangeEpk = kwargs.pop('feed_range') change_feed_state_context["feedRange"] = feed_range._feed_range_internal - request_context["feedRange"] = feed_range feed_options["containerProperties"] = self._get_properties() feed_options["changeFeedStateContext"] = change_feed_state_context @@ -712,8 +706,7 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.QueryItemsChangeFeed( - self.container_link, options=feed_options, response_hook=response_hook, - request_context=request_context, **kwargs + self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) if response_hook: @@ -722,18 +715,18 @@ def query_items_change_feed( # pylint: disable=unused-argument @distributed_trace_async async def upsert_item( - self, - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Insert or update the specified item. @@ -756,9 +749,9 @@ async def upsert_item( before high priority requests start getting throttled. Feature must first be enabled at the account level. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the upsert operation went through. The dict will be empty if + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after the upsert operation went through. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -787,26 +780,25 @@ async def upsert_item( database_or_container_link=self.container_link, document=body, options=request_options, - request_context={}, **kwargs ) return result or {} @distributed_trace_async async def replace_item( - self, - item: Union[str, Mapping[str, Any]], - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces the specified item if it exists in the container. @@ -831,9 +823,9 @@ async def replace_item( :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace failed or the item with given id does not exist. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -860,27 +852,26 @@ async def replace_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = await self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, options=request_options, - request_context={}, **kwargs + document_link=item_link, new_document=body, options=request_options, **kwargs ) return result or {} @distributed_trace_async async def patch_item( - self, - item: Union[str, Dict[str, Any]], - partition_key: PartitionKeyType, - patch_operations: List[Dict[str, Any]], - *, - filter_predicate: Optional[str] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Dict[str, Any]], + partition_key: PartitionKeyType, + patch_operations: List[Dict[str, Any]], + *, + filter_predicate: Optional[str] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """ Patches the specified item with the provided operations if it exists in the container. @@ -905,9 +896,9 @@ async def patch_item( 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 no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the patch operation went through. The dict will be empty if + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after the patch operation went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. @@ -937,24 +928,23 @@ async def patch_item( item_link = self._get_document_link(item) result = await self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, options=request_options, - request_context={}, **kwargs) + document_link=item_link, operations=patch_operations, options=request_options, **kwargs) return result or {} @distributed_trace_async async def delete_item( - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> None: """Delete the specified item from the container. @@ -1001,8 +991,7 @@ async def delete_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) - await self.client_connection.DeleteItem(document_link=document_link, options=request_options, - request_context={}, **kwargs) + await self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) @distributed_trace_async async def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @@ -1036,9 +1025,9 @@ async def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @distributed_trace_async async def replace_throughput( - self, - throughput: Union[int, ThroughputProperties], - **kwargs: Any + self, + throughput: Union[int, ThroughputProperties], + **kwargs: Any ) -> ThroughputProperties: """Replace the container's throughput. @@ -1073,10 +1062,10 @@ async def replace_throughput( @distributed_trace def list_conflicts( - self, - *, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + *, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """List all the conflicts in the container. @@ -1102,13 +1091,13 @@ def list_conflicts( @distributed_trace def query_conflicts( - self, - query: str, - *, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + query: str, + *, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Return all conflicts matching a given `query`. @@ -1147,10 +1136,10 @@ def query_conflicts( @distributed_trace_async async def get_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any, + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any, ) -> Dict[str, Any]: """Get the conflict identified by `conflict`. @@ -1175,10 +1164,10 @@ async def get_conflict( @distributed_trace_async async def delete_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any, + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any, ) -> None: """Delete a specified conflict from the container. @@ -1203,15 +1192,15 @@ async def delete_conflict( @distributed_trace_async async def delete_all_items_by_partition_key( - self, - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - **kwargs: Any + self, + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + **kwargs: Any ) -> None: """The delete by partition key feature is an asynchronous, background operation that allows you to delete all documents with the same logical partition key value, using the Cosmos SDK. The delete by partition key @@ -1251,17 +1240,17 @@ async def delete_all_items_by_partition_key( @distributed_trace_async async def execute_item_batch( - self, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """ Executes the transactional batch for the specified partition key. @@ -1303,8 +1292,7 @@ async def execute_item_batch( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] return await self.client_connection.Batch( - collection_link=self.container_link, batch_operations=batch_operations, options=request_options, - request_context={}, **kwargs) + collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) async def read_feed_ranges( self, @@ -1330,17 +1318,18 @@ async def read_feed_ranges( [Range("", "FF", True, False)], **kwargs) - return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) + return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange), self.container_link) for partitionKeyRange in partition_key_ranges] - async def get_updated_session_token(self, + def get_updated_session_token(self, feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], target_feed_range: FeedRange - ) -> "Session Token": + ) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a response from any crud operation. This should only be used if maintaining own session token or else the sdk will keep track of - session token. + session token. Session tokens and feed ranges are scoped to a container. Only input session tokens + and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple(str, FeedRange)] :param target_feed_range: feed range to get most up to date session token. @@ -1348,7 +1337,7 @@ async def get_updated_session_token(self, :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range, self.container_link) async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. @@ -1357,9 +1346,9 @@ async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) - :returns: a feed range :rtype: FeedRange """ - return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key)) + return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key), self.container_link) - async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. :param parent_feed_range: left feed range :type parent_feed_range: FeedRange @@ -1368,4 +1357,7 @@ async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_ra :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ + if (child_feed_range._container_link != self.container_link or + parent_feed_range._container_link != self.container_link): + raise ValueError("Feed ranges must be from the same container.") return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_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 58dc4eea0dcd..088602860468 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 @@ -326,9 +326,9 @@ def _check_if_account_session_consistency(self, database_account: DatabaseAccoun return None def _GetDatabaseIdWithPathForUser( - self, - database_link: str, - user: Mapping[str, Any] + self, + database_link: str, + user: Mapping[str, Any] ) -> Tuple[Optional[str], str]: base._validate_resource(user) path = base.GetPathFromLink(database_link, "users") @@ -336,9 +336,9 @@ def _GetDatabaseIdWithPathForUser( return database_id, path def _GetContainerIdWithPathForSproc( - self, - collection_link: str, - sproc: Mapping[str, Any] + self, + collection_link: str, + sproc: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(sproc) sproc = dict(sproc) @@ -351,9 +351,9 @@ def _GetContainerIdWithPathForSproc( return collection_id, path, sproc def _GetContainerIdWithPathForTrigger( - self, - collection_link: str, - trigger: Mapping[str, Any] + self, + collection_link: str, + trigger: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(trigger) trigger = dict(trigger) @@ -367,9 +367,9 @@ def _GetContainerIdWithPathForTrigger( return collection_id, path, trigger def _GetContainerIdWithPathForUDF( - self, - collection_link: str, - udf: Mapping[str, Any] + self, + collection_link: str, + udf: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(udf) udf = dict(udf) @@ -383,9 +383,9 @@ def _GetContainerIdWithPathForUDF( return collection_id, path, udf async def GetDatabaseAccount( - self, - url_connection: Optional[str] = None, - **kwargs: Any + self, + url_connection: Optional[str] = None, + **kwargs: Any ) -> documents.DatabaseAccount: """Gets database account info. @@ -433,10 +433,10 @@ async def GetDatabaseAccount( return database_account async def CreateDatabase( - self, - database: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a database. @@ -457,11 +457,11 @@ async def CreateDatabase( return await self.Create(database, path, "dbs", None, None, options, **kwargs) async def CreateUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user. @@ -484,11 +484,11 @@ async def CreateUser( return await self.Create(user, path, "users", database_id, None, options, **kwargs) async def CreateContainer( - self, - database_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ): """Creates a collection in a database. @@ -511,12 +511,11 @@ async def CreateContainer( return await self.Create(collection, path, "colls", database_id, None, options, **kwargs) async def CreateItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -526,7 +525,6 @@ async def CreateItem( The Azure Cosmos document to create. :param dict options: The request options for the request. - :param dict request_context: :return: The created Document. :rtype: @@ -552,15 +550,14 @@ async def CreateItem( if base.IsItemContainerLink(database_or_container_link): options = await self._AddPartitionKey(database_or_container_link, document, options) - return await self.Create(document, path, "docs", collection_id, None, options, - request_context, **kwargs) + return await self.Create(document, path, "docs", collection_id, None, options, **kwargs) async def CreatePermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a permission for a user. @@ -583,11 +580,11 @@ async def CreatePermission( return await self.Create(permission, path, "permissions", user_id, None, options, **kwargs) async def CreateUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user-defined function in a collection. @@ -609,11 +606,11 @@ async def CreateUserDefinedFunction( return await self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) async def CreateTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a trigger in a collection. @@ -635,11 +632,11 @@ async def CreateTrigger( return await self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) async def CreateStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Creates a stored procedure in a collection. @@ -661,11 +658,11 @@ async def CreateStoredProcedure( return await self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) async def ExecuteStoredProcedure( - self, - sproc_link: str, - params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Executes a store procedure. @@ -700,15 +697,14 @@ async def ExecuteStoredProcedure( return result async def Create( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -719,7 +715,6 @@ async def Create( :param dict initial_headers: :param dict options: The request options for the request. - :param dict request_context: :return: The created Azure Cosmos resource. :rtype: @@ -740,18 +735,16 @@ async def Create( # update session for write request self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result async def UpsertUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user. @@ -772,11 +765,11 @@ async def UpsertUser( return await self.Upsert(user, path, "users", database_id, None, options, **kwargs) async def UpsertPermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a permission for a user. @@ -799,12 +792,11 @@ async def UpsertPermission( return await self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) async def UpsertItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a document in a collection. @@ -814,7 +806,6 @@ async def UpsertItem( The Azure Cosmos document to upsert. :param dict options: The request options for the request. - :param dict request_context: :return: The upserted Document. :rtype: @@ -839,19 +830,17 @@ async def UpsertItem( collection_id, document, path = self._GetContainerIdWithPathForItem( database_or_container_link, document, options ) - return await self.Upsert(document, path, "docs", collection_id, None, options, - request_context, **kwargs) + return await self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) async def Upsert( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -862,7 +851,6 @@ async def Upsert( :param dict initial_headers: :param dict options: The request options for the request. - :param dict request_context: :return: The upserted Azure Cosmos resource. :rtype: @@ -884,19 +872,17 @@ async def Upsert( self.last_response_headers = last_response_headers # update session for write request self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result async def __Post( - self, - path: str, - request_params: _request_object.RequestObject, - body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'POST' async http request. @@ -920,10 +906,10 @@ async def __Post( ) async def ReadDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a database. @@ -944,10 +930,10 @@ async def ReadDatabase( return await self.Read(path, "dbs", database_id, None, options, **kwargs) async def ReadContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a collection. @@ -970,11 +956,10 @@ async def ReadContainer( return await self.Read(path, "colls", collection_id, None, options, **kwargs) async def ReadItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a document. @@ -982,7 +967,6 @@ async def ReadItem( The link to the document. :param dict options: The request options for the request. - :param dict request_context: :return: The read Document. @@ -995,13 +979,13 @@ async def ReadItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - return await self.Read(path, "docs", document_id, None, options, request_context, **kwargs) + return await self.Read(path, "docs", document_id, None, options, **kwargs) async def ReadUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user. @@ -1024,10 +1008,10 @@ async def ReadUser( return await self.Read(path, "users", user_id, None, options, **kwargs) async def ReadPermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a permission. @@ -1050,10 +1034,10 @@ async def ReadPermission( return await self.Read(path, "permissions", permission_id, None, options, **kwargs) async def ReadUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user-defined function. @@ -1076,10 +1060,10 @@ async def ReadUserDefinedFunction( return await self.Read(path, "udfs", udf_id, None, options, **kwargs) async def ReadStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a stored procedure. @@ -1101,10 +1085,10 @@ async def ReadStoredProcedure( return await self.Read(path, "sprocs", sproc_id, None, options, **kwargs) async def ReadTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a trigger. @@ -1127,10 +1111,10 @@ async def ReadTrigger( return await self.Read(path, "triggers", trigger_id, None, options, **kwargs) async def ReadConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a conflict. @@ -1152,14 +1136,13 @@ async def ReadConflict( return await self.Read(path, "conflicts", conflict_id, None, options, **kwargs) async def Read( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -1169,7 +1152,7 @@ async def Read( :param dict initial_headers: :param dict options: The request options for the request. - :param dict request_context: + :return: The upserted Azure Cosmos resource. :rtype: @@ -1186,18 +1169,16 @@ async def Read( request_params = _request_object.RequestObject(typ, documents._OperationType.Read) result, last_response_headers = await self.__Get(path, request_params, headers, **kwargs) self.last_response_headers = last_response_headers - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result async def __Get( - self, - path: str, - request_params: _request_object.RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'GET' async http request. @@ -1220,11 +1201,11 @@ async def __Get( ) async def ReplaceUser( - self, - user_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user and return it. @@ -1248,11 +1229,11 @@ async def ReplaceUser( return await self.Replace(user, path, "users", user_id, None, options, **kwargs) async def ReplacePermission( - self, - permission_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a permission and return it. @@ -1276,11 +1257,11 @@ async def ReplacePermission( return await self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) async def ReplaceContainer( - self, - collection_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a collection and return it. @@ -1305,11 +1286,11 @@ async def ReplaceContainer( return await self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) async def ReplaceUserDefinedFunction( - self, - udf_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user-defined function and returns it. @@ -1339,11 +1320,11 @@ async def ReplaceUserDefinedFunction( return await self.Replace(udf, path, "udfs", udf_id, None, options, **kwargs) async def ReplaceTrigger( - self, - trigger_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + trigger_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Replaces a trigger and returns it. @@ -1373,12 +1354,11 @@ async def ReplaceTrigger( return await self.Replace(trigger, path, "triggers", trigger_id, None, options, **kwargs) async def ReplaceItem( - self, - document_link: str, - new_document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + new_document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a document and returns it. @@ -1387,7 +1367,6 @@ async def ReplaceItem( :param dict new_document: :param dict options: The request options for the request. - :param dict request_context: :return: The new Document. :rtype: @@ -1412,23 +1391,20 @@ async def ReplaceItem( collection_link = base.GetItemContainerLink(document_link) options = await self._AddPartitionKey(collection_link, new_document, options) - return await self.Replace(new_document, path, "docs", document_id, None, options, - request_context, **kwargs) + return await self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) async def PatchItem( - self, - document_link: str, - operations: List[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + operations: List[Dict[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Patches a document and returns it. :param str document_link: The link to the document. :param list operations: The operations for the patch request. :param dict options: The request options for the request. - :param dict request_context: The request context for the operation :return: The new Document. :rtype: @@ -1456,17 +1432,15 @@ async def PatchItem( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result async def ReplaceOffer( - self, - offer_link: str, - offer: Dict[str, Any], - **kwargs: Any + self, + offer_link: str, + offer: Dict[str, Any], + **kwargs: Any ) -> Dict[str, Any]: """Replaces an offer and returns it. @@ -1485,11 +1459,11 @@ async def ReplaceOffer( return await self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) async def ReplaceStoredProcedure( - self, - sproc_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a stored procedure and returns it. @@ -1519,15 +1493,14 @@ async def ReplaceStoredProcedure( return await self.Replace(sproc, path, "sprocs", sproc_id, None, options, **kwargs) async def Replace( - self, - resource: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + resource: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -1538,7 +1511,6 @@ async def Replace( :param dict initial_headers: :param dict options: The request options for the request. - :param dict request_context: :return: The new Azure Cosmos resource. :rtype: @@ -1558,19 +1530,17 @@ async def Replace( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, self.last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, result) return result async def __Put( - self, - path: str, - request_params: _request_object.RequestObject, - body: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + body: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PUT' async http request. @@ -1594,12 +1564,12 @@ async def __Put( ) async def __Patch( - self, - path: str, - request_params: _request_object.RequestObject, - request_data: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + request_data: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PATCH' http request. @@ -1623,10 +1593,10 @@ async def __Patch( ) async def DeleteDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a database. @@ -1647,10 +1617,10 @@ async def DeleteDatabase( await self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) async def DeleteUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user. @@ -1672,10 +1642,10 @@ async def DeleteUser( await self.DeleteResource(path, "users", user_id, None, options, **kwargs) async def DeletePermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a permission. @@ -1697,10 +1667,10 @@ async def DeletePermission( await self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) async def DeleteContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a collection. @@ -1722,11 +1692,10 @@ async def DeleteContainer( await self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) async def DeleteItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a document. @@ -1734,7 +1703,6 @@ async def DeleteItem( The link to the document. :param dict options: The request options for the request. - :param dict request_context: :return: The deleted Document. :rtype: @@ -1746,14 +1714,13 @@ async def DeleteItem( path = base.GetPathFromLink(document_link) document_id = base.GetResourceIdOrFullNameFromLink(document_link) - await self.DeleteResource(path, "docs", document_id, None, options, - request_context, **kwargs) + await self.DeleteResource(path, "docs", document_id, None, options, **kwargs) async def DeleteUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user-defined function. @@ -1774,10 +1741,10 @@ async def DeleteUserDefinedFunction( await self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) async def DeleteTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a trigger. @@ -1799,10 +1766,10 @@ async def DeleteTrigger( await self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) async def DeleteStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a stored procedure. @@ -1824,10 +1791,10 @@ async def DeleteStoredProcedure( await self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) async def DeleteConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a conflict. @@ -1849,14 +1816,13 @@ async def DeleteConflict( await self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) async def DeleteResource( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -1866,7 +1832,6 @@ async def DeleteResource( :param dict initial_headers: :param dict options: The request options for the request. - :param dict request_context: :return: The deleted Azure Cosmos resource. :rtype: @@ -1886,17 +1851,15 @@ async def DeleteResource( # update session for request mutates data on server side self._UpdateSessionIfRequired(headers, result, last_response_headers) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, None) async def __Delete( - self, - path: str, - request_params: _request_object.RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[None, Dict[str, Any]]: """Azure Cosmos 'DELETE' async http request. @@ -1919,19 +1882,17 @@ async def __Delete( ) async def Batch( - self, - collection_link: str, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """Executes the given operations in transactional batch. :param str collection_link: The link to the collection :param list batch_operations: The batch of operations for the batch request. :param dict options: The request options for the request. - :param dict request_context: The request context for the operation :return: The result of the batch operation. @@ -1977,22 +1938,20 @@ async def Batch( " index {}. Error message: {}".format( str(error_index), Constants.ERROR_TRANSLATIONS.get(error_status) - ), + ), operation_responses=final_responses ) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(last_response_headers, final_responses) return final_responses async def _Batch( - self, - batch_operations: List[Dict[str, Any]], - path: str, - collection_id: Optional[str], - options: Mapping[str, Any], - **kwargs: Any + self, + batch_operations: List[Dict[str, Any]], + path: str, + collection_id: Optional[str], + options: Mapping[str, Any], + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) @@ -2002,10 +1961,10 @@ async def _Batch( return cast(Tuple[List[Dict[str, Any]], Dict[str, Any]], result) def _ReadPartitionKeyRanges( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads Partition Key Ranges. @@ -2024,11 +1983,11 @@ def _ReadPartitionKeyRanges( return self._QueryPartitionKeyRanges(collection_link, None, feed_options, **kwargs) def _QueryPartitionKeyRanges( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries Partition Key Ranges in a collection. @@ -2063,9 +2022,9 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadDatabases( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all databases. @@ -2083,10 +2042,10 @@ def ReadDatabases( return self.QueryDatabases(None, options, **kwargs) def QueryDatabases( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries databases. @@ -2115,10 +2074,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadContainers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all collections in a database. @@ -2137,11 +2096,11 @@ def ReadContainers( return self.QueryContainers(database_link, None, options, **kwargs) def QueryContainers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries collections in a database. @@ -2175,19 +2134,17 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadItems( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all documents in a collection. :param str collection_link: The link to the document collection. :param dict feed_options: The additional options for the operation. :param response_hook: A callable invoked with the response metadata. - :param dict request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. :rtype: query_iterable.QueryIterable @@ -2196,18 +2153,16 @@ def ReadItems( if feed_options is None: feed_options = {} - return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, - request_context=request_context, **kwargs) + return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) def QueryItems( - self, - database_or_container_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key: Optional[PartitionKeyType] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key: Optional[PartitionKeyType] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -2217,7 +2172,6 @@ def QueryItems( :param dict options: The request options for the request. :param str partition_key: Partition key for the query(default value None) :param response_hook: A callable invoked with the response metadata. - :param dict request_context: The request context for the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2254,7 +2208,6 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di query, options, response_hook=response_hook, - request_context=request_context, **kwargs ), self.last_response_headers, @@ -2270,19 +2223,17 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def QueryItemsChangeFeed( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. :param str collection_link: The link to the document collection. :param dict options: The request options for the request. :param response_hook: A callable invoked with the response metadata. - :param dict request_context: The request_context of the operation :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2296,8 +2247,7 @@ def QueryItemsChangeFeed( partition_key_range_id = options["partitionKeyRangeId"] return self._QueryChangeFeed( - collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, - request_context=request_context, **kwargs + collection_link, "Documents", options, partition_key_range_id, response_hook=response_hook, **kwargs ) def _QueryChangeFeed( @@ -2307,7 +2257,6 @@ def _QueryChangeFeed( options: Optional[Mapping[str, Any]] = None, partition_key_range_id: Optional[str] = None, response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - request_context: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries change feed of a resource in a collection. @@ -2317,7 +2266,6 @@ def _QueryChangeFeed( :param dict options: The request options for the request. :param str partition_key_range_id: Specifies partition key range id. :param response_hook: A callable invoked with the response metadata - :param dict request_context: The request context of the operation. :type response_hook: Callable[[Dict[str, str], Dict[str, Any]] :return: Query Iterable of Documents. @@ -2356,7 +2304,6 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di options, partition_key_range_id, response_hook=response_hook, - request_context=request_context, **kwargs ), self.last_response_headers, @@ -2371,10 +2318,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def QueryOffers( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Query for all offers. @@ -2407,10 +2354,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadUsers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all users in a database. @@ -2430,11 +2377,11 @@ def ReadUsers( return self.QueryUsers(database_link, None, options, **kwargs) def QueryUsers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries users in a database. @@ -2469,10 +2416,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadPermissions( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all permissions for a user. @@ -2492,11 +2439,11 @@ def ReadPermissions( return self.QueryPermissions(user_link, None, options, **kwargs) def QueryPermissions( - self, - user_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + user_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> AsyncItemPaged[Dict[str, Any]]: """Queries permissions for a user. @@ -2530,10 +2477,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadStoredProcedures( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all store procedures in a collection. @@ -2553,11 +2500,11 @@ def ReadStoredProcedures( return self.QueryStoredProcedures(collection_link, None, options, **kwargs) def QueryStoredProcedures( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries stored procedures in a collection. @@ -2592,10 +2539,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadTriggers( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all triggers in a collection. @@ -2615,11 +2562,11 @@ def ReadTriggers( return self.QueryTriggers(collection_link, None, options, **kwargs) def QueryTriggers( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries triggers in a collection. @@ -2653,10 +2600,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadUserDefinedFunctions( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all user-defined functions in a collection. @@ -2676,11 +2623,11 @@ def ReadUserDefinedFunctions( return self.QueryUserDefinedFunctions(collection_link, None, options, **kwargs) def QueryUserDefinedFunctions( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries user-defined functions in a collection. @@ -2715,10 +2662,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadConflicts( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> AsyncItemPaged[Dict[str, Any]]: """Reads conflicts. @@ -2737,11 +2684,11 @@ def ReadConflicts( return self.QueryConflicts(collection_link, None, feed_options, **kwargs) def QueryConflicts( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries conflicts in a collection. @@ -2776,13 +2723,13 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) async def QueryFeed( - self, - path: str, - collection_id: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Mapping[str, Any], - partition_key_range_id: Optional[str] = None, - **kwargs: Any + self, + path: str, + collection_id: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Mapping[str, Any], + partition_key_range_id: Optional[str] = None, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query Feed for Document Collection resource. @@ -2810,19 +2757,18 @@ async def QueryFeed( ) async def __QueryFeed( # pylint: disable=too-many-branches,too-many-statements,too-many-locals - self, - path: str, - typ: str, - id_: Optional[str], - result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], - create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - is_query_plan: bool = False, - request_context: Optional[Dict[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id_: Optional[str], + result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], + create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + is_query_plan: bool = False, + **kwargs: Any ) -> List[Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -2839,7 +2785,6 @@ async def __QueryFeed( # pylint: disable=too-many-branches,too-many-statements, :param function response_hook: :param bool is_query_plan: Specifies if the call is to fetch query plan - :param dict request_context: :returns: A list of the queried resources. :rtype: list :raises SystemError: If the query compatibility mode is undefined. @@ -2873,8 +2818,6 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: await change_feed_state.populate_request_headers_async(self._routing_map_provider, headers) result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, result) return __GetBodiesFromQueryResult(result) @@ -2959,8 +2902,6 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: results["Documents"].extend(partial_result["Documents"]) else: results = partial_result - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, partial_result) # if the prefix partition query has results lets return it @@ -2972,16 +2913,14 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: INDEX_METRICS_HEADER = http_constants.HttpHeaders.IndexUtilization index_metrics_raw = self.last_response_headers[INDEX_METRICS_HEADER] self.last_response_headers[INDEX_METRICS_HEADER] = _utils.get_index_metrics_info(index_metrics_raw) - if request_context is not None: - self._add_request_context(request_context) if response_hook: response_hook(self.last_response_headers, result) return __GetBodiesFromQueryResult(result) def __CheckAndUnifyQueryFormat( - self, - query_body: Union[str, Dict[str, Any]] + self, + query_body: Union[str, Dict[str, Any]] ) -> Union[str, Dict[str, Any]]: """Checks and unifies the format of the query body. @@ -3015,10 +2954,10 @@ def __CheckAndUnifyQueryFormat( return query_body def _UpdateSessionIfRequired( - self, - request_headers: Mapping[str, Any], - response_result: Optional[Mapping[str, Any]], - response_headers: Optional[Mapping[str, Any]] + self, + request_headers: Mapping[str, Any], + response_result: Optional[Mapping[str, Any]], + response_headers: Optional[Mapping[str, Any]] ) -> None: """ Updates session if necessary. @@ -3083,9 +3022,9 @@ def _GetUserIdWithPathForPermission(self, permission, user_link): return path, user_id def RegisterPartitionResolver( - self, - database_link: str, - partition_resolver: RangePartitionResolver + self, + database_link: str, + partition_resolver: RangePartitionResolver ) -> None: """Registers the partition resolver associated with the database link @@ -3235,10 +3174,10 @@ async def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kw ) async def DeleteAllItemsByPartitionKey( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Exposes an API to delete all items with a single partition key without the user having to explicitly call delete on each record in the partition key. @@ -3266,7 +3205,7 @@ async def DeleteAllItemsByPartitionKey( headers = base.GetHeaders(self, initial_headers, "post", path, collection_id, "partitionkey", options) request_params = _request_object.RequestObject("partitionkey", documents._OperationType.Delete) _, last_response_headers = await self.__Post(path=path, request_params=request_params, - req_headers=headers, body=None, **kwargs) + req_headers=headers, body=None, **kwargs) self.last_response_headers = last_response_headers if response_hook: response_hook(last_response_headers, None) @@ -3283,8 +3222,3 @@ async def _get_partition_key_definition(self, collection_link: str) -> Optional[ partition_key_definition = container.get("partitionKey") self.__container_properties_cache[collection_link] = _set_properties_cache(container) return partition_key_definition - - def _add_request_context(self, request_context): - if http_constants.HttpHeaders.SessionToken in self.last_response_headers: - request_context['session_token'] = self.last_response_headers[http_constants.HttpHeaders.SessionToken] - self.last_response_headers["request_context"] = request_context diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index eda6d09b6565..96725d32174a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -76,11 +76,11 @@ class ContainerProxy: # pylint: disable=too-many-public-methods """ def __init__( - self, - client_connection: CosmosClientConnection, - database_link: str, - id: str, - properties: Optional[Dict[str, Any]] = None + self, + client_connection: CosmosClientConnection, + database_link: str, + id: str, + properties: Optional[Dict[str, Any]] = None ) -> None: self.id = id self.container_link = "{}/colls/{}".format(database_link, self.id) @@ -125,8 +125,8 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> return conflict_or_link["_self"] def _set_partition_key( - self, - partition_key: PartitionKeyType + self, + partition_key: PartitionKeyType ) -> Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(self.is_system_key) @@ -144,15 +144,15 @@ def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: @distributed_trace def read( # pylint:disable=docstring-missing-param - self, - populate_query_metrics: Optional[bool] = None, - populate_partition_key_range_statistics: Optional[bool] = None, - populate_quota_info: Optional[bool] = None, - *, - session_token: Optional[str] = None, - priority: Optional[Literal["High", "Low"]] = None, - initial_headers: Optional[Dict[str, str]] = None, - **kwargs: Any + self, + populate_query_metrics: Optional[bool] = None, + populate_partition_key_range_statistics: Optional[bool] = None, + populate_quota_info: Optional[bool] = None, + *, + session_token: Optional[str] = None, + priority: Optional[Literal["High", "Low"]] = None, + initial_headers: Optional[Dict[str, str]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Read the container properties. @@ -194,17 +194,17 @@ def read( # pylint:disable=docstring-missing-param @distributed_trace def read_item( # pylint:disable=docstring-missing-param - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - populate_query_metrics: Optional[bool] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + populate_query_metrics: Optional[bool] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Get the item identified by `item`. @@ -258,20 +258,19 @@ def read_item( # pylint:disable=docstring-missing-param request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] - return self.client_connection.ReadItem(document_link=doc_link, request_context={}, - options=request_options, **kwargs) + return self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs) @distributed_trace def read_all_items( # pylint:disable=docstring-missing-param - self, - max_item_count: Optional[int] = None, - populate_query_metrics: Optional[bool] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + max_item_count: Optional[int] = None, + populate_query_metrics: Optional[bool] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """List all the items in the container. @@ -315,8 +314,7 @@ def read_all_items( # pylint:disable=docstring-missing-param feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] items = self.client_connection.ReadItems( - collection_link=self.container_link, feed_options=feed_options, - response_hook=response_hook, request_context={}, **kwargs) + collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs) if response_hook: response_hook(self.client_connection.last_response_headers, items) return items @@ -522,17 +520,16 @@ def query_items_change_feed( feed_options["maxItemCount"] = kwargs.pop('max_item_count') except KeyError: feed_options["maxItemCount"] = args[3] - request_context = {} + if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] =\ + change_feed_state_context["partitionKey"] = \ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) - change_feed_state_context["partitionKeyFeedRange"] =\ + change_feed_state_context["partitionKeyFeedRange"] = \ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: feed_range: FeedRangeEpk = kwargs.pop('feed_range') change_feed_state_context["feedRange"] = feed_range._feed_range_internal - request_context["feed_range"] = feed_range container_properties = self._get_properties() feed_options["changeFeedStateContext"] = change_feed_state_context @@ -543,11 +540,7 @@ def query_items_change_feed( response_hook.clear() result = self.client_connection.QueryItemsChangeFeed( - self.container_link, - options=feed_options, - response_hook=response_hook, - request_context=request_context, - **kwargs + self.container_link, options=feed_options, response_hook=response_hook, **kwargs ) if response_hook: response_hook(self.client_connection.last_response_headers, result) @@ -555,22 +548,22 @@ def query_items_change_feed( @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, - *, - populate_index_metrics: Optional[bool] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, - **kwargs: Any + 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, + *, + populate_index_metrics: Optional[bool] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + continuation_token_limit: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Return all results matching the given `query`. @@ -671,7 +664,6 @@ def query_items( # pylint:disable=docstring-missing-param options=feed_options, partition_key=partition_key, response_hook=response_hook, - request_context={}, **kwargs ) if response_hook: @@ -679,28 +671,29 @@ def query_items( # pylint:disable=docstring-missing-param return items def __is_prefix_partitionkey( - self, partition_key: PartitionKeyType) -> bool: + self, partition_key: PartitionKeyType) -> bool: properties = self._get_properties() pk_properties = properties["partitionKey"] partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) return partition_key_definition._is_prefix_partition_key(partition_key) + @distributed_trace def replace_item( # pylint:disable=docstring-missing-param - self, - item: Union[str, Mapping[str, Any]], - body: Dict[str, Any], - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + body: Dict[str, Any], + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces the specified item if it exists in the container. @@ -722,9 +715,9 @@ def replace_item( # pylint:disable=docstring-missing-param 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 no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` + sending response payloads. When not specified explicitly here, the default value will be determined from + kwargs or when also not specified there from client-level kwargs. + :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace operation failed or the item with given id does not exist. @@ -758,26 +751,26 @@ def replace_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + result = self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, request_context={}, options=request_options, **kwargs - ) + document_link=item_link, new_document=body, options=request_options, **kwargs) return result or {} @distributed_trace def upsert_item( # pylint:disable=docstring-missing-param - self, - body: Dict[str, Any], - populate_query_metrics: Optional[bool]=None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + populate_query_metrics: Optional[bool]=None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Insert or update the specified item. @@ -797,9 +790,9 @@ def upsert_item( # pylint:disable=docstring-missing-param :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or - when also not specified there from client-level kwargs. + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or + when also not specified there from client-level kwargs. :returns: A dict representing the upserted item. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. :rtype: Dict[str, Any] @@ -830,10 +823,10 @@ def upsert_item( # pylint:disable=docstring-missing-param request_options["populateQueryMetrics"] = populate_query_metrics if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] + result = self.client_connection.UpsertItem( database_or_container_link=self.container_link, document=body, - request_context={}, options=request_options, **kwargs ) @@ -841,21 +834,21 @@ def upsert_item( # pylint:disable=docstring-missing-param @distributed_trace def create_item( # pylint:disable=docstring-missing-param - self, - body: Dict[str, Any], - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - indexing_directive: Optional[int] = None, - *, - enable_automatic_id_generation: bool = False, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + indexing_directive: Optional[int] = None, + *, + enable_automatic_id_generation: bool = False, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Create an item in the container. @@ -879,8 +872,8 @@ def create_item( # pylint:disable=docstring-missing-param :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. @@ -915,26 +908,25 @@ def create_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, - options=request_options, request_context={}, **kwargs) + database_or_container_link=self.container_link, document=body, options=request_options, **kwargs) return result or {} @distributed_trace def patch_item( - self, - item: Union[str, Dict[str, Any]], - partition_key: PartitionKeyType, - patch_operations: List[Dict[str, Any]], - *, - filter_predicate: Optional[str] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Dict[str, Any]], + partition_key: PartitionKeyType, + patch_operations: List[Dict[str, Any]], + *, + filter_predicate: Optional[str] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """ Patches the specified item with the provided operations if it exists in the container. @@ -958,10 +950,10 @@ def patch_item( :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after the patch operations went through. The dict will be empty + :returns: A dict representing the item after the patch operations went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. @@ -991,23 +983,22 @@ def patch_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] item_link = self._get_document_link(item) result = self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, - request_context={}, options=request_options, **kwargs) + document_link=item_link, operations=patch_operations, options=request_options, **kwargs) return result or {} @distributed_trace def execute_item_batch( - self, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """ Executes the transactional batch for the specified partition key. @@ -1045,27 +1036,25 @@ def execute_item_batch( request_options = build_options(kwargs) request_options["partitionKey"] = self._set_partition_key(partition_key) request_options["disableAutomaticIdGeneration"] = True - result = self.client_connection.Batch( - collection_link=self.container_link, batch_operations=batch_operations, - request_context={}, - options=request_options, **kwargs) - return result + + return self.client_connection.Batch( + collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs) @distributed_trace def delete_item( # pylint:disable=docstring-missing-param - self, - item: Union[Mapping[str, Any], str], - partition_key: PartitionKeyType, - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[Mapping[str, Any], str], + partition_key: PartitionKeyType, + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> None: """Delete the specified item from the container. @@ -1116,9 +1105,7 @@ def delete_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] document_link = self._get_document_link(item) - self.client_connection.DeleteItem(document_link=document_link, - request_context={}, - options=request_options, **kwargs) + self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs) @distributed_trace def read_offer(self, **kwargs: Any) -> Offer: @@ -1167,9 +1154,9 @@ def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @distributed_trace def replace_throughput( - self, - throughput: Union[int, ThroughputProperties], - **kwargs: Any + self, + throughput: Union[int, ThroughputProperties], + **kwargs: Any ) -> ThroughputProperties: """Replace the container's throughput. @@ -1201,9 +1188,9 @@ def replace_throughput( @distributed_trace def list_conflicts( - self, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """List all the conflicts in the container. @@ -1228,13 +1215,13 @@ def list_conflicts( @distributed_trace def query_conflicts( - self, - query: str, - parameters: Optional[List[Dict[str, object]]] = None, - enable_cross_partition_query: Optional[bool] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + query: str, + parameters: Optional[List[Dict[str, object]]] = None, + enable_cross_partition_query: Optional[bool] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Return all conflicts matching a given `query`. @@ -1274,10 +1261,10 @@ def query_conflicts( @distributed_trace def get_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any ) -> Dict[str, Any]: """Get the conflict identified by `conflict`. @@ -1302,10 +1289,10 @@ def get_conflict( @distributed_trace def delete_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any ) -> None: """Delete a specified conflict from the container. @@ -1332,15 +1319,15 @@ def delete_conflict( @distributed_trace def delete_all_items_by_partition_key( - self, - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - **kwargs: Any + self, + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + **kwargs: Any ) -> None: """The delete by partition key feature is an asynchronous, background operation that allows you to delete all documents with the same logical partition key value, using the Cosmos SDK. The delete by partition key @@ -1401,17 +1388,18 @@ def read_feed_ranges( [Range("", "FF", True, False)], # default to full range **kwargs) - return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) + return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange), self.container_link) for partitionKeyRange in partition_key_ranges] - def get_updated_session_token(self, - feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], - target_feed_range: FeedRange - ) -> "Session Token": + def get_updated_session_token( + self, + feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], + target_feed_range: FeedRange) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a response from any crud operation. This should only be used if maintaining own session token or else the sdk will keep track of - session token. + session token. Session tokens and feed ranges are scoped to a container. Only input session tokens + and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple(str, Range)] :param target_feed_range: feed range to get most up to date session token. @@ -1419,7 +1407,7 @@ def get_updated_session_token(self, :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range, self.container_link) def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. @@ -1428,7 +1416,7 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Feed :returns: a feed range :rtype: FeedRange """ - return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) + return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key), self.container_link) def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. @@ -1439,5 +1427,8 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ + if (child_feed_range._container_link != self.container_link or + parent_feed_range._container_link != self.container_link): + raise ValueError("Feed ranges must be from the same container.") return child_feed_range._feed_range_internal.get_normalized_range().is_subset( parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index d8d44b6ceb6e..f75b49520459 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -98,12 +98,13 @@ def test_partition_key_to_feed_range(self, setup): @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_subset_ranges) def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset): - epk_parent_feed_range = FeedRangeEpk(parent_feed_range) - epk_child_feed_range = FeedRangeEpk(child_feed_range) + epk_parent_feed_range = FeedRangeEpk(parent_feed_range, setup["created_collection"].container_link) + epk_child_feed_range = FeedRangeEpk(child_feed_range, setup["created_collection"].container_link) assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) == is_subset def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) @@ -111,6 +112,21 @@ def test_feed_range_is_subset_from_pk(self, setup): def test_overlaps(self, setup, range1, range2, overlaps): assert Range.overlaps(range1, range2) == overlaps + def test_is_subset_with_wrong_feed_range(self, setup): + wrong_container = "wrong_container_link" + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + wrong_container) + epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) + with pytest.raises(ValueError, match="Feed ranges must be from the same container."): + setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) + epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), + wrong_container) + with pytest.raises(ValueError, match="Feed ranges must be from the same container."): + setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py new file mode 100644 index 000000000000..231db3eac598 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py @@ -0,0 +1,78 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import unittest +import uuid + +import pytest +import pytest_asyncio + +import azure.cosmos.partition_key as partition_key +import test_config +from azure.cosmos._feed_range import FeedRangeEpk +from azure.cosmos._routing.routing_range import Range +from azure.cosmos.aio import CosmosClient + + +@pytest_asyncio.fixture() +def setup(): + if (TestFeedRangeAsync.masterKey == '[YOUR_KEY_HERE]' or + TestFeedRangeAsync.host == '[YOUR_ENDPOINT_HERE]'): + raise Exception( + "You must specify your Azure Cosmos account values for " + "'masterKey' and 'host' at the top of this class to run the " + "tests.") + test_client = CosmosClient(TestFeedRangeAsync.host, test_config.TestConfig.masterKey), + created_db = test_client[0].get_database_client(TestFeedRangeAsync.TEST_DATABASE_ID) + return { + "created_db": created_db, + "created_collection": created_db.get_container_client(TestFeedRangeAsync.TEST_CONTAINER_ID) + } + +@pytest.mark.cosmosEmulator +@pytest.mark.asyncio +@pytest.mark.usefixtures("setup") +class TestFeedRangeAsync: + """Tests to verify methods for operations on feed ranges + """ + + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID + TEST_CONTAINER_ID = test_config.TestConfig.TEST_MULTI_PARTITION_CONTAINER_ID + + + async def test_partition_key_to_feed_range(self, setup): + created_container = await setup["created_db"].create_container( + id='container_' + str(uuid.uuid4()), + partition_key=partition_key.PartitionKey(path="/id") + ) + feed_range = await created_container.feed_range_from_partition_key("1") + assert feed_range._feed_range_internal.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", + "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) + await setup["created_db"].delete_container(created_container) + + async def test_feed_range_is_subset_from_pk(self, setup): + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) + epk_child_feed_range = await setup["created_collection"].feed_range_from_partition_key("1") + assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + + def test_is_subset_with_wrong_feed_range(self, setup): + wrong_container = "wrong_container_link" + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + wrong_container) + epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) + with pytest.raises(ValueError, match="Feed ranges must be from the same container."): + setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), + setup["created_collection"].container_link) + epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), + wrong_container) + with pytest.raises(ValueError, match="Feed ranges must be from the same container."): + setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) + +if __name__ == '__main__': + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context.py b/sdk/cosmos/azure-cosmos/test/test_request_context.py deleted file mode 100644 index 1c46152dc0c3..000000000000 --- a/sdk/cosmos/azure-cosmos/test/test_request_context.py +++ /dev/null @@ -1,97 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) Microsoft Corporation. All rights reserved. -import time -import unittest -import uuid - -import pytest - -import azure.cosmos.cosmos_client as cosmos_client -import test_config - -@pytest.fixture(scope="class") -def setup(): - if (TestRequestContext.masterKey == '[YOUR_KEY_HERE]' or - TestRequestContext.host == '[YOUR_ENDPOINT_HERE]'): - raise Exception( - "You must specify your Azure Cosmos account values for " - "'masterKey' and 'host' at the top of this class to run the " - "tests.") - test_client = cosmos_client.CosmosClient(TestRequestContext.host, test_config.TestConfig.masterKey), - created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) - return { - "created_db": created_db, - "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) - } - -def validate_request_context(collection): - request_context = collection.client_connection.last_response_headers["request_context"] - keys_expected = ["session_token"] - assert request_context is not None - for key in keys_expected: - assert request_context[key] is not None - -def createItem(id = None, pk='A', name='sample'): - if id is None: - id = 'item' + str(uuid.uuid4()) - item = { - 'id': id, - 'name': name, - 'pk': pk - } - return item - -@pytest.mark.cosmosEmulator -@pytest.mark.unittest -@pytest.mark.usefixtures("setup") -class TestRequestContext: - """Tests to verify request context gets populated correctly - """ - - host = test_config.TestConfig.host - masterKey = test_config.TestConfig.masterKey - TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID - TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID - - def test_crud_request_context(self, setup): - item = createItem() - setup["created_collection"].create_item(item) - validate_request_context(setup["created_collection"]) - - setup["created_collection"].read_item(item['id'], item['pk']) - validate_request_context(setup["created_collection"]) - - new_item = createItem(item['id'], name='sample_replaced') - setup["created_collection"].replace_item(item['id'], new_item) - validate_request_context(setup["created_collection"]) - operations = [ - {"op": "add", "path": "/favorite_color", "value": "red"}, - {"op": "replace", "path": "/name", "value": 14}, - ] - setup["created_collection"].patch_item(item['id'], item['pk'], operations) - validate_request_context(setup["created_collection"]) - - items = list(setup["created_collection"].read_all_items()) - assert len(items) > 0 - validate_request_context(setup["created_collection"]) - - setup["created_collection"].upsert_item(createItem()) - validate_request_context(setup["created_collection"]) - - for i in range(100): - setup["created_collection"].create_item(createItem()) - items = list(setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)) - assert len(items) > 100 - validate_request_context(setup["created_collection"]) - - items = list(setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", - parameters=[dict(name="@id", value=item['id'])], - partition_key=item['pk'])) - assert len(items) == 1 - validate_request_context(setup["created_collection"]) - - setup["created_collection"].delete_item(item['id'], item['pk']) - validate_request_context(setup["created_collection"]) - -if __name__ == '__main__': - unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py b/sdk/cosmos/azure-cosmos/test/test_request_context_async.py deleted file mode 100644 index 5d772ed09eca..000000000000 --- a/sdk/cosmos/azure-cosmos/test/test_request_context_async.py +++ /dev/null @@ -1,99 +0,0 @@ -# The MIT License (MIT) -# Copyright (c) Microsoft Corporation. All rights reserved. - -import unittest -import uuid - -import pytest -import pytest_asyncio - -import test_config -from azure.cosmos.aio import CosmosClient - -@pytest_asyncio.fixture() -async def setup(): - if (TestRequestContext.masterKey == '[YOUR_KEY_HERE]' or - TestRequestContext.host == '[YOUR_ENDPOINT_HERE]'): - raise Exception( - "You must specify your Azure Cosmos account values for " - "'masterKey' and 'host' at the top of this class to run the " - "tests.") - test_client = CosmosClient(TestRequestContext.host, test_config.TestConfig.masterKey), - created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID) - return { - "created_db": created_db, - "created_collection": created_db.get_container_client(TestRequestContext.TEST_CONTAINER_ID) - } - -def validate_request_context(collection): - request_context = collection.client_connection.last_response_headers["request_context"] - keys_expected = ["session_token"] - assert request_context is not None - for key in keys_expected: - assert request_context[key] is not None - -def createItem(id = None, pk='A', name='sample'): - if id is None: - id = 'item' + str(uuid.uuid4()) - item = { - 'id': id, - 'name': name, - 'pk': pk - } - return item - - -@pytest.mark.cosmosEmulator -@pytest.mark.asyncio -@pytest.mark.usefixtures("setup") -class TestRequestContext: - """Tests to verify request context gets populated correctly - """ - - host = test_config.TestConfig.host - masterKey = test_config.TestConfig.masterKey - TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID - TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID - - async def test_crud_request_context(self, setup): - item = createItem() - await setup["created_collection"].create_item(item) - validate_request_context(setup["created_collection"]) - - await setup["created_collection"].read_item(item['id'], item['pk']) - validate_request_context(setup["created_collection"]) - - new_item = createItem(item['id'], name='sample_replaced') - await setup["created_collection"].replace_item(item['id'], new_item) - validate_request_context(setup["created_collection"]) - operations = [ - {"op": "add", "path": "/favorite_color", "value": "red"}, - {"op": "replace", "path": "/name", "value": 14}, - ] - await setup["created_collection"].patch_item(item['id'], item['pk'], operations) - validate_request_context(setup["created_collection"]) - - items = [item async for item in setup["created_collection"].read_all_items()] - assert len(items) > 0 - validate_request_context(setup["created_collection"]) - - await setup["created_collection"].upsert_item(createItem()) - validate_request_context(setup["created_collection"]) - - for i in range(100): - await setup["created_collection"].create_item(createItem()) - items = [item async for item in setup["created_collection"].query_items_change_feed(is_start_from_beginning=True)] - assert len(items) > 100 - validate_request_context(setup["created_collection"]) - - items = [item async for item in setup["created_collection"].query_items("SELECT * FROM c WHERE c.id = @id", - parameters=[dict(name="@id", value=item['id'])], - partition_key=item['pk'])] - assert len(items) == 1 - validate_request_context(setup["created_collection"]) - - await setup["created_collection"].delete_item(item['id'], item['pk']) - validate_request_context(setup["created_collection"]) - -if __name__ == '__main__': - unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 114a60831ee5..778de06f2c3a 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -2,9 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. import random -import time import unittest -import uuid import pytest @@ -32,8 +30,7 @@ def setup(): } def create_split_ranges(): - # add one with several ranges being equal to one - test_params = [ # split with two children + return [ # split with two children ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), # split with one child @@ -47,6 +44,10 @@ def create_split_ranges(): ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=52"), + # several ranges being equal to one range + ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), + (("BB", "CC"),"1:1#66#3=52"), (("CC", "DD"),"1:1#55#3=52")], + ("AA", "DD"), "0:1#60#3=52,1:1#66#3=52"), # merge with one child ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], ("AA", "DD"), "0:1#55#3=52"), @@ -66,15 +67,6 @@ def create_split_ranges(): ([(("AA", "BB"), "0:1#54#3=52"), (("AA", "BB"),"0:2#57#3=53")], ("AA", "BB"), "0:2#57#3=53") ] - actual_test_params = [] - for test_param in test_params: - split_ranges = [] - for feed_range, session_token in test_param[0]: - split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], - True, False)), session_token)) - target_feed_range = FeedRangeEpk(Range(test_param[1][0], test_param[1][1], True, False)) - actual_test_params.append((split_ranges, target_feed_range, test_param[2])) - return actual_test_params @pytest.mark.cosmosEmulator @pytest.mark.unittest @@ -91,7 +83,8 @@ class TestSessionTokenHelpers: TEST_COLLECTION_ID = configs.TEST_SINGLE_PARTITION_CONTAINER_ID def test_get_session_token_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False), + setup[COLLECTION].container_link) session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] session_token = "0:1#51#3=52" @@ -100,7 +93,8 @@ def test_get_session_token_update(self, setup): assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False), + setup[COLLECTION].container_link) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) @@ -112,15 +106,18 @@ def test_many_session_tokens_update_same_range(self, setup): assert updated_session_token == session_token def test_many_session_tokens_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False), + setup[COLLECTION].container_link) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) feed_ranges_and_session_tokens.append((feed_range, session_token)) # adding irrelevant feed ranges - feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False)) - feed_range2 = FeedRangeEpk(Range("00", "55", True, False)) + feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False), + setup[COLLECTION].container_link) + feed_range2 = FeedRangeEpk(Range("00", "55", True, False), + setup[COLLECTION].container_link) for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) if i % 2 == 0: @@ -135,16 +132,30 @@ def test_many_session_tokens_update(self, setup): @pytest.mark.parametrize("split_ranges, target_feed_range, expected_session_token", create_split_ranges()) def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, expected_session_token): - updated_session_token = setup[COLLECTION].get_updated_session_token(split_ranges, target_feed_range) + actual_split_ranges = [] + for feed_range, session_token in split_ranges: + actual_split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], + True, False), + setup[COLLECTION].container_link), session_token)) + target_feed_range = FeedRangeEpk(Range(target_feed_range[0], target_feed_range[1][1], + True, False), + setup[COLLECTION].container_link) + updated_session_token = setup[COLLECTION].get_updated_session_token(actual_split_ranges, target_feed_range) assert updated_session_token == expected_session_token def test_invalid_feed_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False), + setup[COLLECTION].container_link) session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, - FeedRangeEpk(Range("CC", "FF", True, False))) + FeedRangeEpk(Range( + "CC", + "FF", + True, + False), + setup[COLLECTION].container_link)) if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py index 849d00be0d1c..6e8397bdccd4 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py @@ -11,6 +11,7 @@ import test_config from azure.cosmos import DatabaseProxy, PartitionKey from azure.cosmos._session_token_helpers import is_compound_session_token, parse_session_token +from azure.cosmos.http_constants import HttpHeaders class TestUpdatedSessionTokenSplit(unittest.TestCase): @@ -120,12 +121,12 @@ def create_items_logical_pk(container, target_pk, previous_session_token, feed_r 'pk': 'A' + str(random.randint(1, 10)) } container.create_item(item, session_token=previous_session_token) - request_context = container.client_connection.last_response_headers["request_context"] + session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] if item['pk'] == target_pk: - target_session_token = request_context["session_token"] - previous_session_token = request_context["session_token"] + target_session_token = session_token + previous_session_token = session_token feed_ranges_and_session_tokens.append((container.feed_range_from_partition_key(item['pk']), - request_context["session_token"])) + session_token)) return target_session_token, previous_session_token @staticmethod @@ -145,12 +146,12 @@ def create_items_physical_pk(container, pk_feed_range, previous_session_token, f 'pk': 'A' + str(random.randint(1, 10)) } container.create_item(item, session_token=previous_session_token) - request_context = container.client_connection.last_response_headers["request_context"] + session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] curr_feed_range = container.feed_range_from_partition_key(item['pk']) if container.is_feed_range_subset(target_feed_range, curr_feed_range): - target_session_token = request_context["session_token"] - previous_session_token = request_context["session_token"] - feed_ranges_and_session_tokens.append((curr_feed_range, request_context["session_token"])) + target_session_token = session_token + previous_session_token = session_token + feed_ranges_and_session_tokens.append((curr_feed_range, session_token)) return target_session_token, target_feed_range, previous_session_token From 5552912757c0237724a32a330180071b742d1002 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 9 Oct 2024 23:21:58 -0700 Subject: [PATCH 37/59] Reacting to comments --- .../azure/cosmos/_cosmos_client_connection.py | 956 +++++++++--------- .../aio/_cosmos_client_connection_async.py | 802 +++++++-------- 2 files changed, 879 insertions(+), 879 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 cf62e92f6f2c..aa0241d7f289 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -108,12 +108,12 @@ class _QueryCompatibilityMode: _DefaultStringRangePrecision = -1 def __init__( - self, - url_connection: str, - auth: CredentialDict, - connection_policy: Optional[ConnectionPolicy] = None, - consistency_level: Optional[str] = None, - **kwargs: Any + self, + url_connection: str, + auth: CredentialDict, + connection_policy: Optional[ConnectionPolicy] = None, + consistency_level: Optional[str] = None, + **kwargs: Any ) -> None: """ :param str url_connection: @@ -255,9 +255,9 @@ def _set_container_properties_cache(self, container_link: str, properties: Optio self.__container_properties_cache[container_link] = {} def _set_client_consistency_level( - self, - database_account: DatabaseAccount, - consistency_level: Optional[str], + self, + database_account: DatabaseAccount, + consistency_level: Optional[str], ) -> None: """Checks if consistency level param was passed in by user and sets it to that value or to the account default. @@ -348,10 +348,10 @@ def GetPartitionResolver(self, database_link: str) -> Optional[RangePartitionRes return self.partition_resolvers.get(base.TrimBeginningAndEndingSlashes(database_link)) def CreateDatabase( - self, - database: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a database. @@ -371,10 +371,10 @@ def CreateDatabase( return self.Create(database, path, "dbs", None, None, options, **kwargs) def ReadDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a database. @@ -396,9 +396,9 @@ def ReadDatabase( return self.Read(path, "dbs", database_id, None, options, **kwargs) def ReadDatabases( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all databases. @@ -417,10 +417,10 @@ def ReadDatabases( return self.QueryDatabases(None, options, **kwargs) def QueryDatabases( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries databases. @@ -438,18 +438,18 @@ def QueryDatabases( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - "/dbs", "dbs", "", lambda r: r["Databases"], - lambda _, b: b, query, options, **kwargs) + "/dbs", "dbs", "", lambda r: r["Databases"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def ReadContainers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all collections in a database. @@ -469,11 +469,11 @@ def ReadContainers( return self.QueryContainers(database_link, None, options, **kwargs) def QueryContainers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries collections in a database. @@ -496,19 +496,19 @@ def QueryContainers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "colls", database_id, lambda r: r["DocumentCollections"], - lambda _, body: body, query, options, **kwargs) + path, "colls", database_id, lambda r: r["DocumentCollections"], + lambda _, body: body, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateContainer( - self, - database_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a collection in a database. @@ -532,11 +532,11 @@ def CreateContainer( return self.Create(collection, path, "colls", database_id, None, options, **kwargs) def ReplaceContainer( - self, - collection_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a collection and return it. @@ -562,10 +562,10 @@ def ReplaceContainer( return self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) def ReadContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a collection. @@ -588,11 +588,11 @@ def ReadContainer( return self.Read(path, "colls", collection_id, None, options, **kwargs) def CreateUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user. @@ -616,11 +616,11 @@ def CreateUser( return self.Create(user, path, "users", database_id, None, options, **kwargs) def UpsertUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user. @@ -648,10 +648,10 @@ def _GetDatabaseIdWithPathForUser(self, database_link: str, user: Mapping[str, A return database_id, path def ReadUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user. @@ -674,10 +674,10 @@ def ReadUser( return self.Read(path, "users", user_id, None, options, **kwargs) def ReadUsers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all users in a database. @@ -697,11 +697,11 @@ def ReadUsers( return self.QueryUsers(database_link, None, options, **kwargs) def QueryUsers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries users in a database. @@ -725,18 +725,18 @@ def QueryUsers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "users", database_id, lambda r: r["Users"], - lambda _, b: b, query, options, **kwargs) + path, "users", database_id, lambda r: r["Users"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def DeleteDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a database. @@ -759,11 +759,11 @@ def DeleteDatabase( self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) def CreatePermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a permission for a user. @@ -787,11 +787,11 @@ def CreatePermission( return self.Create(permission, path, "permissions", user_id, None, options, **kwargs) def UpsertPermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a permission for a user. @@ -815,9 +815,9 @@ def UpsertPermission( return self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) def _GetUserIdWithPathForPermission( - self, - permission: Mapping[str, Any], - user_link: str + self, + permission: Mapping[str, Any], + user_link: str ) -> Tuple[str, Optional[str]]: base._validate_resource(permission) path = base.GetPathFromLink(user_link, "permissions") @@ -825,10 +825,10 @@ def _GetUserIdWithPathForPermission( return path, user_id def ReadPermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a permission. @@ -851,10 +851,10 @@ def ReadPermission( return self.Read(path, "permissions", permission_id, None, options, **kwargs) def ReadPermissions( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all permissions for a user. @@ -875,11 +875,11 @@ def ReadPermissions( return self.QueryPermissions(user_link, None, options, **kwargs) def QueryPermissions( - self, - user_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries permissions for a user. @@ -903,19 +903,19 @@ def QueryPermissions( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "permissions", user_id, lambda r: r["Permissions"], - lambda _, b: b, query, options, **kwargs) + path, "permissions", user_id, lambda r: r["Permissions"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def ReplaceUser( - self, - user_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user and return it. @@ -940,10 +940,10 @@ def ReplaceUser( return self.Replace(user, path, "users", user_id, None, options, **kwargs) def DeleteUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user. @@ -966,11 +966,11 @@ def DeleteUser( self.DeleteResource(path, "users", user_id, None, options, **kwargs) def ReplacePermission( - self, - permission_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a permission and return it. @@ -995,10 +995,10 @@ def ReplacePermission( return self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) def DeletePermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a permission. @@ -1021,11 +1021,11 @@ def DeletePermission( self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) def ReadItems( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all documents in a collection. :param str collection_link: The link to the document collection. @@ -1041,13 +1041,13 @@ def ReadItems( return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) def QueryItems( - self, - database_or_container_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key: Optional[PartitionKeyType] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + database_or_container_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key: Optional[PartitionKeyType] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -1086,15 +1086,15 @@ def QueryItems( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - response_hook=response_hook, - **kwargs) + path, + "docs", + collection_id, + lambda r: r["Documents"], + lambda _, b: b, + query, + options, + response_hook=response_hook, + **kwargs) return ItemPaged( self, @@ -1106,11 +1106,11 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def QueryItemsChangeFeed( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. @@ -1135,13 +1135,13 @@ def QueryItemsChangeFeed( ) def _QueryChangeFeed( - self, - collection_link: str, - resource_type: str, - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + collection_link: str, + resource_type: str, + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries change feed of a resource in a collection. @@ -1199,10 +1199,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def _ReadPartitionKeyRanges( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads Partition Key Ranges. @@ -1218,11 +1218,11 @@ def _ReadPartitionKeyRanges( return self._QueryPartitionKeyRanges(collection_link, None, feed_options, **kwargs) def _QueryPartitionKeyRanges( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries Partition Key Ranges in a collection. @@ -1246,19 +1246,19 @@ def _QueryPartitionKeyRanges( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], - lambda _, b: b, query, options, **kwargs) + path, "pkranges", collection_id, lambda r: r["PartitionKeyRanges"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -1290,11 +1290,11 @@ def CreateItem( return self.Create(document, path, "docs", collection_id, None, options, **kwargs) def UpsertItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a document in a collection. @@ -1334,10 +1334,10 @@ def UpsertItem( # Gets the collection id and path for the document def _GetContainerIdWithPathForItem( - self, - database_or_container_link: str, - document: Mapping[str, Any], - options: Mapping[str, Any] + self, + database_or_container_link: str, + document: Mapping[str, Any], + options: Mapping[str, Any] ) -> Tuple[Optional[str], Dict[str, Any], str]: if not database_or_container_link: @@ -1366,10 +1366,10 @@ def _GetContainerIdWithPathForItem( return collection_id, document, path def ReadItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Reads a document. @@ -1392,10 +1392,10 @@ def ReadItem( return self.Read(path, "docs", document_id, None, options, **kwargs) def ReadTriggers( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all triggers in a collection. @@ -1416,11 +1416,11 @@ def ReadTriggers( return self.QueryTriggers(collection_link, None, options, **kwargs) def QueryTriggers( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries triggers in a collection. @@ -1452,11 +1452,11 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def CreateTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a trigger in a collection. @@ -1479,9 +1479,9 @@ def CreateTrigger( return self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForTrigger( - self, - collection_link: str, - trigger: Mapping[str, Any] + self, + collection_link: str, + trigger: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(trigger) trigger = dict(trigger) @@ -1495,10 +1495,10 @@ def _GetContainerIdWithPathForTrigger( return collection_id, path, trigger def ReadTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a trigger. @@ -1521,11 +1521,11 @@ def ReadTrigger( return self.Read(path, "triggers", trigger_id, None, options, **kwargs) def UpsertTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a trigger in a collection. :param str collection_link: @@ -1545,10 +1545,10 @@ def UpsertTrigger( return self.Upsert(trigger, path, "triggers", collection_id, None, options, **kwargs) def ReadUserDefinedFunctions( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all user-defined functions in a collection. @@ -1569,11 +1569,11 @@ def ReadUserDefinedFunctions( return self.QueryUserDefinedFunctions(collection_link, None, options, **kwargs) def QueryUserDefinedFunctions( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries user-defined functions in a collection. @@ -1597,19 +1597,19 @@ def QueryUserDefinedFunctions( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], - lambda _, b: b, query, options, **kwargs) + path, "udfs", collection_id, lambda r: r["UserDefinedFunctions"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def CreateUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user-defined function in a collection. @@ -1632,11 +1632,11 @@ def CreateUserDefinedFunction( return self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) def UpsertUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user-defined function in a collection. @@ -1659,9 +1659,9 @@ def UpsertUserDefinedFunction( return self.Upsert(udf, path, "udfs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForUDF( - self, - collection_link: str, - udf: Mapping[str, Any] + self, + collection_link: str, + udf: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(udf) udf = dict(udf) @@ -1675,10 +1675,10 @@ def _GetContainerIdWithPathForUDF( return collection_id, path, udf def ReadUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user-defined function. @@ -1701,10 +1701,10 @@ def ReadUserDefinedFunction( return self.Read(path, "udfs", udf_id, None, options, **kwargs) def ReadStoredProcedures( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all store procedures in a collection. @@ -1725,11 +1725,11 @@ def ReadStoredProcedures( return self.QueryStoredProcedures(collection_link, None, options, **kwargs) def QueryStoredProcedures( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries stored procedures in a collection. @@ -1761,11 +1761,11 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def CreateStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a stored procedure in a collection. @@ -1788,11 +1788,11 @@ def CreateStoredProcedure( return self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) def UpsertStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a stored procedure in a collection. @@ -1815,9 +1815,9 @@ def UpsertStoredProcedure( return self.Upsert(sproc, path, "sprocs", collection_id, None, options, **kwargs) def _GetContainerIdWithPathForSproc( - self, - collection_link: str, - sproc: Mapping[str, Any] + self, + collection_link: str, + sproc: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(sproc) sproc = dict(sproc) @@ -1830,10 +1830,10 @@ def _GetContainerIdWithPathForSproc( return collection_id, path, sproc def ReadStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a stored procedure. @@ -1856,10 +1856,10 @@ def ReadStoredProcedure( return self.Read(path, "sprocs", sproc_id, None, options, **kwargs) def ReadConflicts( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads conflicts. @@ -1879,11 +1879,11 @@ def ReadConflicts( return self.QueryConflicts(collection_link, None, feed_options, **kwargs) def QueryConflicts( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Queries conflicts in a collection. @@ -1915,10 +1915,10 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str ) def ReadConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a conflict. @@ -1940,10 +1940,10 @@ def ReadConflict( return self.Read(path, "conflicts", conflict_id, None, options, **kwargs) def DeleteContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a collection. @@ -1966,11 +1966,11 @@ def DeleteContainer( self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) def ReplaceItem( - self, - document_link: str, - new_document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + new_document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a document and returns it. @@ -2007,11 +2007,11 @@ def ReplaceItem( return self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) def PatchItem( - self, - document_link: str, - operations: List[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + operations: List[Dict[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Patches a document and returns it. @@ -2049,11 +2049,11 @@ def PatchItem( return result def Batch( - self, - collection_link: str, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """Executes the given operations in transactional batch. @@ -2111,12 +2111,12 @@ def Batch( return final_responses def _Batch( - self, - batch_operations: List[Dict[str, Any]], - path: str, - collection_id: Optional[str], - options: Mapping[str, Any], - **kwargs: Any + self, + batch_operations: List[Dict[str, Any]], + path: str, + collection_id: Optional[str], + options: Mapping[str, Any], + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) @@ -2128,10 +2128,10 @@ def _Batch( ) def DeleteItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a document. @@ -2154,10 +2154,10 @@ def DeleteItem( self.DeleteResource(path, "docs", document_id, None, options, **kwargs) def DeleteAllItemsByPartitionKey( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Exposes an API to delete all items with a single partition key without the user having to explicitly call delete on each record in the partition key. @@ -2193,11 +2193,11 @@ def DeleteAllItemsByPartitionKey( response_hook(last_response_headers, None) def ReplaceTrigger( - self, - trigger_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a trigger and returns it. @@ -2228,10 +2228,10 @@ def ReplaceTrigger( return self.Replace(trigger, path, "triggers", trigger_id, None, options, **kwargs) def DeleteTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a trigger. @@ -2254,11 +2254,11 @@ def DeleteTrigger( self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) def ReplaceUserDefinedFunction( - self, - udf_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user-defined function and returns it. @@ -2289,10 +2289,10 @@ def ReplaceUserDefinedFunction( return self.Replace(udf, path, "udfs", udf_id, None, options, **kwargs) def DeleteUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user-defined function. @@ -2315,11 +2315,11 @@ def DeleteUserDefinedFunction( self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) def ExecuteStoredProcedure( - self, - sproc_link: str, - params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Executes a store procedure. @@ -2355,11 +2355,11 @@ def ExecuteStoredProcedure( return result def ReplaceStoredProcedure( - self, - sproc_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a stored procedure and returns it. @@ -2390,10 +2390,10 @@ def ReplaceStoredProcedure( return self.Replace(sproc, path, "sprocs", sproc_id, None, options, **kwargs) def DeleteStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a stored procedure. @@ -2416,10 +2416,10 @@ def DeleteStoredProcedure( self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) def DeleteConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a conflict. @@ -2442,10 +2442,10 @@ def DeleteConflict( self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) def ReplaceOffer( - self, - offer_link: str, - offer: Dict[str, Any], - **kwargs: Any + self, + offer_link: str, + offer: Dict[str, Any], + **kwargs: Any ) -> Dict[str, Any]: """Replaces an offer and returns it. @@ -2465,9 +2465,9 @@ def ReplaceOffer( return self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) def ReadOffer( - self, - offer_link: str, - **kwargs: Any + self, + offer_link: str, + **kwargs: Any ) -> Dict[str, Any]: """Reads an offer. :param str offer_link: @@ -2482,9 +2482,9 @@ def ReadOffer( return self.Read(path, "offers", offer_id, None, {}, **kwargs) def ReadOffers( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Reads all offers. :param dict options: @@ -2500,10 +2500,10 @@ def ReadOffers( return self.QueryOffers(None, options, **kwargs) def QueryOffers( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Query for all offers. @@ -2522,17 +2522,17 @@ def QueryOffers( def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: return self.__QueryFeed( - "/offers", "offers", "", lambda r: r["Offers"], - lambda _, b: b, query, options, **kwargs) + "/offers", "offers", "", lambda r: r["Offers"], + lambda _, b: b, query, options, **kwargs) return ItemPaged( self, query, options, fetch_function=fetch_fn, page_iterator_class=query_iterable.QueryIterable ) def GetDatabaseAccount( - self, - url_connection: Optional[str] = None, - **kwargs: Any + self, + url_connection: Optional[str] = None, + **kwargs: Any ) -> DatabaseAccount: """Gets database account info. @@ -2579,14 +2579,14 @@ def GetDatabaseAccount( return database_account def Create( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -2623,14 +2623,14 @@ def Create( return result def Upsert( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -2667,14 +2667,14 @@ def Upsert( return result def Replace( - self, - resource: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + resource: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -2710,13 +2710,13 @@ def Replace( return result def Read( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -2748,13 +2748,13 @@ def Read( return result def DeleteResource( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -2788,11 +2788,11 @@ def DeleteResource( response_hook(last_response_headers, None) def __Get( - self, - path: str, - request_params: RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'GET' http request. @@ -2815,12 +2815,12 @@ def __Get( ) def __Post( - self, - path: str, - request_params: RequestObject, - body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'POST' http request. @@ -2844,12 +2844,12 @@ def __Post( ) def __Put( - self, - path: str, - request_params: RequestObject, - body: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + body: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PUT' http request. @@ -2873,12 +2873,12 @@ def __Put( ) def __Patch( - self, - path: str, - request_params: RequestObject, - request_data: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + request_data: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PATCH' http request. @@ -2902,11 +2902,11 @@ def __Patch( ) def __Delete( - self, - path: str, - request_params: RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[None, Dict[str, Any]]: """Azure Cosmos 'DELETE' http request. @@ -2929,13 +2929,13 @@ def __Delete( ) def QueryFeed( - self, - path: str, - collection_id: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Mapping[str, Any], - partition_key_range_id: Optional[str] = None, - **kwargs: Any + self, + path: str, + collection_id: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Mapping[str, Any], + partition_key_range_id: Optional[str] = None, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query Feed for Document Collection resource. @@ -2948,29 +2948,29 @@ def QueryFeed( :rtype: tuple of (dict, dict) """ return self.__QueryFeed( - path, - "docs", - collection_id, - lambda r: r["Documents"], - lambda _, b: b, - query, - options, - partition_key_range_id, - **kwargs) + path, + "docs", + collection_id, + lambda r: r["Documents"], + lambda _, b: b, + query, + options, + partition_key_range_id, + **kwargs) def __QueryFeed( # pylint: disable=too-many-locals, too-many-statements - self, - path: str, - resource_type: str, - resource_id: Optional[str], - result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], - create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - is_query_plan: bool = False, - **kwargs: Any + self, + path: str, + resource_type: str, + resource_id: Optional[str], + result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], + create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + is_query_plan: bool = False, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -3205,10 +3205,10 @@ def __CheckAndUnifyQueryFormat(self, query_body: Union[str, Dict[str, Any]]) -> # Adds the partition key to options def _AddPartitionKey( - self, - collection_link: str, - document: Mapping[str, Any], - options: Mapping[str, Any] + self, + collection_link: str, + document: Mapping[str, Any], + options: Mapping[str, Any] ) -> Dict[str, Any]: collection_link = base.TrimBeginningAndEndingSlashes(collection_link) partitionKeyDefinition = self._get_partition_key_definition(collection_link) @@ -3223,9 +3223,9 @@ def _AddPartitionKey( # Extracts the partition key from the document using the partitionKey definition def _ExtractPartitionKey( - self, - partitionKeyDefinition: Mapping[str, Any], - document: Mapping[str, Any] + self, + partitionKeyDefinition: Mapping[str, Any], + document: Mapping[str, Any] ) -> Union[List[Optional[Union[str, float, bool]]], str, float, bool, _Empty, _Undefined]: if partitionKeyDefinition["kind"] == "MultiHash": ret: List[Optional[Union[str, float, bool]]] = [] @@ -3252,10 +3252,10 @@ def _ExtractPartitionKey( # Navigates the document to retrieve the partitionKey specified in the partition key parts def _retrieve_partition_key( - self, - partition_key_parts: List[str], - document: Mapping[str, Any], - is_system_key: bool + self, + partition_key_parts: List[str], + document: Mapping[str, Any], + is_system_key: bool ) -> Union[str, float, bool, _Empty, _Undefined]: expected_matchCount = len(partition_key_parts) matchCount = 0 @@ -3287,10 +3287,10 @@ def _refresh_container_properties_cache(self, container_link: str): self._set_container_properties_cache(container_link, _set_properties_cache(container)) def _UpdateSessionIfRequired( - self, - request_headers: Mapping[str, Any], - response_result: Optional[Mapping[str, Any]], - response_headers: Optional[Mapping[str, Any]] + self, + request_headers: Mapping[str, Any], + response_result: Optional[Mapping[str, Any]], + response_headers: Optional[Mapping[str, Any]] ) -> None: """ Updates session if necessary. 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 088602860468..eeb67225660a 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 @@ -326,9 +326,9 @@ def _check_if_account_session_consistency(self, database_account: DatabaseAccoun return None def _GetDatabaseIdWithPathForUser( - self, - database_link: str, - user: Mapping[str, Any] + self, + database_link: str, + user: Mapping[str, Any] ) -> Tuple[Optional[str], str]: base._validate_resource(user) path = base.GetPathFromLink(database_link, "users") @@ -336,9 +336,9 @@ def _GetDatabaseIdWithPathForUser( return database_id, path def _GetContainerIdWithPathForSproc( - self, - collection_link: str, - sproc: Mapping[str, Any] + self, + collection_link: str, + sproc: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(sproc) sproc = dict(sproc) @@ -351,9 +351,9 @@ def _GetContainerIdWithPathForSproc( return collection_id, path, sproc def _GetContainerIdWithPathForTrigger( - self, - collection_link: str, - trigger: Mapping[str, Any] + self, + collection_link: str, + trigger: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(trigger) trigger = dict(trigger) @@ -367,9 +367,9 @@ def _GetContainerIdWithPathForTrigger( return collection_id, path, trigger def _GetContainerIdWithPathForUDF( - self, - collection_link: str, - udf: Mapping[str, Any] + self, + collection_link: str, + udf: Mapping[str, Any] ) -> Tuple[Optional[str], str, Dict[str, Any]]: base._validate_resource(udf) udf = dict(udf) @@ -383,9 +383,9 @@ def _GetContainerIdWithPathForUDF( return collection_id, path, udf async def GetDatabaseAccount( - self, - url_connection: Optional[str] = None, - **kwargs: Any + self, + url_connection: Optional[str] = None, + **kwargs: Any ) -> documents.DatabaseAccount: """Gets database account info. @@ -433,10 +433,10 @@ async def GetDatabaseAccount( return database_account async def CreateDatabase( - self, - database: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a database. @@ -457,11 +457,11 @@ async def CreateDatabase( return await self.Create(database, path, "dbs", None, None, options, **kwargs) async def CreateUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user. @@ -484,11 +484,11 @@ async def CreateUser( return await self.Create(user, path, "users", database_id, None, options, **kwargs) async def CreateContainer( - self, - database_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ): """Creates a collection in a database. @@ -511,11 +511,11 @@ async def CreateContainer( return await self.Create(collection, path, "colls", database_id, None, options, **kwargs) async def CreateItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a document in a collection. @@ -553,11 +553,11 @@ async def CreateItem( return await self.Create(document, path, "docs", collection_id, None, options, **kwargs) async def CreatePermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a permission for a user. @@ -580,11 +580,11 @@ async def CreatePermission( return await self.Create(permission, path, "permissions", user_id, None, options, **kwargs) async def CreateUserDefinedFunction( - self, - collection_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a user-defined function in a collection. @@ -606,11 +606,11 @@ async def CreateUserDefinedFunction( return await self.Create(udf, path, "udfs", collection_id, None, options, **kwargs) async def CreateTrigger( - self, - collection_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates a trigger in a collection. @@ -632,11 +632,11 @@ async def CreateTrigger( return await self.Create(trigger, path, "triggers", collection_id, None, options, **kwargs) async def CreateStoredProcedure( - self, - collection_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + collection_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Creates a stored procedure in a collection. @@ -658,11 +658,11 @@ async def CreateStoredProcedure( return await self.Create(sproc, path, "sprocs", collection_id, None, options, **kwargs) async def ExecuteStoredProcedure( - self, - sproc_link: str, - params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + params: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Executes a store procedure. @@ -697,14 +697,14 @@ async def ExecuteStoredProcedure( return result async def Create( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Creates an Azure Cosmos resource and returns it. @@ -740,11 +740,11 @@ async def Create( return result async def UpsertUser( - self, - database_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a user. @@ -765,11 +765,11 @@ async def UpsertUser( return await self.Upsert(user, path, "users", database_id, None, options, **kwargs) async def UpsertPermission( - self, - user_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a permission for a user. @@ -792,11 +792,11 @@ async def UpsertPermission( return await self.Upsert(permission, path, "permissions", user_id, None, options, **kwargs) async def UpsertItem( - self, - database_or_container_link: str, - document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_or_container_link: str, + document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts a document in a collection. @@ -833,14 +833,14 @@ async def UpsertItem( return await self.Upsert(document, path, "docs", collection_id, None, options, **kwargs) async def Upsert( - self, - body: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + body: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Upserts an Azure Cosmos resource and returns it. @@ -877,12 +877,12 @@ async def Upsert( return result async def __Post( - self, - path: str, - request_params: _request_object.RequestObject, - body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + body: Optional[Union[str, List[Dict[str, Any]], Dict[str, Any]]], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'POST' async http request. @@ -906,10 +906,10 @@ async def __Post( ) async def ReadDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a database. @@ -930,10 +930,10 @@ async def ReadDatabase( return await self.Read(path, "dbs", database_id, None, options, **kwargs) async def ReadContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a collection. @@ -956,10 +956,10 @@ async def ReadContainer( return await self.Read(path, "colls", collection_id, None, options, **kwargs) async def ReadItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a document. @@ -982,10 +982,10 @@ async def ReadItem( return await self.Read(path, "docs", document_id, None, options, **kwargs) async def ReadUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user. @@ -1008,10 +1008,10 @@ async def ReadUser( return await self.Read(path, "users", user_id, None, options, **kwargs) async def ReadPermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a permission. @@ -1034,10 +1034,10 @@ async def ReadPermission( return await self.Read(path, "permissions", permission_id, None, options, **kwargs) async def ReadUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a user-defined function. @@ -1060,10 +1060,10 @@ async def ReadUserDefinedFunction( return await self.Read(path, "udfs", udf_id, None, options, **kwargs) async def ReadStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a stored procedure. @@ -1085,10 +1085,10 @@ async def ReadStoredProcedure( return await self.Read(path, "sprocs", sproc_id, None, options, **kwargs) async def ReadTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a trigger. @@ -1111,10 +1111,10 @@ async def ReadTrigger( return await self.Read(path, "triggers", trigger_id, None, options, **kwargs) async def ReadConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads a conflict. @@ -1136,13 +1136,13 @@ async def ReadConflict( return await self.Read(path, "conflicts", conflict_id, None, options, **kwargs) async def Read( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Reads an Azure Cosmos resource and returns it. @@ -1174,11 +1174,11 @@ async def Read( return result async def __Get( - self, - path: str, - request_params: _request_object.RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'GET' async http request. @@ -1201,11 +1201,11 @@ async def __Get( ) async def ReplaceUser( - self, - user_link: str, - user: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + user: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user and return it. @@ -1229,11 +1229,11 @@ async def ReplaceUser( return await self.Replace(user, path, "users", user_id, None, options, **kwargs) async def ReplacePermission( - self, - permission_link: str, - permission: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + permission: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a permission and return it. @@ -1257,11 +1257,11 @@ async def ReplacePermission( return await self.Replace(permission, path, "permissions", permission_id, None, options, **kwargs) async def ReplaceContainer( - self, - collection_link: str, - collection: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + collection: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a collection and return it. @@ -1286,11 +1286,11 @@ async def ReplaceContainer( return await self.Replace(collection, path, "colls", collection_id, None, options, **kwargs) async def ReplaceUserDefinedFunction( - self, - udf_link: str, - udf: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + udf: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a user-defined function and returns it. @@ -1320,11 +1320,11 @@ async def ReplaceUserDefinedFunction( return await self.Replace(udf, path, "udfs", udf_id, None, options, **kwargs) async def ReplaceTrigger( - self, - trigger_link: str, - trigger: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + trigger_link: str, + trigger: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> Dict[str, Any]: """Replaces a trigger and returns it. @@ -1354,11 +1354,11 @@ async def ReplaceTrigger( return await self.Replace(trigger, path, "triggers", trigger_id, None, options, **kwargs) async def ReplaceItem( - self, - document_link: str, - new_document: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + new_document: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a document and returns it. @@ -1394,11 +1394,11 @@ async def ReplaceItem( return await self.Replace(new_document, path, "docs", document_id, None, options, **kwargs) async def PatchItem( - self, - document_link: str, - operations: List[Dict[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + operations: List[Dict[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Patches a document and returns it. @@ -1437,10 +1437,10 @@ async def PatchItem( return result async def ReplaceOffer( - self, - offer_link: str, - offer: Dict[str, Any], - **kwargs: Any + self, + offer_link: str, + offer: Dict[str, Any], + **kwargs: Any ) -> Dict[str, Any]: """Replaces an offer and returns it. @@ -1459,11 +1459,11 @@ async def ReplaceOffer( return await self.Replace(offer, path, "offers", offer_id, None, None, **kwargs) async def ReplaceStoredProcedure( - self, - sproc_link: str, - sproc: Dict[str, Any], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + sproc: Dict[str, Any], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces a stored procedure and returns it. @@ -1493,14 +1493,14 @@ async def ReplaceStoredProcedure( return await self.Replace(sproc, path, "sprocs", sproc_id, None, options, **kwargs) async def Replace( - self, - resource: Dict[str, Any], - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + resource: Dict[str, Any], + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces an Azure Cosmos resource and returns it. @@ -1535,12 +1535,12 @@ async def Replace( return result async def __Put( - self, - path: str, - request_params: _request_object.RequestObject, - body: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + body: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PUT' async http request. @@ -1564,12 +1564,12 @@ async def __Put( ) async def __Patch( - self, - path: str, - request_params: _request_object.RequestObject, - request_data: Dict[str, Any], - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + request_data: Dict[str, Any], + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[Dict[str, Any], Dict[str, Any]]: """Azure Cosmos 'PATCH' http request. @@ -1593,10 +1593,10 @@ async def __Patch( ) async def DeleteDatabase( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a database. @@ -1617,10 +1617,10 @@ async def DeleteDatabase( await self.DeleteResource(path, "dbs", database_id, None, options, **kwargs) async def DeleteUser( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user. @@ -1642,10 +1642,10 @@ async def DeleteUser( await self.DeleteResource(path, "users", user_id, None, options, **kwargs) async def DeletePermission( - self, - permission_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + permission_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a permission. @@ -1667,10 +1667,10 @@ async def DeletePermission( await self.DeleteResource(path, "permissions", permission_id, None, options, **kwargs) async def DeleteContainer( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a collection. @@ -1692,10 +1692,10 @@ async def DeleteContainer( await self.DeleteResource(path, "colls", collection_id, None, options, **kwargs) async def DeleteItem( - self, - document_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + document_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a document. @@ -1717,10 +1717,10 @@ async def DeleteItem( await self.DeleteResource(path, "docs", document_id, None, options, **kwargs) async def DeleteUserDefinedFunction( - self, - udf_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + udf_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a user-defined function. @@ -1741,10 +1741,10 @@ async def DeleteUserDefinedFunction( await self.DeleteResource(path, "udfs", udf_id, None, options, **kwargs) async def DeleteTrigger( - self, - trigger_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + trigger_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a trigger. @@ -1766,10 +1766,10 @@ async def DeleteTrigger( await self.DeleteResource(path, "triggers", trigger_id, None, options, **kwargs) async def DeleteStoredProcedure( - self, - sproc_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + sproc_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a stored procedure. @@ -1791,10 +1791,10 @@ async def DeleteStoredProcedure( await self.DeleteResource(path, "sprocs", sproc_id, None, options, **kwargs) async def DeleteConflict( - self, - conflict_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + conflict_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes a conflict. @@ -1816,13 +1816,13 @@ async def DeleteConflict( await self.DeleteResource(path, "conflicts", conflict_id, None, options, **kwargs) async def DeleteResource( - self, - path: str, - typ: str, - id: Optional[str], - initial_headers: Optional[Mapping[str, Any]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + path: str, + typ: str, + id: Optional[str], + initial_headers: Optional[Mapping[str, Any]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Deletes an Azure Cosmos resource and returns it. @@ -1855,11 +1855,11 @@ async def DeleteResource( response_hook(last_response_headers, None) async def __Delete( - self, - path: str, - request_params: _request_object.RequestObject, - req_headers: Dict[str, Any], - **kwargs: Any + self, + path: str, + request_params: _request_object.RequestObject, + req_headers: Dict[str, Any], + **kwargs: Any ) -> Tuple[None, Dict[str, Any]]: """Azure Cosmos 'DELETE' async http request. @@ -1882,11 +1882,11 @@ async def __Delete( ) async def Batch( - self, - collection_link: str, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """Executes the given operations in transactional batch. @@ -1938,7 +1938,7 @@ async def Batch( " index {}. Error message: {}".format( str(error_index), Constants.ERROR_TRANSLATIONS.get(error_status) - ), + ), operation_responses=final_responses ) if response_hook: @@ -1946,12 +1946,12 @@ async def Batch( return final_responses async def _Batch( - self, - batch_operations: List[Dict[str, Any]], - path: str, - collection_id: Optional[str], - options: Mapping[str, Any], - **kwargs: Any + self, + batch_operations: List[Dict[str, Any]], + path: str, + collection_id: Optional[str], + options: Mapping[str, Any], + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: initial_headers = self.default_headers.copy() base._populate_batch_headers(initial_headers) @@ -1961,10 +1961,10 @@ async def _Batch( return cast(Tuple[List[Dict[str, Any]], Dict[str, Any]], result) def _ReadPartitionKeyRanges( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads Partition Key Ranges. @@ -1983,11 +1983,11 @@ def _ReadPartitionKeyRanges( return self._QueryPartitionKeyRanges(collection_link, None, feed_options, **kwargs) def _QueryPartitionKeyRanges( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries Partition Key Ranges in a collection. @@ -2022,9 +2022,9 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadDatabases( - self, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all databases. @@ -2042,10 +2042,10 @@ def ReadDatabases( return self.QueryDatabases(None, options, **kwargs) def QueryDatabases( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries databases. @@ -2074,10 +2074,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadContainers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all collections in a database. @@ -2096,11 +2096,11 @@ def ReadContainers( return self.QueryContainers(database_link, None, options, **kwargs) def QueryContainers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries collections in a database. @@ -2134,11 +2134,11 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadItems( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all documents in a collection. @@ -2156,13 +2156,13 @@ def ReadItems( return self.QueryItems(collection_link, None, feed_options, response_hook=response_hook, **kwargs) def QueryItems( - self, - database_or_container_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key: Optional[PartitionKeyType] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + database_or_container_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key: Optional[PartitionKeyType] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents in a collection. @@ -2223,11 +2223,11 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def QueryItemsChangeFeed( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries documents change feed in a collection. @@ -2318,10 +2318,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def QueryOffers( - self, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Query for all offers. @@ -2354,10 +2354,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadUsers( - self, - database_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all users in a database. @@ -2377,11 +2377,11 @@ def ReadUsers( return self.QueryUsers(database_link, None, options, **kwargs) def QueryUsers( - self, - database_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + database_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries users in a database. @@ -2416,10 +2416,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadPermissions( - self, - user_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + user_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all permissions for a user. @@ -2439,11 +2439,11 @@ def ReadPermissions( return self.QueryPermissions(user_link, None, options, **kwargs) def QueryPermissions( - self, - user_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + user_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> AsyncItemPaged[Dict[str, Any]]: """Queries permissions for a user. @@ -2477,10 +2477,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadStoredProcedures( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all store procedures in a collection. @@ -2500,11 +2500,11 @@ def ReadStoredProcedures( return self.QueryStoredProcedures(collection_link, None, options, **kwargs) def QueryStoredProcedures( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries stored procedures in a collection. @@ -2539,10 +2539,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadTriggers( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all triggers in a collection. @@ -2562,11 +2562,11 @@ def ReadTriggers( return self.QueryTriggers(collection_link, None, options, **kwargs) def QueryTriggers( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries triggers in a collection. @@ -2600,10 +2600,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadUserDefinedFunctions( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Reads all user-defined functions in a collection. @@ -2623,11 +2623,11 @@ def ReadUserDefinedFunctions( return self.QueryUserDefinedFunctions(collection_link, None, options, **kwargs) def QueryUserDefinedFunctions( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries user-defined functions in a collection. @@ -2662,10 +2662,10 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) def ReadConflicts( - self, - collection_link: str, - feed_options: Optional[Mapping[str, Any]] = None, - **kwargs + self, + collection_link: str, + feed_options: Optional[Mapping[str, Any]] = None, + **kwargs ) -> AsyncItemPaged[Dict[str, Any]]: """Reads conflicts. @@ -2684,11 +2684,11 @@ def ReadConflicts( return self.QueryConflicts(collection_link, None, feed_options, **kwargs) def QueryConflicts( - self, - collection_link: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Queries conflicts in a collection. @@ -2723,13 +2723,13 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Di ) async def QueryFeed( - self, - path: str, - collection_id: str, - query: Optional[Union[str, Dict[str, Any]]], - options: Mapping[str, Any], - partition_key_range_id: Optional[str] = None, - **kwargs: Any + self, + path: str, + collection_id: str, + query: Optional[Union[str, Dict[str, Any]]], + options: Mapping[str, Any], + partition_key_range_id: Optional[str] = None, + **kwargs: Any ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Query Feed for Document Collection resource. @@ -2757,18 +2757,18 @@ async def QueryFeed( ) async def __QueryFeed( # pylint: disable=too-many-branches,too-many-statements,too-many-locals - self, - path: str, - typ: str, - id_: Optional[str], - result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], - create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], - query: Optional[Union[str, Dict[str, Any]]], - options: Optional[Mapping[str, Any]] = None, - partition_key_range_id: Optional[str] = None, - response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, - is_query_plan: bool = False, - **kwargs: Any + self, + path: str, + typ: str, + id_: Optional[str], + result_fn: Callable[[Dict[str, Any]], List[Dict[str, Any]]], + create_fn: Optional[Callable[['CosmosClientConnection', Dict[str, Any]], Dict[str, Any]]], + query: Optional[Union[str, Dict[str, Any]]], + options: Optional[Mapping[str, Any]] = None, + partition_key_range_id: Optional[str] = None, + response_hook: Optional[Callable[[Mapping[str, Any], Mapping[str, Any]], None]] = None, + is_query_plan: bool = False, + **kwargs: Any ) -> List[Dict[str, Any]]: """Query for more than one Azure Cosmos resources. @@ -2919,8 +2919,8 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]: return __GetBodiesFromQueryResult(result) def __CheckAndUnifyQueryFormat( - self, - query_body: Union[str, Dict[str, Any]] + self, + query_body: Union[str, Dict[str, Any]] ) -> Union[str, Dict[str, Any]]: """Checks and unifies the format of the query body. @@ -2954,10 +2954,10 @@ def __CheckAndUnifyQueryFormat( return query_body def _UpdateSessionIfRequired( - self, - request_headers: Mapping[str, Any], - response_result: Optional[Mapping[str, Any]], - response_headers: Optional[Mapping[str, Any]] + self, + request_headers: Mapping[str, Any], + response_result: Optional[Mapping[str, Any]], + response_headers: Optional[Mapping[str, Any]] ) -> None: """ Updates session if necessary. @@ -3022,9 +3022,9 @@ def _GetUserIdWithPathForPermission(self, permission, user_link): return path, user_id def RegisterPartitionResolver( - self, - database_link: str, - partition_resolver: RangePartitionResolver + self, + database_link: str, + partition_resolver: RangePartitionResolver ) -> None: """Registers the partition resolver associated with the database link @@ -3174,10 +3174,10 @@ async def _GetQueryPlanThroughGateway(self, query: str, resource_link: str, **kw ) async def DeleteAllItemsByPartitionKey( - self, - collection_link: str, - options: Optional[Mapping[str, Any]] = None, - **kwargs: Any + self, + collection_link: str, + options: Optional[Mapping[str, Any]] = None, + **kwargs: Any ) -> None: """Exposes an API to delete all items with a single partition key without the user having to explicitly call delete on each record in the partition key. @@ -3205,7 +3205,7 @@ async def DeleteAllItemsByPartitionKey( headers = base.GetHeaders(self, initial_headers, "post", path, collection_id, "partitionkey", options) request_params = _request_object.RequestObject("partitionkey", documents._OperationType.Delete) _, last_response_headers = await self.__Post(path=path, request_params=request_params, - req_headers=headers, body=None, **kwargs) + req_headers=headers, body=None, **kwargs) self.last_response_headers = last_response_headers if response_hook: response_hook(last_response_headers, None) From 1bbbd0f8f6140bfe7e52e1a69135a8cb75a5454b Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 10 Oct 2024 13:41:35 -0700 Subject: [PATCH 38/59] pylint and added hpk tests --- .../azure/cosmos/_session_token_helpers.py | 14 +- .../azure/cosmos/aio/_container.py | 337 ++++++++-------- .../azure-cosmos/azure/cosmos/container.py | 370 +++++++++--------- .../test/test_session_token_helpers.py | 6 +- .../test/test_updated_session_token_split.py | 92 ++++- 5 files changed, 438 insertions(+), 381 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 5c94a9f3e094..ad07de732eb4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -22,8 +22,8 @@ """Internal Helper functions for manipulating session tokens. """ from azure.cosmos._routing.routing_range import Range -from ._feed_range import FeedRange from azure.cosmos._vector_session_token import VectorSessionToken +from ._feed_range import FeedRange # pylint: disable=protected-access @@ -55,7 +55,7 @@ def split_compound_session_tokens(compound_session_tokens: [(Range, str)]) -> [s session_tokens.append(session_token) return session_tokens -def merge_session_tokens_for_same_physical_pk(session_tokens: [str]) -> [str]: +def merge_session_tokens_for_same_partition(session_tokens: [str]) -> [str]: i = 0 while i < len(session_tokens): j = i + 1 @@ -75,7 +75,7 @@ def merge_session_tokens_for_same_physical_pk(session_tokens: [str]) -> [str]: return session_tokens -def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: +def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: # pylint: disable=too-many-nested-blocks processed_ranges = [] while len(overlapping_ranges) != 0: feed_range_cmp, session_token_cmp = overlapping_ranges[0] @@ -191,14 +191,14 @@ def get_updated_session_token(feed_ranges_to_session_tokens: [(FeedRange, str)], if len(remaining_session_tokens) == 1: return remaining_session_tokens[0] # merging any session tokens with same physical partition key range id - remaining_session_tokens = merge_session_tokens_for_same_physical_pk(remaining_session_tokens) + remaining_session_tokens = merge_session_tokens_for_same_partition(remaining_session_tokens) updated_session_token = "" # compound the remaining session tokens - for i in range(len(remaining_session_tokens)): + for i, remaining_session_token in enumerate(remaining_session_tokens): if i == len(remaining_session_tokens) - 1: - updated_session_token += remaining_session_tokens[i] + updated_session_token += remaining_session_token else: - updated_session_token += remaining_session_tokens[i] + "," + updated_session_token += remaining_session_token + "," return updated_session_token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index c973ebd9ae95..6e159cd0454c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -77,11 +77,11 @@ class ContainerProxy: """ def __init__( - self, - client_connection: CosmosClientConnection, - database_link: str, - id: str, - properties: Optional[Dict[str, Any]] = None + self, + client_connection: CosmosClientConnection, + database_link: str, + id: str, + properties: Optional[Dict[str, Any]] = None ) -> None: self.client_connection = client_connection self.id = id @@ -130,8 +130,8 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> return conflict_or_link["_self"] async def _set_partition_key( - self, - partition_key: PartitionKeyType + self, + partition_key: PartitionKeyType ) -> Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(await self.is_system_key) @@ -147,14 +147,14 @@ async def _get_epk_range_for_partition_key(self, partition_key_value: PartitionK @distributed_trace_async async def read( - self, - *, - populate_partition_key_range_statistics: Optional[bool] = None, - populate_quota_info: Optional[bool] = None, - session_token: Optional[str] = None, - priority: Optional[Literal["High", "Low"]] = None, - initial_headers: Optional[Dict[str, str]] = None, - **kwargs: Any + self, + *, + populate_partition_key_range_statistics: Optional[bool] = None, + populate_quota_info: Optional[bool] = None, + session_token: Optional[str] = None, + priority: Optional[Literal["High", "Low"]] = None, + initial_headers: Optional[Dict[str, str]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Read the container properties. @@ -191,20 +191,20 @@ async def read( @distributed_trace_async async def create_item( - self, - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - indexing_directive: Optional[int] = None, - enable_automatic_id_generation: bool = False, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + indexing_directive: Optional[int] = None, + enable_automatic_id_generation: bool = False, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Create an item in the container. @@ -231,8 +231,8 @@ async def create_item( before high priority requests start getting throttled. Feature must first be enabled at the account level. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -266,16 +266,16 @@ async def create_item( @distributed_trace_async async def read_item( - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - *, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + *, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Get the item identified by `item`. @@ -330,14 +330,14 @@ async def read_item( @distributed_trace def read_all_items( - self, - *, - max_item_count: Optional[int] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + *, + max_item_count: Optional[int] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """List all the items in the container. @@ -382,21 +382,21 @@ def read_all_items( @distributed_trace def query_items( - self, - query: str, - *, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: 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, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, - **kwargs: Any + self, + query: str, + *, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: 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, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + continuation_token_limit: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Return all results matching the given `query`. @@ -686,7 +686,7 @@ def query_items_change_feed( # pylint: disable=unused-argument feed_options["maxItemCount"] = kwargs.pop('max_item_count') if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = \ + change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get("partition_key"))) change_feed_state_context["partitionKeyFeedRange"] = \ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) @@ -715,18 +715,18 @@ def query_items_change_feed( # pylint: disable=unused-argument @distributed_trace_async async def upsert_item( - self, - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Insert or update the specified item. @@ -749,9 +749,9 @@ async def upsert_item( before high priority requests start getting throttled. Feature must first be enabled at the account level. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the upsert operation went through. The dict will be empty if + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after the upsert operation went through. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -786,19 +786,19 @@ async def upsert_item( @distributed_trace_async async def replace_item( - self, - item: Union[str, Mapping[str, Any]], - body: Dict[str, Any], - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + body: Dict[str, Any], + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces the specified item if it exists in the container. @@ -823,9 +823,9 @@ async def replace_item( :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace failed or the item with given id does not exist. :keyword bool no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` is specified. :rtype: Dict[str, Any] """ @@ -858,20 +858,20 @@ async def replace_item( @distributed_trace_async async def patch_item( - self, - item: Union[str, Dict[str, Any]], - partition_key: PartitionKeyType, - patch_operations: List[Dict[str, Any]], - *, - filter_predicate: Optional[str] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Dict[str, Any]], + partition_key: PartitionKeyType, + patch_operations: List[Dict[str, Any]], + *, + filter_predicate: Optional[str] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """ Patches the specified item with the provided operations if it exists in the container. @@ -896,9 +896,9 @@ async def patch_item( 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 no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - client-level options. - :returns: A dict representing the item after the patch operation went through. The dict will be empty if + sending response payloads. When not specified explicitly here, the default value will be determined from + client-level options. + :returns: A dict representing the item after the patch operation went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. @@ -933,18 +933,18 @@ async def patch_item( @distributed_trace_async async def delete_item( - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> None: """Delete the specified item from the container. @@ -1025,9 +1025,9 @@ async def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @distributed_trace_async async def replace_throughput( - self, - throughput: Union[int, ThroughputProperties], - **kwargs: Any + self, + throughput: Union[int, ThroughputProperties], + **kwargs: Any ) -> ThroughputProperties: """Replace the container's throughput. @@ -1062,10 +1062,10 @@ async def replace_throughput( @distributed_trace def list_conflicts( - self, - *, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + *, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """List all the conflicts in the container. @@ -1091,13 +1091,13 @@ def list_conflicts( @distributed_trace def query_conflicts( - self, - query: str, - *, - parameters: Optional[List[Dict[str, object]]] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + query: str, + *, + parameters: Optional[List[Dict[str, object]]] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> AsyncItemPaged[Dict[str, Any]]: """Return all conflicts matching a given `query`. @@ -1136,10 +1136,10 @@ def query_conflicts( @distributed_trace_async async def get_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any, + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any, ) -> Dict[str, Any]: """Get the conflict identified by `conflict`. @@ -1164,10 +1164,10 @@ async def get_conflict( @distributed_trace_async async def delete_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any, + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any, ) -> None: """Delete a specified conflict from the container. @@ -1192,15 +1192,15 @@ async def delete_conflict( @distributed_trace_async async def delete_all_items_by_partition_key( - self, - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - **kwargs: Any + self, + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + **kwargs: Any ) -> None: """The delete by partition key feature is an asynchronous, background operation that allows you to delete all documents with the same logical partition key value, using the Cosmos SDK. The delete by partition key @@ -1240,17 +1240,17 @@ async def delete_all_items_by_partition_key( @distributed_trace_async async def execute_item_batch( - self, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """ Executes the transactional batch for the specified partition key. @@ -1360,4 +1360,5 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F if (child_feed_range._container_link != self.container_link or parent_feed_range._container_link != self.container_link): raise ValueError("Feed ranges must be from the same container.") - return child_feed_range._feed_range_internal.get_normalized_range().is_subset(parent_feed_range._feed_range_internal.get_normalized_range()) + return child_feed_range._feed_range_internal.get_normalized_range().is_subset( + parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 96725d32174a..a2699dfc8b08 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -76,11 +76,11 @@ class ContainerProxy: # pylint: disable=too-many-public-methods """ def __init__( - self, - client_connection: CosmosClientConnection, - database_link: str, - id: str, - properties: Optional[Dict[str, Any]] = None + self, + client_connection: CosmosClientConnection, + database_link: str, + id: str, + properties: Optional[Dict[str, Any]] = None ) -> None: self.id = id self.container_link = "{}/colls/{}".format(database_link, self.id) @@ -125,8 +125,8 @@ def _get_conflict_link(self, conflict_or_link: Union[str, Mapping[str, Any]]) -> return conflict_or_link["_self"] def _set_partition_key( - self, - partition_key: PartitionKeyType + self, + partition_key: PartitionKeyType ) -> Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined]: if partition_key == NonePartitionKeyValue: return _return_undefined_or_empty_partition_key(self.is_system_key) @@ -144,15 +144,15 @@ def __get_client_container_caches(self) -> Dict[str, Dict[str, Any]]: @distributed_trace def read( # pylint:disable=docstring-missing-param - self, - populate_query_metrics: Optional[bool] = None, - populate_partition_key_range_statistics: Optional[bool] = None, - populate_quota_info: Optional[bool] = None, - *, - session_token: Optional[str] = None, - priority: Optional[Literal["High", "Low"]] = None, - initial_headers: Optional[Dict[str, str]] = None, - **kwargs: Any + self, + populate_query_metrics: Optional[bool] = None, + populate_partition_key_range_statistics: Optional[bool] = None, + populate_quota_info: Optional[bool] = None, + *, + session_token: Optional[str] = None, + priority: Optional[Literal["High", "Low"]] = None, + initial_headers: Optional[Dict[str, str]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Read the container properties. @@ -194,17 +194,17 @@ def read( # pylint:disable=docstring-missing-param @distributed_trace def read_item( # pylint:disable=docstring-missing-param - self, - item: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - populate_query_metrics: Optional[bool] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + populate_query_metrics: Optional[bool] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> Dict[str, Any]: """Get the item identified by `item`. @@ -262,15 +262,15 @@ def read_item( # pylint:disable=docstring-missing-param @distributed_trace def read_all_items( # pylint:disable=docstring-missing-param - self, - max_item_count: Optional[int] = None, - populate_query_metrics: Optional[bool] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + max_item_count: Optional[int] = None, + populate_query_metrics: Optional[bool] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """List all the items in the container. @@ -522,9 +522,9 @@ def query_items_change_feed( feed_options["maxItemCount"] = args[3] if kwargs.get("partition_key") is not None: - change_feed_state_context["partitionKey"] = \ + change_feed_state_context["partitionKey"] =\ self._set_partition_key(cast(PartitionKeyType, kwargs.get('partition_key'))) - change_feed_state_context["partitionKeyFeedRange"] = \ + change_feed_state_context["partitionKeyFeedRange"] =\ self._get_epk_range_for_partition_key(kwargs.pop('partition_key')) if kwargs.get("feed_range") is not None: @@ -548,22 +548,22 @@ def query_items_change_feed( @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, - *, - populate_index_metrics: Optional[bool] = None, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - max_integrated_cache_staleness_in_ms: Optional[int] = None, - priority: Optional[Literal["High", "Low"]] = None, - continuation_token_limit: Optional[int] = None, - **kwargs: Any + 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, + *, + populate_index_metrics: Optional[bool] = None, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + max_integrated_cache_staleness_in_ms: Optional[int] = None, + priority: Optional[Literal["High", "Low"]] = None, + continuation_token_limit: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Return all results matching the given `query`. @@ -671,7 +671,7 @@ def query_items( # pylint:disable=docstring-missing-param return items def __is_prefix_partitionkey( - self, partition_key: PartitionKeyType) -> bool: + self, partition_key: PartitionKeyType) -> bool: properties = self._get_properties() pk_properties = properties["partitionKey"] partition_key_definition = PartitionKey(path=pk_properties["paths"], kind=pk_properties["kind"]) @@ -680,20 +680,20 @@ def __is_prefix_partitionkey( @distributed_trace def replace_item( # pylint:disable=docstring-missing-param - self, - item: Union[str, Mapping[str, Any]], - body: Dict[str, Any], - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Mapping[str, Any]], + body: Dict[str, Any], + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Replaces the specified item if it exists in the container. @@ -715,9 +715,9 @@ def replace_item( # pylint:disable=docstring-missing-param 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 no_response: Indicates whether service should be instructed to skip - sending response payloads. When not specified explicitly here, the default value will be determined from - kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` + sending response payloads. When not specified explicitly here, the default value will be determined from + kwargs or when also not specified there from client-level kwargs. + :returns: A dict representing the item after replace went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The replace operation failed or the item with given id does not exist. @@ -753,24 +753,24 @@ def replace_item( # pylint:disable=docstring-missing-param request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.ReplaceItem( - document_link=item_link, new_document=body, options=request_options, **kwargs) + document_link=item_link, new_document=body, options=request_options, **kwargs) return result or {} @distributed_trace def upsert_item( # pylint:disable=docstring-missing-param - self, - body: Dict[str, Any], - populate_query_metrics: Optional[bool]=None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + populate_query_metrics: Optional[bool]=None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Insert or update the specified item. @@ -790,9 +790,9 @@ def upsert_item( # pylint:disable=docstring-missing-param :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or - when also not specified there from client-level kwargs. + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or + when also not specified there from client-level kwargs. :returns: A dict representing the upserted item. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The given item could not be upserted. :rtype: Dict[str, Any] @@ -825,30 +825,30 @@ def upsert_item( # pylint:disable=docstring-missing-param request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.UpsertItem( - database_or_container_link=self.container_link, - document=body, - options=request_options, - **kwargs - ) + database_or_container_link=self.container_link, + document=body, + options=request_options, + **kwargs + ) return result or {} @distributed_trace def create_item( # pylint:disable=docstring-missing-param - self, - body: Dict[str, Any], - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - indexing_directive: Optional[int] = None, - *, - enable_automatic_id_generation: bool = False, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + body: Dict[str, Any], + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + indexing_directive: Optional[int] = None, + *, + enable_automatic_id_generation: bool = False, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """Create an item in the container. @@ -872,8 +872,8 @@ def create_item( # pylint:disable=docstring-missing-param :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. :returns: A dict representing the new item. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists. @@ -908,25 +908,25 @@ def create_item( # pylint:disable=docstring-missing-param if self.container_link in self.__get_client_container_caches(): request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] result = self.client_connection.CreateItem( - database_or_container_link=self.container_link, document=body, options=request_options, **kwargs) + database_or_container_link=self.container_link, document=body, options=request_options, **kwargs) return result or {} @distributed_trace def patch_item( - self, - item: Union[str, Dict[str, Any]], - partition_key: PartitionKeyType, - patch_operations: List[Dict[str, Any]], - *, - filter_predicate: Optional[str] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - no_response: Optional[bool] = None, - **kwargs: Any + self, + item: Union[str, Dict[str, Any]], + partition_key: PartitionKeyType, + patch_operations: List[Dict[str, Any]], + *, + filter_predicate: Optional[str] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + no_response: Optional[bool] = None, + **kwargs: Any ) -> Dict[str, Any]: """ Patches the specified item with the provided operations if it exists in the container. @@ -950,10 +950,10 @@ def patch_item( :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 no_response: Indicates whether service should be instructed to skip sending - response payloads. When not specified explicitly here, the default value will be determined from kwargs or + :keyword bool no_response: Indicates whether service should be instructed to skip sending + response payloads. When not specified explicitly here, the default value will be determined from kwargs or when also not specified there from client-level kwargs. - :returns: A dict representing the item after the patch operations went through. The dict will be empty + :returns: A dict representing the item after the patch operations went through. The dict will be empty if `no_response` is specified. :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: The patch operations failed or the item with given id does not exist. @@ -983,22 +983,22 @@ def patch_item( request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"] item_link = self._get_document_link(item) result = self.client_connection.PatchItem( - document_link=item_link, operations=patch_operations, options=request_options, **kwargs) + document_link=item_link, operations=patch_operations, options=request_options, **kwargs) return result or {} @distributed_trace def execute_item_batch( - self, - batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + batch_operations: Sequence[Union[Tuple[str, Tuple[Any, ...]], Tuple[str, Tuple[Any, ...], Dict[str, Any]]]], + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> List[Dict[str, Any]]: """ Executes the transactional batch for the specified partition key. @@ -1042,19 +1042,19 @@ def execute_item_batch( @distributed_trace def delete_item( # pylint:disable=docstring-missing-param - self, - item: Union[Mapping[str, Any], str], - partition_key: PartitionKeyType, - populate_query_metrics: Optional[bool] = None, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - *, - session_token: Optional[str] = None, - initial_headers: Optional[Dict[str, str]] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - priority: Optional[Literal["High", "Low"]] = None, - **kwargs: Any + self, + item: Union[Mapping[str, Any], str], + partition_key: PartitionKeyType, + populate_query_metrics: Optional[bool] = None, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + *, + session_token: Optional[str] = None, + initial_headers: Optional[Dict[str, str]] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + priority: Optional[Literal["High", "Low"]] = None, + **kwargs: Any ) -> None: """Delete the specified item from the container. @@ -1154,9 +1154,9 @@ def get_throughput(self, **kwargs: Any) -> ThroughputProperties: @distributed_trace def replace_throughput( - self, - throughput: Union[int, ThroughputProperties], - **kwargs: Any + self, + throughput: Union[int, ThroughputProperties], + **kwargs: Any ) -> ThroughputProperties: """Replace the container's throughput. @@ -1188,9 +1188,9 @@ def replace_throughput( @distributed_trace def list_conflicts( - self, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """List all the conflicts in the container. @@ -1215,13 +1215,13 @@ def list_conflicts( @distributed_trace def query_conflicts( - self, - query: str, - parameters: Optional[List[Dict[str, object]]] = None, - enable_cross_partition_query: Optional[bool] = None, - partition_key: Optional[PartitionKeyType] = None, - max_item_count: Optional[int] = None, - **kwargs: Any + self, + query: str, + parameters: Optional[List[Dict[str, object]]] = None, + enable_cross_partition_query: Optional[bool] = None, + partition_key: Optional[PartitionKeyType] = None, + max_item_count: Optional[int] = None, + **kwargs: Any ) -> ItemPaged[Dict[str, Any]]: """Return all conflicts matching a given `query`. @@ -1261,10 +1261,10 @@ def query_conflicts( @distributed_trace def get_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any ) -> Dict[str, Any]: """Get the conflict identified by `conflict`. @@ -1289,10 +1289,10 @@ def get_conflict( @distributed_trace def delete_conflict( - self, - conflict: Union[str, Mapping[str, Any]], - partition_key: PartitionKeyType, - **kwargs: Any + self, + conflict: Union[str, Mapping[str, Any]], + partition_key: PartitionKeyType, + **kwargs: Any ) -> None: """Delete a specified conflict from the container. @@ -1319,15 +1319,15 @@ def delete_conflict( @distributed_trace def delete_all_items_by_partition_key( - self, - partition_key: PartitionKeyType, - *, - pre_trigger_include: Optional[str] = None, - post_trigger_include: Optional[str] = None, - session_token: Optional[str] = None, - etag: Optional[str] = None, - match_condition: Optional[MatchConditions] = None, - **kwargs: Any + self, + partition_key: PartitionKeyType, + *, + pre_trigger_include: Optional[str] = None, + post_trigger_include: Optional[str] = None, + session_token: Optional[str] = None, + etag: Optional[str] = None, + match_condition: Optional[MatchConditions] = None, + **kwargs: Any ) -> None: """The delete by partition key feature is an asynchronous, background operation that allows you to delete all documents with the same logical partition key value, using the Cosmos SDK. The delete by partition key @@ -1389,7 +1389,7 @@ def read_feed_ranges( **kwargs) return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange), self.container_link) - for partitionKeyRange in partition_key_ranges] + for partitionKeyRange in partition_key_ranges] def get_updated_session_token( self, diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 778de06f2c3a..f7b765c9aa65 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -65,7 +65,11 @@ def create_split_ranges(): ("AA", "EE"), "0:1#54#3=52,2:1#51#3=52"), # different version numbers ([(("AA", "BB"), "0:1#54#3=52"), (("AA", "BB"),"0:2#57#3=53")], - ("AA", "BB"), "0:2#57#3=53") + ("AA", "BB"), "0:2#57#3=53"), + # mixed scenarios + ([(("AA", "DD"), "0:1#60#3=53"), (("AA", "BB"), "1:1#54#3=52"), (("AA", "BB"), "1:1#52#3=53"), + (("BB", "CC"),"1:1#53#3=52"), (("BB", "CC"),"1:1#70#3=55,3:1#90#3=52"), + (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=53,1:1#70#3=55,3:1#90#3=52") ] @pytest.mark.cosmosEmulator diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py index 6e8397bdccd4..b42a7d87d664 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py @@ -14,6 +14,24 @@ from azure.cosmos.http_constants import HttpHeaders +def create_item(hpk): + if hpk: + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'state': 'CA', + 'city': 'LA' + str(random.randint(1, 10)), + 'zipcode': '90001' + } + else: + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + return item + + class TestUpdatedSessionTokenSplit(unittest.TestCase): """Test for session token helpers""" @@ -44,8 +62,9 @@ def test_updated_session_token_from_logical_pk(self): session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token + feed_ranges_and_session_tokens.append((target_feed_range, session_token)) - self.trigger_split(container) + self.trigger_split(container, 11000) target_session_token, _ = self.create_items_logical_pk(container, target_pk, session_token, feed_ranges_and_session_tokens) @@ -56,7 +75,7 @@ def test_updated_session_token_from_logical_pk(self): self.database.delete_container(container.id) def test_updated_session_token_from_physical_pk(self): - container = self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), + container = self.database.create_container("test_updated_session_token_from_physical_pk" + str(uuid.uuid4()), PartitionKey(path="/pk"), offer_throughput=400) feed_ranges_and_session_tokens = [] @@ -69,7 +88,7 @@ def test_updated_session_token_from_physical_pk(self): session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token - self.trigger_split(container) + self.trigger_split(container, 11000) _, target_feed_range, previous_session_token = self.create_items_physical_pk(container, pk_feed_range, session_token, @@ -88,11 +107,47 @@ def test_updated_session_token_from_physical_pk(self): assert '2' in pk_range_ids self.database.delete_container(container.id) + def test_updated_session_token_hpk(self): + container = self.database.create_container("test_updated_session_token_hpk" + str(uuid.uuid4()), + PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + pk = ['CA', 'LA1', '90001'] + pk_feed_range = container.feed_range_from_partition_key(pk) + target_session_token, target_feed_range, previous_session_token = self.create_items_physical_pk(container, + pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens, + True) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert session_token == target_session_token + self.database.delete_container(container.id) + + + def test_updated_session_token_logical_hpk(self): + container = self.database.create_container("test_updated_session_token_from_logical_hpk" + str(uuid.uuid4()), + PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + target_pk = ['CA1', 'LA1', '900011'] + target_session_token, previous_session_token = self.create_items_logical_pk(container, target_pk, + previous_session_token, + feed_ranges_and_session_tokens, + True) + target_feed_range = container.feed_range_from_partition_key(target_pk) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + self.database.delete_container(container.id) + @staticmethod - def trigger_split(container): + def trigger_split(container, throughput): print("Triggering a split in session token helpers") - container.replace_throughput(11000) + container.replace_throughput(throughput) print("changed offer to 11k") print("--------------------------------") print("Waiting for split to complete") @@ -112,25 +167,22 @@ def trigger_split(container): print("Split in session token helpers has completed") @staticmethod - def create_items_logical_pk(container, target_pk, previous_session_token, feed_ranges_and_session_tokens): + def create_items_logical_pk(container, target_pk, previous_session_token, feed_ranges_and_session_tokens, hpk=False): target_session_token = "" for i in range(100): - item = { - 'id': 'item' + str(uuid.uuid4()), - 'name': 'sample', - 'pk': 'A' + str(random.randint(1, 10)) - } + item = create_item(hpk) container.create_item(item, session_token=previous_session_token) session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] - if item['pk'] == target_pk: + pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] + if pk == target_pk: target_session_token = session_token previous_session_token = session_token - feed_ranges_and_session_tokens.append((container.feed_range_from_partition_key(item['pk']), + feed_ranges_and_session_tokens.append((container.feed_range_from_partition_key(pk), session_token)) return target_session_token, previous_session_token @staticmethod - def create_items_physical_pk(container, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens): + def create_items_physical_pk(container, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens, hpk=False): target_session_token = "" container_feed_ranges = container.read_feed_ranges() target_feed_range = None @@ -140,14 +192,14 @@ def create_items_physical_pk(container, pk_feed_range, previous_session_token, f break for i in range(100): - item = { - 'id': 'item' + str(uuid.uuid4()), - 'name': 'sample', - 'pk': 'A' + str(random.randint(1, 10)) - } + item = create_item(hpk) container.create_item(item, session_token=previous_session_token) session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] - curr_feed_range = container.feed_range_from_partition_key(item['pk']) + if hpk: + pk = [item['state'], item['city'], item['zipcode']] + curr_feed_range = container.feed_range_from_partition_key(pk) + else: + curr_feed_range = container.feed_range_from_partition_key(item['pk']) if container.is_feed_range_subset(target_feed_range, curr_feed_range): target_session_token = session_token previous_session_token = session_token From a9299ab431e407aced47228da4ebeaa4886dccc1 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 10 Oct 2024 17:01:50 -0700 Subject: [PATCH 39/59] reacting to comments --- .../cosmos/_change_feed/change_feed_state.py | 6 +- .../_change_feed/feed_range_internal.py | 39 ++-- .../azure-cosmos/azure/cosmos/_feed_range.py | 6 +- .../azure/cosmos/_session_token_helpers.py | 8 +- .../azure/cosmos/aio/_container.py | 4 +- .../azure-cosmos/azure/cosmos/container.py | 4 +- .../test/test_change_feed_split_async.py | 2 +- ...split.py => test_updated_session_token.py} | 25 +- .../test/test_updated_session_token_async.py | 213 ++++++++++++++++++ 9 files changed, 267 insertions(+), 40 deletions(-) rename sdk/cosmos/azure-cosmos/test/{test_updated_session_token_split.py => test_updated_session_token.py} (93%) create mode 100644 sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index a3adecbf34c2..d7cc5a0270c6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -389,7 +389,8 @@ def from_initial_state( feed_range =\ FeedRangeInternalPartitionKey( change_feed_state_context["partitionKey"], - change_feed_state_context["partitionKeyFeedRange"]) + change_feed_state_context["partitionKeyFeedRange"], + container_link) else: raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") else: @@ -399,7 +400,8 @@ def from_initial_state( "", "FF", True, - False) + False), + container_link ) change_feed_start_from = ( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index c04fda0952f9..471cad11e058 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -51,44 +51,51 @@ def _to_base64_encoded_string(self) -> str: class FeedRangeInternalPartitionKey(FeedRangeInternal): type_property_name = "PK" + container_link_property_name = "Container" def __init__( self, pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], - feed_range: Range) -> None: # pylint: disable=line-too-long + feed_range: Range, + container_link: str) -> None: # pylint: disable=line-too-long if pk_value is None: raise ValueError("PartitionKey cannot be None") if feed_range is None: raise ValueError("Feed range cannot be None") + if container_link is None: + raise ValueError("Container link cannot be None") self._pk_value = pk_value self._feed_range = feed_range + self._container_link = container_link def get_normalized_range(self) -> Range: return self._feed_range.to_normalized_range() def to_dict(self) -> Dict[str, Any]: if isinstance(self._pk_value, _Undefined): - return { self.type_property_name: [{}] } + return { self.type_property_name: [{}], self.container_link_property_name: self._container_link } if isinstance(self._pk_value, _Empty): - return { self.type_property_name: [] } + return { self.type_property_name: [], self.container_link_property_name: self._container_link} if isinstance(self._pk_value, list): - return { self.type_property_name: list(self._pk_value) } + return { self.type_property_name: list(self._pk_value), + self.container_link_property_name: self._container_link } - return { self.type_property_name: self._pk_value } + return { self.type_property_name: self._pk_value, self.container_link_property_name: self._container_link } @classmethod def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInternalPartitionKey': - if data.get(cls.type_property_name): + if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): pk_value = data.get(cls.type_property_name) + container_link = data.get(cls.container_link_property_name) if not pk_value: - return cls(_Empty(), feed_range) + return cls(_Empty(), feed_range, container_link) if pk_value == [{}]: - return cls(_Undefined(), feed_range) + return cls(_Undefined(), feed_range, container_link) if isinstance(pk_value, list): - return cls(list(pk_value), feed_range) - return cls(data[cls.type_property_name], feed_range) + return cls(list(pk_value), feed_range, container_link) + return cls(data[cls.type_property_name], feed_range, container_link) raise ValueError(f"Can not parse FeedRangeInternalPartitionKey from the json," f" there is no property {cls.type_property_name}") @@ -96,12 +103,14 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInterna class FeedRangeInternalEpk(FeedRangeInternal): type_property_name = "Range" + container_link_property_name = "Container" - def __init__(self, feed_range: Range) -> None: + def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") self._range = feed_range + self._container_link = container_link self._base64_encoded_string: Optional[str] = None def get_normalized_range(self) -> Range: @@ -109,14 +118,16 @@ def get_normalized_range(self) -> Range: def to_dict(self) -> Dict[str, Any]: return { - self.type_property_name: self._range.to_dict() + self.type_property_name: self._range.to_dict(), + self.container_link_property_name: self._container_link } @classmethod def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeInternalEpk': - if data.get(cls.type_property_name): + if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) - return cls(feed_range) + container_link = data.get(cls.container_link_property_name) + return cls(feed_range, container_link) raise ValueError(f"Can not parse FeedRangeInternalEPK from the json," f" there is no property {cls.type_property_name}") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index f4b8dc576cb5..357a91a60d4b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -55,8 +55,7 @@ def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") - self._feed_range_internal = FeedRangeInternalEpk(feed_range) - self._container_link = container_link + self._feed_range_internal = FeedRangeInternalEpk(feed_range, container_link) def __str__(self) -> str: """Get a json representation of the feed range. @@ -68,4 +67,5 @@ def __str__(self) -> str: @classmethod def _from_json(cls, data: Dict[str, Any]) -> 'FeedRange': - return cls(FeedRangeInternalEpk.from_json(data)._range) + feed_range_internal = FeedRangeInternalEpk.from_json(data) + return cls(feed_range_internal._range, feed_range_internal._container_link) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index ad07de732eb4..6efaee3647fa 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -75,9 +75,9 @@ def merge_session_tokens_for_same_partition(session_tokens: [str]) -> [str]: return session_tokens -def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: # pylint: disable=too-many-nested-blocks +def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: processed_ranges = [] - while len(overlapping_ranges) != 0: + while len(overlapping_ranges) != 0: # pylint: disable=too-many-nested-blocks feed_range_cmp, session_token_cmp = overlapping_ranges[0] # compound session tokens are not considered for merging if is_compound_session_token(session_token_cmp): @@ -144,13 +144,13 @@ def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, st def get_updated_session_token(feed_ranges_to_session_tokens: [(FeedRange, str)], target_feed_range: FeedRange, container_link: str): - if target_feed_range._container_link != container_link: + if target_feed_range._feed_range_internal._container_link != container_link: raise ValueError('The target feed range does not belong to the container.') target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges overlapping_ranges = [] for feed_range_to_session_token in feed_ranges_to_session_tokens: - if feed_range_to_session_token[0]._container_link != container_link: + if feed_range_to_session_token[0]._feed_range_internal._container_link != container_link: raise ValueError('The feed range does not belong to the container.') if Range.overlaps(target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range()): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 6e159cd0454c..93b56c76912c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1357,8 +1357,8 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - if (child_feed_range._container_link != self.container_link or - parent_feed_range._container_link != self.container_link): + if (child_feed_range._feed_range_internal._container_link != self.container_link or + parent_feed_range._feed_range_internal._container_link != self.container_link): raise ValueError("Feed ranges must be from the same container.") return child_feed_range._feed_range_internal.get_normalized_range().is_subset( parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index a2699dfc8b08..3810e5a6098b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1427,8 +1427,8 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - if (child_feed_range._container_link != self.container_link or - parent_feed_range._container_link != self.container_link): + if (child_feed_range._feed_range_internal._container_link != self.container_link or + parent_feed_range._feed_range_internal._container_link != self.container_link): raise ValueError("Feed ranges must be from the same container.") return child_feed_range._feed_range_internal.get_normalized_range().is_subset( parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py index 60f7b2810884..7c13f4682d02 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py @@ -88,7 +88,7 @@ async def test_query_change_feed_with_split_async(self): actual_ids.append(item['id']) assert actual_ids == expected_ids - self.created_database.delete_container(created_collection.id) + await self.created_database.delete_container(created_collection.id) if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token.py similarity index 93% rename from sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py rename to sdk/cosmos/azure-cosmos/test/test_updated_session_token.py index b42a7d87d664..efbf9a9ad890 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_split.py +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token.py @@ -32,7 +32,7 @@ def create_item(hpk): return item -class TestUpdatedSessionTokenSplit(unittest.TestCase): +class TestUpdatedSessionToken(unittest.TestCase): """Test for session token helpers""" created_db: DatabaseProxy = None @@ -55,10 +55,10 @@ def test_updated_session_token_from_logical_pk(self): feed_ranges_and_session_tokens = [] previous_session_token = "" target_pk = 'A1' - target_session_token, previous_session_token = self.create_items_logical_pk(container, target_pk, + target_feed_range = container.feed_range_from_partition_key(target_pk) + target_session_token, previous_session_token = self.create_items_logical_pk(container, target_feed_range, previous_session_token, feed_ranges_and_session_tokens) - target_feed_range = container.feed_range_from_partition_key(target_pk) session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token @@ -66,7 +66,7 @@ def test_updated_session_token_from_logical_pk(self): self.trigger_split(container, 11000) - target_session_token, _ = self.create_items_logical_pk(container, target_pk, session_token, + target_session_token, _ = self.create_items_logical_pk(container, target_feed_range, session_token, feed_ranges_and_session_tokens) target_feed_range = container.feed_range_from_partition_key(target_pk) session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) @@ -132,12 +132,12 @@ def test_updated_session_token_logical_hpk(self): offer_throughput=400) feed_ranges_and_session_tokens = [] previous_session_token = "" - target_pk = ['CA1', 'LA1', '900011'] - target_session_token, previous_session_token = self.create_items_logical_pk(container, target_pk, + target_pk = ['CA', 'LA1', '90001'] + target_feed_range = container.feed_range_from_partition_key(target_pk) + target_session_token, previous_session_token = self.create_items_logical_pk(container, target_feed_range, previous_session_token, feed_ranges_and_session_tokens, True) - target_feed_range = container.feed_range_from_partition_key(target_pk) session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token @@ -163,21 +163,22 @@ def trigger_split(container, throughput): time.sleep(60) else: break - - print("Split in session token helpers has completed") + print("Split in session token helpers has completed") @staticmethod - def create_items_logical_pk(container, target_pk, previous_session_token, feed_ranges_and_session_tokens, hpk=False): + def create_items_logical_pk(container, target_pk_range, previous_session_token, feed_ranges_and_session_tokens, hpk=False): target_session_token = "" for i in range(100): item = create_item(hpk) container.create_item(item, session_token=previous_session_token) session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] - if pk == target_pk: + pk_range = container.feed_range_from_partition_key(pk) + if (pk_range._feed_range_internal.get_normalized_range() == + target_pk_range._feed_range_internal.get_normalized_range()): target_session_token = session_token previous_session_token = session_token - feed_ranges_and_session_tokens.append((container.feed_range_from_partition_key(pk), + feed_ranges_and_session_tokens.append((pk_range, session_token)) return target_session_token, previous_session_token diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py new file mode 100644 index 000000000000..b757456a44a4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py @@ -0,0 +1,213 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +import random +import time +import unittest +import uuid + + +import test_config +from azure.cosmos import PartitionKey +from azure.cosmos._session_token_helpers import is_compound_session_token, parse_session_token +from azure.cosmos.aio import DatabaseProxy +from azure.cosmos.aio import CosmosClient +from azure.cosmos.http_constants import HttpHeaders + + +def create_item(hpk): + if hpk: + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'state': 'CA', + 'city': 'LA' + str(random.randint(1, 10)), + 'zipcode': '90001' + } + else: + item = { + 'id': 'item' + str(uuid.uuid4()), + 'name': 'sample', + 'pk': 'A' + str(random.randint(1, 10)) + } + return item + + +class TestUpdatedSessionTokenAsync(unittest.IsolatedAsyncioTestCase): + """Test for session token helpers""" + + created_db: DatabaseProxy = None + client: CosmosClient = None + host = test_config.TestConfig.host + masterKey = test_config.TestConfig.masterKey + configs = test_config.TestConfig + TEST_DATABASE_ID = configs.TEST_DATABASE_ID + + @classmethod + def setUpClass(cls): + cls.client = CosmosClient(cls.host, cls.masterKey) + cls.database = cls.client.get_database_client(cls.TEST_DATABASE_ID) + + + async def test_updated_session_token_from_logical_pk(self): + container = await self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + target_pk = 'A1' + target_feed_range = await container.feed_range_from_partition_key(target_pk) + target_session_token, previous_session_token = await self.create_items_logical_pk(container, target_feed_range, + previous_session_token, + feed_ranges_and_session_tokens) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + feed_ranges_and_session_tokens.append((target_feed_range, session_token)) + + await self.trigger_split(container, 11000) + + target_session_token, _ = await self.create_items_logical_pk(container, target_feed_range, session_token, + feed_ranges_and_session_tokens) + target_feed_range = await container.feed_range_from_partition_key(target_pk) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + await self.database.delete_container(container.id) + + async def test_updated_session_token_from_physical_pk(self): + container = await self.database.create_container("test_updated_session_token_from_physical_pk" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + pk_feed_range = await container.feed_range_from_partition_key('A1') + target_session_token, target_feed_range, previous_session_token = await self.create_items_physical_pk(container, pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert session_token == target_session_token + + await self.trigger_split(container, 11000) + + _, target_feed_range, previous_session_token = await self.create_items_physical_pk(container, pk_feed_range, + session_token, + feed_ranges_and_session_tokens) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert is_compound_session_token(session_token) + session_tokens = session_token.split(",") + assert len(session_tokens) == 2 + pk_range_id1, session_token1 = parse_session_token(session_tokens[0]) + pk_range_id2, session_token2 = parse_session_token(session_tokens[1]) + pk_range_ids = [pk_range_id1, pk_range_id2] + + assert 320 == (session_token1.global_lsn + session_token2.global_lsn) + assert '1' in pk_range_ids + assert '2' in pk_range_ids + await self.database.delete_container(container.id) + + async def test_updated_session_token_hpk(self): + container = await self.database.create_container("test_updated_session_token_hpk" + str(uuid.uuid4()), + PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + pk = ['CA', 'LA1', '90001'] + pk_feed_range = await container.feed_range_from_partition_key(pk) + target_session_token, target_feed_range, previous_session_token = await self.create_items_physical_pk(container, + pk_feed_range, + previous_session_token, + feed_ranges_and_session_tokens, + True) + + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + assert session_token == target_session_token + await self.database.delete_container(container.id) + + + async def test_updated_session_token_logical_hpk(self): + container = await self.database.create_container("test_updated_session_token_from_logical_hpk" + str(uuid.uuid4()), + PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), + offer_throughput=400) + feed_ranges_and_session_tokens = [] + previous_session_token = "" + target_pk = ['CA', 'LA1', '90001'] + target_feed_range = await container.feed_range_from_partition_key(target_pk) + target_session_token, previous_session_token = await self.create_items_logical_pk(container, target_feed_range, + previous_session_token, + feed_ranges_and_session_tokens, + True) + session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + + assert session_token == target_session_token + await self.database.delete_container(container.id) + + + @staticmethod + async def trigger_split(container, throughput): + print("Triggering a split in session token helpers") + await container.replace_throughput(throughput) + print("changed offer to 11k") + print("--------------------------------") + print("Waiting for split to complete") + start_time = time.time() + + while True: + offer = await container.get_throughput() + if offer.properties['content'].get('isOfferReplacePending', False): + if time.time() - start_time > 60 * 25: # timeout test at 25 minutes + unittest.skip("Partition split didn't complete in time.") + else: + print("Waiting for split to complete") + time.sleep(60) + else: + break + print("Split in session token helpers has completed") + + @staticmethod + async def create_items_logical_pk(container, target_pk_range, previous_session_token, feed_ranges_and_session_tokens, hpk=False): + target_session_token = "" + for i in range(100): + item = create_item(hpk) + await container.create_item(item, session_token=previous_session_token) + session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] + pk_feed_range = await container.feed_range_from_partition_key(pk) + if (pk_feed_range._feed_range_internal.get_normalized_range() == + target_pk_range._feed_range_internal.get_normalized_range()): + target_session_token = session_token + previous_session_token = session_token + feed_ranges_and_session_tokens.append((pk_feed_range, + session_token)) + return target_session_token, previous_session_token + + @staticmethod + async def create_items_physical_pk(container, pk_feed_range, previous_session_token, feed_ranges_and_session_tokens, hpk=False): + target_session_token = "" + container_feed_ranges = await container.read_feed_ranges() + target_feed_range = None + for feed_range in container_feed_ranges: + if container.is_feed_range_subset(feed_range, pk_feed_range): + target_feed_range = feed_range + break + + for i in range(100): + item = create_item(hpk) + await container.create_item(item, session_token=previous_session_token) + session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + if hpk: + pk = [item['state'], item['city'], item['zipcode']] + curr_feed_range = await container.feed_range_from_partition_key(pk) + else: + curr_feed_range = await container.feed_range_from_partition_key(item['pk']) + if container.is_feed_range_subset(target_feed_range, curr_feed_range): + target_session_token = session_token + previous_session_token = session_token + feed_ranges_and_session_tokens.append((curr_feed_range, session_token)) + + return target_session_token, target_feed_range, previous_session_token + +if __name__ == '__main__': + unittest.main() From 215501640377aee6c1abe417c579d5741e071655 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 11 Oct 2024 11:16:34 -0700 Subject: [PATCH 40/59] fix tests and mypy --- .../azure/cosmos/_change_feed/feed_range_internal.py | 2 +- sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py | 3 +++ .../azure/cosmos/_session_token_helpers.py | 10 ++++++---- .../test/test_updated_session_token_async.py | 10 ++++++---- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index 471cad11e058..28b1eb59f76a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -105,7 +105,7 @@ class FeedRangeInternalEpk(FeedRangeInternal): type_property_name = "Range" container_link_property_name = "Container" - def __init__(self, feed_range: Range, container_link: str) -> None: + def __init__(self, feed_range: Range, container_link: Optional[any]) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 357a91a60d4b..39de2db2e730 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -32,6 +32,9 @@ class FeedRange(ABC): """Represents a single feed range in an Azure Cosmos DB SQL API container. """ + def __init__(self) -> None: + self._feed_range_internal = None + @staticmethod def from_string(json_str: str) -> 'FeedRange': """ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 6efaee3647fa..f199ed23ad5a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -21,6 +21,8 @@ """Internal Helper functions for manipulating session tokens. """ +from typing import Tuple, List + from azure.cosmos._routing.routing_range import Range from azure.cosmos._vector_session_token import VectorSessionToken from ._feed_range import FeedRange @@ -40,11 +42,11 @@ def merge_session_tokens_with_same_range(session_token1: str, session_token2: st def is_compound_session_token(session_token: str) -> bool: return "," in session_token -def parse_session_token(session_token: str) -> (str, VectorSessionToken): +def parse_session_token(session_token: str) -> Tuple[str, VectorSessionToken]: tokens = session_token.split(":") return tokens[0], VectorSessionToken.create(tokens[1]) -def split_compound_session_tokens(compound_session_tokens: [(Range, str)]) -> [str]: +def split_compound_session_tokens(compound_session_tokens: List[Tuple[Range, str]]) -> List[str]: session_tokens = [] for _, session_token in compound_session_tokens: if is_compound_session_token(session_token): @@ -55,7 +57,7 @@ def split_compound_session_tokens(compound_session_tokens: [(Range, str)]) -> [s session_tokens.append(session_token) return session_tokens -def merge_session_tokens_for_same_partition(session_tokens: [str]) -> [str]: +def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[str]: i = 0 while i < len(session_tokens): j = i + 1 @@ -75,7 +77,7 @@ def merge_session_tokens_for_same_partition(session_tokens: [str]) -> [str]: return session_tokens -def merge_ranges_with_subsets(overlapping_ranges: [(Range, str)]) -> [(Range, str)]: +def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> List[Tuple[Range, str]]: processed_ranges = [] while len(overlapping_ranges) != 0: # pylint: disable=too-many-nested-blocks feed_range_cmp, session_token_cmp = overlapping_ranges[0] diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py index b757456a44a4..3b8d0eb115cb 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py @@ -43,11 +43,13 @@ class TestUpdatedSessionTokenAsync(unittest.IsolatedAsyncioTestCase): configs = test_config.TestConfig TEST_DATABASE_ID = configs.TEST_DATABASE_ID - @classmethod - def setUpClass(cls): - cls.client = CosmosClient(cls.host, cls.masterKey) - cls.database = cls.client.get_database_client(cls.TEST_DATABASE_ID) + async def asyncSetUp(self): + self.client = CosmosClient(self.host, self.masterKey) + self.database = self.client.get_database_client(self.TEST_DATABASE_ID) + async def tearDown(self): + await self.client.delete_database(self.database.id) + await self.client.close() async def test_updated_session_token_from_logical_pk(self): container = await self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), From 0436355875388605d5ff61a7f105b4321cc7431e Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 11 Oct 2024 12:54:07 -0700 Subject: [PATCH 41/59] fix mypy --- .../cosmos/_change_feed/feed_range_internal.py | 17 ++++++++++++++--- .../azure-cosmos/azure/cosmos/_feed_range.py | 14 ++++++++++---- .../azure/cosmos/_session_token_helpers.py | 2 +- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index 28b1eb59f76a..f4b0607c3a31 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -33,6 +33,11 @@ class FeedRangeInternal(ABC): + @property + @abstractmethod + def _container_link(self) -> str: + pass + @abstractmethod def get_normalized_range(self) -> Range: pass @@ -88,7 +93,7 @@ def to_dict(self) -> Dict[str, Any]: def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInternalPartitionKey': if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): pk_value = data.get(cls.type_property_name) - container_link = data.get(cls.container_link_property_name) + container_link = str(data.get(cls.container_link_property_name)) if not pk_value: return cls(_Empty(), feed_range, container_link) if pk_value == [{}]: @@ -100,12 +105,15 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInterna raise ValueError(f"Can not parse FeedRangeInternalPartitionKey from the json," f" there is no property {cls.type_property_name}") + def _container_link(self) -> str: + return self._container_link + class FeedRangeInternalEpk(FeedRangeInternal): type_property_name = "Range" container_link_property_name = "Container" - def __init__(self, feed_range: Range, container_link: Optional[any]) -> None: + def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") @@ -126,7 +134,7 @@ def to_dict(self) -> Dict[str, Any]: def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeInternalEpk': if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) - container_link = data.get(cls.container_link_property_name) + container_link = str(data.get(cls.container_link_property_name)) return cls(feed_range, container_link) raise ValueError(f"Can not parse FeedRangeInternalEPK from the json," f" there is no property {cls.type_property_name}") @@ -141,3 +149,6 @@ def __str__(self) -> str: self._base64_encoded_string = self._to_base64_encoded_string() return self._base64_encoded_string + + def _container_link(self) -> str: + return self._container_link diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 39de2db2e730..56eb98239d03 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -21,10 +21,10 @@ import base64 import json -from abc import ABC +from abc import ABC, abstractmethod from typing import Any, Dict -from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk, FeedRangeInternal from azure.cosmos._routing.routing_range import Range # pylint: disable=protected-access @@ -32,8 +32,11 @@ class FeedRange(ABC): """Represents a single feed range in an Azure Cosmos DB SQL API container. """ - def __init__(self) -> None: - self._feed_range_internal = None + + @property + @abstractmethod + def _feed_range_internal(self) -> FeedRangeInternal: + pass @staticmethod def from_string(json_str: str) -> 'FeedRange': @@ -72,3 +75,6 @@ def __str__(self) -> str: def _from_json(cls, data: Dict[str, Any]) -> 'FeedRange': feed_range_internal = FeedRangeInternalEpk.from_json(data) return cls(feed_range_internal._range, feed_range_internal._container_link) + + def _feed_range_internal(self): + return self._feed_range_internal diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index f199ed23ad5a..680272f02515 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -144,7 +144,7 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li overlapping_ranges.remove(overlapping_ranges[0]) return processed_ranges -def get_updated_session_token(feed_ranges_to_session_tokens: [(FeedRange, str)], target_feed_range: FeedRange, +def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange, container_link: str): if target_feed_range._feed_range_internal._container_link != container_link: raise ValueError('The target feed range does not belong to the container.') From 103eb41ab762e9ccfabf1cdacc94eaad96e99d86 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 11 Oct 2024 14:17:34 -0700 Subject: [PATCH 42/59] fix mypy --- .../_change_feed/feed_range_internal.py | 20 ++++++------------- .../azure-cosmos/azure/cosmos/_feed_range.py | 14 ++++--------- .../azure/cosmos/aio/_container.py | 4 ++-- .../azure-cosmos/azure/cosmos/container.py | 6 +++--- 4 files changed, 15 insertions(+), 29 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index f4b0607c3a31..5d1aa8a8aa52 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -33,10 +33,10 @@ class FeedRangeInternal(ABC): - @property - @abstractmethod - def _container_link(self) -> str: - pass + def __init__(self, container_link: str) -> None: + if container_link is None: + raise ValueError("container_link cannot be None") + self._container_link = container_link @abstractmethod def get_normalized_range(self) -> Range: @@ -68,12 +68,9 @@ def __init__( raise ValueError("PartitionKey cannot be None") if feed_range is None: raise ValueError("Feed range cannot be None") - if container_link is None: - raise ValueError("Container link cannot be None") - + super().__init__(container_link) self._pk_value = pk_value self._feed_range = feed_range - self._container_link = container_link def get_normalized_range(self) -> Range: return self._feed_range.to_normalized_range() @@ -105,8 +102,6 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInterna raise ValueError(f"Can not parse FeedRangeInternalPartitionKey from the json," f" there is no property {cls.type_property_name}") - def _container_link(self) -> str: - return self._container_link class FeedRangeInternalEpk(FeedRangeInternal): @@ -116,9 +111,9 @@ class FeedRangeInternalEpk(FeedRangeInternal): def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") + super().__init__(container_link) self._range = feed_range - self._container_link = container_link self._base64_encoded_string: Optional[str] = None def get_normalized_range(self) -> Range: @@ -149,6 +144,3 @@ def __str__(self) -> str: self._base64_encoded_string = self._to_base64_encoded_string() return self._base64_encoded_string - - def _container_link(self) -> str: - return self._container_link diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 56eb98239d03..2e6b4fb2a5a5 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -21,7 +21,7 @@ import base64 import json -from abc import ABC, abstractmethod +from abc import ABC from typing import Any, Dict from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk, FeedRangeInternal @@ -33,10 +33,8 @@ class FeedRange(ABC): """ - @property - @abstractmethod - def _feed_range_internal(self) -> FeedRangeInternal: - pass + def __init__(self, feed_range_internal: FeedRangeInternal) -> None: + self._feed_range_internal = feed_range_internal @staticmethod def from_string(json_str: str) -> 'FeedRange': @@ -60,8 +58,7 @@ class FeedRangeEpk(FeedRange): def __init__(self, feed_range: Range, container_link: str) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") - - self._feed_range_internal = FeedRangeInternalEpk(feed_range, container_link) + super().__init__(FeedRangeInternalEpk(feed_range, container_link)) def __str__(self) -> str: """Get a json representation of the feed range. @@ -75,6 +72,3 @@ def __str__(self) -> str: def _from_json(cls, data: Dict[str, Any]) -> 'FeedRange': feed_range_internal = FeedRangeInternalEpk.from_json(data) return cls(feed_range_internal._range, feed_range_internal._container_link) - - def _feed_range_internal(self): - return self._feed_range_internal diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 93b56c76912c..a2c8ad456b24 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1322,7 +1322,7 @@ async def read_feed_ranges( for partitionKeyRange in partition_key_ranges] def get_updated_session_token(self, - feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], + feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange ) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples @@ -1331,7 +1331,7 @@ def get_updated_session_token(self, session token. Session tokens and feed ranges are scoped to a container. Only input session tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. - :type feed_ranges_to_session_tokens: List[Tuple(str, FeedRange)] + :type feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]] :param target_feed_range: feed range to get most up to date session token. :type target_feed_range: FeedRange :returns: a session token diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 3810e5a6098b..213b50ae2f94 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1393,7 +1393,7 @@ def read_feed_ranges( def get_updated_session_token( self, - feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]], + feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a response from any crud operation. @@ -1401,9 +1401,9 @@ def get_updated_session_token( session token. Session tokens and feed ranges are scoped to a container. Only input session tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. - :type feed_ranges_to_session_tokens: List[Tuple(str, Range)] + :type feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]] :param target_feed_range: feed range to get most up to date session token. - :type target_feed_range: Range + :type target_feed_range: FeedRange :returns: a session token :rtype: str """ From 76451df8e6968bc2ab62addaeae790e91201cb62 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 15 Oct 2024 13:42:43 -0700 Subject: [PATCH 43/59] reacting to comments --- .../azure-cosmos/azure/cosmos/_feed_range.py | 1 + .../azure/cosmos/_routing/routing_range.py | 5 ++--- .../azure/cosmos/_session_token_helpers.py | 19 +++++++++++++++++++ .../test/test_session_token_helpers.py | 5 ++++- 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 2e6b4fb2a5a5..2afb583728f0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -65,6 +65,7 @@ def __str__(self) -> str: The returned json string can be used to create a new feed range from it. :return: A json representation of the feed range. + :rtype: str """ return self._feed_range_internal.__str__() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py index ba1a88b2d4f3..452bc32e5b34 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py @@ -224,6 +224,5 @@ def merge(self, other: 'Range') -> 'Range': def is_subset(self, parent_range: 'Range') -> bool: normalized_parent_range = parent_range.to_normalized_range() normalized_child_range = self.to_normalized_range() - return normalized_parent_range.contains(normalized_child_range.min) and \ - (normalized_parent_range.contains(normalized_child_range.max) - or normalized_parent_range.max == normalized_child_range.max) + return (normalized_parent_range.min <= normalized_child_range.min and + normalized_parent_range.max >= normalized_child_range.max) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 680272f02515..d31bdbf5f689 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -29,10 +29,18 @@ # pylint: disable=protected-access + +# ex inputs: +# 1. "1:1#51", "1:1#55" -> "1:1#55" +# 2. "0:1#57", "1:1#52" -> "0:1#57" def merge_session_tokens_with_same_range(session_token1: str, session_token2: str) -> str: pk_range_id1, vector_session_token1 = parse_session_token(session_token1) pk_range_id2, vector_session_token2 = parse_session_token(session_token2) pk_range_id = pk_range_id1 + # The pkrangeid could be different in this scenario + # Ex. get_updated_session_token([("AA", "BB"), "1:1#51"], ("AA", "DD")) -> "1:1#51" + # Then we input this back into get_updated_session_token after a merge happened + # get_updated_session_token([("AA", "DD"), "1:1#51", ("AA", "DD"), "0:1#55"], ("AA", "DD")) -> "0:1#55" if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 @@ -57,6 +65,8 @@ def split_compound_session_tokens(compound_session_tokens: List[Tuple[Range, str session_tokens.append(session_token) return session_tokens +# ex inputs: +# ["1:1#51", "1:1#55", "1:1#57"] -> ["1:1#57"] def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[str]: i = 0 while i < len(session_tokens): @@ -77,6 +87,14 @@ def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[s return session_tokens +# ex inputs: +# 1. [(("AA", "BB"), "1:1#51"), (("BB", "DD"), "2:1#51"), (("AA", "DD"), "0:1#55")] -> +# [("AA", "DD"), "0:1#55"] +# 2. [(("AA", "BB"), "1:1#57"), (("BB", "DD"), "2:1#58"), (("AA", "DD"), "0:1#55")] -> +# [("AA", "DD"), "1:1#57,2:1#58"] +# 3. [(("AA", "BB"), "1:1#57"), (("BB", "DD"), "2:1#52"), (("AA", "DD"), "0:1#55")] -> +# [("AA", "DD"), "1:1#57,2:1#52,0:1#55"] +# compound session tokens are not considered will just pass them along def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> List[Tuple[Range, str]]: processed_ranges = [] while len(overlapping_ranges) != 0: # pylint: disable=too-many-nested-blocks @@ -146,6 +164,7 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange, container_link: str): + if target_feed_range._feed_range_internal._container_link != container_link: raise ValueError('The target feed range does not belong to the container.') target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index f7b765c9aa65..394422d7ed07 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -30,9 +30,12 @@ def setup(): } def create_split_ranges(): - return [ # split with two children + return [ # split with two children ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), + # same range different pkrangeid + ([(("AA", "DD"), "1:1#51#3=52"), (("AA", "DD"),"0:1#55#3=52")], + ("AA", "DD"), "0:1#55#3=52"), # split with one child ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), From 7b0f4b782e1d47a3483c0a5f98db5ec0203470eb Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 15 Oct 2024 14:36:15 -0700 Subject: [PATCH 44/59] reacting to comments --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- .../cosmos/_change_feed/change_feed_state.py | 6 +-- .../_change_feed/feed_range_internal.py | 44 +++++++------------ .../azure-cosmos/azure/cosmos/_feed_range.py | 7 ++- .../azure/cosmos/_session_token_helpers.py | 8 +--- .../azure/cosmos/aio/_container.py | 14 +++--- .../azure-cosmos/azure/cosmos/container.py | 9 ++-- .../azure-cosmos/test/test_feed_range.py | 23 ++-------- .../test/test_feed_range_async.py | 19 +------- .../test/test_session_token_helpers.py | 35 +++------------ .../test/test_updated_session_token_async.py | 16 +++---- 11 files changed, 51 insertions(+), 132 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 933c07cc8e67..e0c46a723d73 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -8,7 +8,7 @@ * Added option to disable write payload on writes. See [PR 37365](https://github.com/Azure/azure-sdk-for-python/pull/37365) * Added get feed ranges API. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) * Added feed range support in `query_items_change_feed`. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687) -* Added helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971) +* Added **preview** helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971) #### Bugs Fixed * Consolidated Container Properties Cache to be in the Client to cache partition key definition and container rid to avoid unnecessary container reads. See [PR 35731](https://github.com/Azure/azure-sdk-for-python/pull/35731). diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index d7cc5a0270c6..a3adecbf34c2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -389,8 +389,7 @@ def from_initial_state( feed_range =\ FeedRangeInternalPartitionKey( change_feed_state_context["partitionKey"], - change_feed_state_context["partitionKeyFeedRange"], - container_link) + change_feed_state_context["partitionKeyFeedRange"]) else: raise ValueError("partitionKey is in the changeFeedStateContext, but missing partitionKeyFeedRange") else: @@ -400,8 +399,7 @@ def from_initial_state( "", "FF", True, - False), - container_link + False) ) change_feed_start_from = ( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py index 5d1aa8a8aa52..c04fda0952f9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py @@ -33,11 +33,6 @@ class FeedRangeInternal(ABC): - def __init__(self, container_link: str) -> None: - if container_link is None: - raise ValueError("container_link cannot be None") - self._container_link = container_link - @abstractmethod def get_normalized_range(self) -> Range: pass @@ -56,19 +51,17 @@ def _to_base64_encoded_string(self) -> str: class FeedRangeInternalPartitionKey(FeedRangeInternal): type_property_name = "PK" - container_link_property_name = "Container" def __init__( self, pk_value: Union[str, int, float, bool, List[Union[str, int, float, bool]], _Empty, _Undefined], - feed_range: Range, - container_link: str) -> None: # pylint: disable=line-too-long + feed_range: Range) -> None: # pylint: disable=line-too-long if pk_value is None: raise ValueError("PartitionKey cannot be None") if feed_range is None: raise ValueError("Feed range cannot be None") - super().__init__(container_link) + self._pk_value = pk_value self._feed_range = feed_range @@ -77,41 +70,36 @@ def get_normalized_range(self) -> Range: def to_dict(self) -> Dict[str, Any]: if isinstance(self._pk_value, _Undefined): - return { self.type_property_name: [{}], self.container_link_property_name: self._container_link } + return { self.type_property_name: [{}] } if isinstance(self._pk_value, _Empty): - return { self.type_property_name: [], self.container_link_property_name: self._container_link} + return { self.type_property_name: [] } if isinstance(self._pk_value, list): - return { self.type_property_name: list(self._pk_value), - self.container_link_property_name: self._container_link } + return { self.type_property_name: list(self._pk_value) } - return { self.type_property_name: self._pk_value, self.container_link_property_name: self._container_link } + return { self.type_property_name: self._pk_value } @classmethod def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangeInternalPartitionKey': - if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): + if data.get(cls.type_property_name): pk_value = data.get(cls.type_property_name) - container_link = str(data.get(cls.container_link_property_name)) if not pk_value: - return cls(_Empty(), feed_range, container_link) + return cls(_Empty(), feed_range) if pk_value == [{}]: - return cls(_Undefined(), feed_range, container_link) + return cls(_Undefined(), feed_range) if isinstance(pk_value, list): - return cls(list(pk_value), feed_range, container_link) - return cls(data[cls.type_property_name], feed_range, container_link) + return cls(list(pk_value), feed_range) + return cls(data[cls.type_property_name], feed_range) raise ValueError(f"Can not parse FeedRangeInternalPartitionKey from the json," f" there is no property {cls.type_property_name}") - class FeedRangeInternalEpk(FeedRangeInternal): type_property_name = "Range" - container_link_property_name = "Container" - def __init__(self, feed_range: Range, container_link: str) -> None: + def __init__(self, feed_range: Range) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") - super().__init__(container_link) self._range = feed_range self._base64_encoded_string: Optional[str] = None @@ -121,16 +109,14 @@ def get_normalized_range(self) -> Range: def to_dict(self) -> Dict[str, Any]: return { - self.type_property_name: self._range.to_dict(), - self.container_link_property_name: self._container_link + self.type_property_name: self._range.to_dict() } @classmethod def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeInternalEpk': - if data.get(cls.type_property_name) and data.get(cls.container_link_property_name): + if data.get(cls.type_property_name): feed_range = Range.ParseFromDict(data.get(cls.type_property_name)) - container_link = str(data.get(cls.container_link_property_name)) - return cls(feed_range, container_link) + return cls(feed_range) raise ValueError(f"Can not parse FeedRangeInternalEPK from the json," f" there is no property {cls.type_property_name}") diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 2afb583728f0..49e19c1a27d4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -55,10 +55,10 @@ def from_string(json_str: str) -> 'FeedRange': class FeedRangeEpk(FeedRange): type_property_name = "Range" - def __init__(self, feed_range: Range, container_link: str) -> None: + def __init__(self, feed_range: Range) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") - super().__init__(FeedRangeInternalEpk(feed_range, container_link)) + super().__init__(FeedRangeInternalEpk(feed_range)) def __str__(self) -> str: """Get a json representation of the feed range. @@ -71,5 +71,4 @@ def __str__(self) -> str: @classmethod def _from_json(cls, data: Dict[str, Any]) -> 'FeedRange': - feed_range_internal = FeedRangeInternalEpk.from_json(data) - return cls(feed_range_internal._range, feed_range_internal._container_link) + return cls(FeedRangeInternalEpk.from_json(data)._range) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index d31bdbf5f689..6d38bdd7a304 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -94,6 +94,7 @@ def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[s # [("AA", "DD"), "1:1#57,2:1#58"] # 3. [(("AA", "BB"), "1:1#57"), (("BB", "DD"), "2:1#52"), (("AA", "DD"), "0:1#55")] -> # [("AA", "DD"), "1:1#57,2:1#52,0:1#55"] +# goal here is to detect any obvious merges or splits that happened # compound session tokens are not considered will just pass them along def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> List[Tuple[Range, str]]: processed_ranges = [] @@ -162,17 +163,12 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li overlapping_ranges.remove(overlapping_ranges[0]) return processed_ranges -def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange, - container_link: str): +def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange): - if target_feed_range._feed_range_internal._container_link != container_link: - raise ValueError('The target feed range does not belong to the container.') target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges overlapping_ranges = [] for feed_range_to_session_token in feed_ranges_to_session_tokens: - if feed_range_to_session_token[0]._feed_range_internal._container_link != container_link: - raise ValueError('The feed range does not belong to the container.') if Range.overlaps(target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range()): overlapping_ranges.append((feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index a2c8ad456b24..f5e0dd50e9bd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1318,10 +1318,10 @@ async def read_feed_ranges( [Range("", "FF", True, False)], **kwargs) - return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange), self.container_link) + return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) for partitionKeyRange in partition_key_ranges] - def get_updated_session_token(self, + async def get_updated_session_token(self, feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange ) -> str: @@ -1337,7 +1337,7 @@ def get_updated_session_token(self, :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range, self.container_link) + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. @@ -1346,9 +1346,9 @@ async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) - :returns: a feed range :rtype: FeedRange """ - return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key), self.container_link) + return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key)) - def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. :param parent_feed_range: left feed range :type parent_feed_range: FeedRange @@ -1357,8 +1357,6 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - if (child_feed_range._feed_range_internal._container_link != self.container_link or - parent_feed_range._feed_range_internal._container_link != self.container_link): - raise ValueError("Feed ranges must be from the same container.") + return child_feed_range._feed_range_internal.get_normalized_range().is_subset( parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 213b50ae2f94..0c55de3802fb 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1388,7 +1388,7 @@ def read_feed_ranges( [Range("", "FF", True, False)], # default to full range **kwargs) - return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange), self.container_link) + return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) for partitionKeyRange in partition_key_ranges] def get_updated_session_token( @@ -1407,7 +1407,7 @@ def get_updated_session_token( :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range, self.container_link) + return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. @@ -1416,7 +1416,7 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Feed :returns: a feed range :rtype: FeedRange """ - return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key), self.container_link) + return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: """Checks if child feed range is a subset of parent feed range. @@ -1427,8 +1427,5 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - if (child_feed_range._feed_range_internal._container_link != self.container_link or - parent_feed_range._feed_range_internal._container_link != self.container_link): - raise ValueError("Feed ranges must be from the same container.") return child_feed_range._feed_range_internal.get_normalized_range().is_subset( parent_feed_range._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index f75b49520459..a4983b41ac19 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -98,13 +98,12 @@ def test_partition_key_to_feed_range(self, setup): @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_subset_ranges) def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset): - epk_parent_feed_range = FeedRangeEpk(parent_feed_range, setup["created_collection"].container_link) - epk_child_feed_range = FeedRangeEpk(child_feed_range, setup["created_collection"].container_link) + epk_parent_feed_range = FeedRangeEpk(parent_feed_range) + epk_child_feed_range = FeedRangeEpk(child_feed_range) assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) == is_subset def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) @@ -112,21 +111,5 @@ def test_feed_range_is_subset_from_pk(self, setup): def test_overlaps(self, setup, range1, range2, overlaps): assert Range.overlaps(range1, range2) == overlaps - def test_is_subset_with_wrong_feed_range(self, setup): - wrong_container = "wrong_container_link" - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - wrong_container) - epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) - with pytest.raises(ValueError, match="Feed ranges must be from the same container."): - setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) - epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), - wrong_container) - with pytest.raises(ValueError, match="Feed ranges must be from the same container."): - setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py index 231db3eac598..fa639b238558 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py @@ -53,26 +53,9 @@ async def test_partition_key_to_feed_range(self, setup): await setup["created_db"].delete_container(created_container) async def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) + epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) epk_child_feed_range = await setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - def test_is_subset_with_wrong_feed_range(self, setup): - wrong_container = "wrong_container_link" - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - wrong_container) - epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) - with pytest.raises(ValueError, match="Feed ranges must be from the same container."): - setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False), - setup["created_collection"].container_link) - epk_child_feed_range = FeedRangeEpk(Range("", "FF", True, False), - wrong_container) - with pytest.raises(ValueError, match="Feed ranges must be from the same container."): - setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) - if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 394422d7ed07..1adab22977c4 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -90,8 +90,7 @@ class TestSessionTokenHelpers: TEST_COLLECTION_ID = configs.TEST_SINGLE_PARTITION_CONTAINER_ID def test_get_session_token_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False), - setup[COLLECTION].container_link) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] session_token = "0:1#51#3=52" @@ -100,8 +99,7 @@ def test_get_session_token_update(self, setup): assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False), - setup[COLLECTION].container_link) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) @@ -113,18 +111,15 @@ def test_many_session_tokens_update_same_range(self, setup): assert updated_session_token == session_token def test_many_session_tokens_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False), - setup[COLLECTION].container_link) + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) feed_ranges_and_session_tokens.append((feed_range, session_token)) # adding irrelevant feed ranges - feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False), - setup[COLLECTION].container_link) - feed_range2 = FeedRangeEpk(Range("00", "55", True, False), - setup[COLLECTION].container_link) + feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False)) + feed_range2 = FeedRangeEpk(Range("00", "55", True, False)) for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) if i % 2 == 0: @@ -142,27 +137,11 @@ def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, e actual_split_ranges = [] for feed_range, session_token in split_ranges: actual_split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], - True, False), - setup[COLLECTION].container_link), session_token)) + True, False)), session_token)) target_feed_range = FeedRangeEpk(Range(target_feed_range[0], target_feed_range[1][1], - True, False), - setup[COLLECTION].container_link) + True, False)) updated_session_token = setup[COLLECTION].get_updated_session_token(actual_split_ranges, target_feed_range) assert updated_session_token == expected_session_token - def test_invalid_feed_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False), - setup[COLLECTION].container_link) - session_token = "0:1#54#3=50" - feed_ranges_and_session_tokens = [(feed_range, session_token)] - with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): - setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, - FeedRangeEpk(Range( - "CC", - "FF", - True, - False), - setup[COLLECTION].container_link)) - if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py index 3b8d0eb115cb..9ad13ca0cf7f 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py @@ -62,7 +62,7 @@ async def test_updated_session_token_from_logical_pk(self): target_session_token, previous_session_token = await self.create_items_logical_pk(container, target_feed_range, previous_session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token feed_ranges_and_session_tokens.append((target_feed_range, session_token)) @@ -72,7 +72,7 @@ async def test_updated_session_token_from_logical_pk(self): target_session_token, _ = await self.create_items_logical_pk(container, target_feed_range, session_token, feed_ranges_and_session_tokens) target_feed_range = await container.feed_range_from_partition_key(target_pk) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) @@ -88,7 +88,7 @@ async def test_updated_session_token_from_physical_pk(self): previous_session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.trigger_split(container, 11000) @@ -97,7 +97,7 @@ async def test_updated_session_token_from_physical_pk(self): session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert is_compound_session_token(session_token) session_tokens = session_token.split(",") assert len(session_tokens) == 2 @@ -124,7 +124,7 @@ async def test_updated_session_token_hpk(self): feed_ranges_and_session_tokens, True) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) @@ -141,7 +141,7 @@ async def test_updated_session_token_logical_hpk(self): previous_session_token, feed_ranges_and_session_tokens, True) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) @@ -191,7 +191,7 @@ async def create_items_physical_pk(container, pk_feed_range, previous_session_to container_feed_ranges = await container.read_feed_ranges() target_feed_range = None for feed_range in container_feed_ranges: - if container.is_feed_range_subset(feed_range, pk_feed_range): + if await container.is_feed_range_subset(feed_range, pk_feed_range): target_feed_range = feed_range break @@ -204,7 +204,7 @@ async def create_items_physical_pk(container, pk_feed_range, previous_session_to curr_feed_range = await container.feed_range_from_partition_key(pk) else: curr_feed_range = await container.feed_range_from_partition_key(item['pk']) - if container.is_feed_range_subset(target_feed_range, curr_feed_range): + if await container.is_feed_range_subset(target_feed_range, curr_feed_range): target_session_token = session_token previous_session_token = session_token feed_ranges_and_session_tokens.append((curr_feed_range, session_token)) From 5d7b978d30fb07087d81b2cce919d633ab45ed43 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 15 Oct 2024 15:19:09 -0700 Subject: [PATCH 45/59] reacting to comments --- sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py | 10 +++------- .../azure/cosmos/_session_token_helpers.py | 13 +++++++------ .../azure-cosmos/azure/cosmos/aio/_container.py | 6 ++++-- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 7 +++++-- .../azure-cosmos/test/test_session_token_helpers.py | 12 ++++++++++++ 5 files changed, 31 insertions(+), 17 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py index 49e19c1a27d4..2bda669b6bc0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.py @@ -24,7 +24,7 @@ from abc import ABC from typing import Any, Dict -from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk, FeedRangeInternal +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._routing.routing_range import Range # pylint: disable=protected-access @@ -32,10 +32,6 @@ class FeedRange(ABC): """Represents a single feed range in an Azure Cosmos DB SQL API container. """ - - def __init__(self, feed_range_internal: FeedRangeInternal) -> None: - self._feed_range_internal = feed_range_internal - @staticmethod def from_string(json_str: str) -> 'FeedRange': """ @@ -58,14 +54,14 @@ class FeedRangeEpk(FeedRange): def __init__(self, feed_range: Range) -> None: if feed_range is None: raise ValueError("feed_range cannot be None") - super().__init__(FeedRangeInternalEpk(feed_range)) + + self._feed_range_internal = FeedRangeInternalEpk(feed_range) def __str__(self) -> str: """Get a json representation of the feed range. The returned json string can be used to create a new feed range from it. :return: A json representation of the feed range. - :rtype: str """ return self._feed_range_internal.__str__() diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 6d38bdd7a304..7bc5bfd94f2e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -21,11 +21,11 @@ """Internal Helper functions for manipulating session tokens. """ -from typing import Tuple, List +from typing import Tuple, List, cast from azure.cosmos._routing.routing_range import Range from azure.cosmos._vector_session_token import VectorSessionToken -from ._feed_range import FeedRange +from ._feed_range import FeedRange, FeedRangeEpk # pylint: disable=protected-access @@ -164,14 +164,15 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li return processed_ranges def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange): - - target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range() + target_feed_range_epk = cast(FeedRangeEpk, target_feed_range) + target_feed_range_normalized = target_feed_range_epk._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges overlapping_ranges = [] for feed_range_to_session_token in feed_ranges_to_session_tokens: + feed_range_epk = cast(FeedRangeEpk, feed_range_to_session_token[0]) if Range.overlaps(target_feed_range_normalized, - feed_range_to_session_token[0]._feed_range_internal.get_normalized_range()): - overlapping_ranges.append((feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), + feed_range_epk._feed_range_internal.get_normalized_range()): + overlapping_ranges.append((feed_range_epk._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1])) if len(overlapping_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 f5e0dd50e9bd..7bbc6de6fcee 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1357,6 +1357,8 @@ async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_ra :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ + child_feed_range_epk = cast(FeedRangeEpk, child_feed_range) + parent_feed_range_epk = cast(FeedRangeEpk, parent_feed_range) - return child_feed_range._feed_range_internal.get_normalized_range().is_subset( - parent_feed_range._feed_range_internal.get_normalized_range()) + return child_feed_range_epk._feed_range_internal.get_normalized_range().is_subset( + parent_feed_range_epk._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 0c55de3802fb..510445aab68a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1427,5 +1427,8 @@ def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: F :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - return child_feed_range._feed_range_internal.get_normalized_range().is_subset( - parent_feed_range._feed_range_internal.get_normalized_range()) + child_feed_range_epk = cast(FeedRangeEpk, child_feed_range) + parent_feed_range_epk = cast(FeedRangeEpk, parent_feed_range) + + return child_feed_range_epk._feed_range_internal.get_normalized_range().is_subset( + parent_feed_range_epk._feed_range_internal.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 1adab22977c4..54244cc40290 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -143,5 +143,17 @@ def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, e updated_session_token = setup[COLLECTION].get_updated_session_token(actual_split_ranges, target_feed_range) assert updated_session_token == expected_session_token + def test_invalid_feed_range(self, setup): + feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + session_token = "0:1#54#3=50" + feed_ranges_and_session_tokens = [(feed_range, session_token)] + with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): + setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, + FeedRangeEpk(Range( + "CC", + "FF", + True, + False))) + if __name__ == '__main__': unittest.main() From d54992f0502417ebf76be85563c6c66a39d9bcf0 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 15 Oct 2024 16:27:44 -0700 Subject: [PATCH 46/59] fix cspell --- sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py | 2 +- sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 7bc5bfd94f2e..6e802519b5b1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -37,7 +37,7 @@ def merge_session_tokens_with_same_range(session_token1: str, session_token2: st pk_range_id1, vector_session_token1 = parse_session_token(session_token1) pk_range_id2, vector_session_token2 = parse_session_token(session_token2) pk_range_id = pk_range_id1 - # The pkrangeid could be different in this scenario + # The partition key range id could be different in this scenario # Ex. get_updated_session_token([("AA", "BB"), "1:1#51"], ("AA", "DD")) -> "1:1#51" # Then we input this back into get_updated_session_token after a merge happened # get_updated_session_token([("AA", "DD"), "1:1#51", ("AA", "DD"), "0:1#55"], ("AA", "DD")) -> "0:1#55" diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 54244cc40290..abe491384f91 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -33,7 +33,7 @@ def create_split_ranges(): return [ # split with two children ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52"), (("BB", "DD"),"2:1#54#3=52")], ("AA", "DD"), "1:1#55#3=52,2:1#54#3=52"), - # same range different pkrangeid + # same range different partition key range ids ([(("AA", "DD"), "1:1#51#3=52"), (("AA", "DD"),"0:1#55#3=52")], ("AA", "DD"), "0:1#55#3=52"), # split with one child From fa16830db9d1cceceafcf92c416b9ca6c02712c9 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 16 Oct 2024 12:26:10 -0700 Subject: [PATCH 47/59] rename method to get_latest_session_token --- .../azure/cosmos/_session_token_helpers.py | 2 +- .../azure/cosmos/aio/_container.py | 20 ++++++++--------- .../azure-cosmos/azure/cosmos/container.py | 14 ++++++------ ..._token.py => test_latest_session_token.py} | 22 +++++++++---------- ....py => test_latest_session_token_async.py} | 22 +++++++++---------- .../test/test_session_token_helpers.py | 16 +++++++------- 6 files changed, 48 insertions(+), 48 deletions(-) rename sdk/cosmos/azure-cosmos/test/{test_updated_session_token.py => test_latest_session_token.py} (91%) rename sdk/cosmos/azure-cosmos/test/{test_updated_session_token_async.py => test_latest_session_token_async.py} (91%) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 6e802519b5b1..5df3c8a93da0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -163,7 +163,7 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li overlapping_ranges.remove(overlapping_ranges[0]) return processed_ranges -def get_updated_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange): +def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange): target_feed_range_epk = cast(FeedRangeEpk, target_feed_range) target_feed_range_normalized = target_feed_range_epk._feed_range_internal.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 7bbc6de6fcee..e4e7ed775ef7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -43,7 +43,7 @@ ) from .._feed_range import FeedRange, FeedRangeEpk from .._routing.routing_range import Range -from .._session_token_helpers import get_updated_session_token +from .._session_token_helpers import get_latest_session_token from ..offer import ThroughputProperties from ..partition_key import ( NonePartitionKeyValue, @@ -1321,15 +1321,15 @@ async def read_feed_ranges( return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) for partitionKeyRange in partition_key_ranges] - async def get_updated_session_token(self, - feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], - target_feed_range: FeedRange - ) -> str: + async def get_latest_session_token(self, + feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], + target_feed_range: FeedRange + ) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples - for a specific target feed range. The feed range can be obtained from a response from any crud operation. - This should only be used if maintaining own session token or else the sdk will keep track of - session token. Session tokens and feed ranges are scoped to a container. Only input session tokens - and feed ranges obtained from the same container. + for a specific target feed range. The feed range can be obtained from a logical partition or by reading the + container feed ranges. This should only be used if maintaining own session token or else the sdk will + keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session + tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]] :param target_feed_range: feed range to get most up to date session token. @@ -1337,7 +1337,7 @@ async def get_updated_session_token(self, :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 510445aab68a..8acb2e228e44 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -41,7 +41,7 @@ from ._cosmos_client_connection import CosmosClientConnection from ._feed_range import FeedRange, FeedRangeEpk from ._routing.routing_range import Range -from ._session_token_helpers import get_updated_session_token +from ._session_token_helpers import get_latest_session_token from .offer import Offer, ThroughputProperties from .partition_key import ( NonePartitionKeyValue, @@ -1391,15 +1391,15 @@ def read_feed_ranges( return [FeedRangeEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)) for partitionKeyRange in partition_key_ranges] - def get_updated_session_token( + def get_latest_session_token( self, feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples - for a specific target feed range. The feed range can be obtained from a response from any crud operation. - This should only be used if maintaining own session token or else the sdk will keep track of - session token. Session tokens and feed ranges are scoped to a container. Only input session tokens - and feed ranges obtained from the same container. + for a specific target feed range. The feed range can be obtained from a logical partition or by reading the + container feed ranges. This should only be used if maintaining own session token or else the sdk will + keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session + tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]] :param target_feed_range: feed range to get most up to date session token. @@ -1407,7 +1407,7 @@ def get_updated_session_token( :returns: a session token :rtype: str """ - return get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range) + return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: """Gets the feed range for a given partition key. diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py similarity index 91% rename from sdk/cosmos/azure-cosmos/test/test_updated_session_token.py rename to sdk/cosmos/azure-cosmos/test/test_latest_session_token.py index efbf9a9ad890..d1ae160dc7a4 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py @@ -32,7 +32,7 @@ def create_item(hpk): return item -class TestUpdatedSessionToken(unittest.TestCase): +class TestLatestSessionToken(unittest.TestCase): """Test for session token helpers""" created_db: DatabaseProxy = None @@ -48,7 +48,7 @@ def setUpClass(cls): cls.database = cls.client.get_database_client(cls.TEST_DATABASE_ID) - def test_updated_session_token_from_logical_pk(self): + def test_latest_session_token_from_logical_pk(self): container = self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), PartitionKey(path="/pk"), offer_throughput=400) @@ -59,7 +59,7 @@ def test_updated_session_token_from_logical_pk(self): target_session_token, previous_session_token = self.create_items_logical_pk(container, target_feed_range, previous_session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token feed_ranges_and_session_tokens.append((target_feed_range, session_token)) @@ -69,12 +69,12 @@ def test_updated_session_token_from_logical_pk(self): target_session_token, _ = self.create_items_logical_pk(container, target_feed_range, session_token, feed_ranges_and_session_tokens) target_feed_range = container.feed_range_from_partition_key(target_pk) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token self.database.delete_container(container.id) - def test_updated_session_token_from_physical_pk(self): + def test_latest_session_token_from_physical_pk(self): container = self.database.create_container("test_updated_session_token_from_physical_pk" + str(uuid.uuid4()), PartitionKey(path="/pk"), offer_throughput=400) @@ -85,7 +85,7 @@ def test_updated_session_token_from_physical_pk(self): previous_session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token self.trigger_split(container, 11000) @@ -94,7 +94,7 @@ def test_updated_session_token_from_physical_pk(self): session_token, feed_ranges_and_session_tokens) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert is_compound_session_token(session_token) session_tokens = session_token.split(",") assert len(session_tokens) == 2 @@ -107,7 +107,7 @@ def test_updated_session_token_from_physical_pk(self): assert '2' in pk_range_ids self.database.delete_container(container.id) - def test_updated_session_token_hpk(self): + def test_latest_session_token_hpk(self): container = self.database.create_container("test_updated_session_token_hpk" + str(uuid.uuid4()), PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), offer_throughput=400) @@ -121,12 +121,12 @@ def test_updated_session_token_hpk(self): feed_ranges_and_session_tokens, True) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token self.database.delete_container(container.id) - def test_updated_session_token_logical_hpk(self): + def test_latest_session_token_logical_hpk(self): container = self.database.create_container("test_updated_session_token_from_logical_hpk" + str(uuid.uuid4()), PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), offer_throughput=400) @@ -138,7 +138,7 @@ def test_updated_session_token_logical_hpk(self): previous_session_token, feed_ranges_and_session_tokens, True) - session_token = container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token self.database.delete_container(container.id) diff --git a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py similarity index 91% rename from sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py rename to sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py index 9ad13ca0cf7f..39917286ca24 100644 --- a/sdk/cosmos/azure-cosmos/test/test_updated_session_token_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py @@ -33,7 +33,7 @@ def create_item(hpk): return item -class TestUpdatedSessionTokenAsync(unittest.IsolatedAsyncioTestCase): +class TestLatestSessionTokenAsync(unittest.IsolatedAsyncioTestCase): """Test for session token helpers""" created_db: DatabaseProxy = None @@ -51,7 +51,7 @@ async def tearDown(self): await self.client.delete_database(self.database.id) await self.client.close() - async def test_updated_session_token_from_logical_pk(self): + async def test_latest_session_token_from_logical_pk(self): container = await self.database.create_container("test_updated_session_token_from_logical_pk" + str(uuid.uuid4()), PartitionKey(path="/pk"), offer_throughput=400) @@ -62,7 +62,7 @@ async def test_updated_session_token_from_logical_pk(self): target_session_token, previous_session_token = await self.create_items_logical_pk(container, target_feed_range, previous_session_token, feed_ranges_and_session_tokens) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token feed_ranges_and_session_tokens.append((target_feed_range, session_token)) @@ -72,12 +72,12 @@ async def test_updated_session_token_from_logical_pk(self): target_session_token, _ = await self.create_items_logical_pk(container, target_feed_range, session_token, feed_ranges_and_session_tokens) target_feed_range = await container.feed_range_from_partition_key(target_pk) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) - async def test_updated_session_token_from_physical_pk(self): + async def test_latest_session_token_from_physical_pk(self): container = await self.database.create_container("test_updated_session_token_from_physical_pk" + str(uuid.uuid4()), PartitionKey(path="/pk"), offer_throughput=400) @@ -88,7 +88,7 @@ async def test_updated_session_token_from_physical_pk(self): previous_session_token, feed_ranges_and_session_tokens) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.trigger_split(container, 11000) @@ -97,7 +97,7 @@ async def test_updated_session_token_from_physical_pk(self): session_token, feed_ranges_and_session_tokens) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert is_compound_session_token(session_token) session_tokens = session_token.split(",") assert len(session_tokens) == 2 @@ -110,7 +110,7 @@ async def test_updated_session_token_from_physical_pk(self): assert '2' in pk_range_ids await self.database.delete_container(container.id) - async def test_updated_session_token_hpk(self): + async def test_latest_session_token_hpk(self): container = await self.database.create_container("test_updated_session_token_hpk" + str(uuid.uuid4()), PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), offer_throughput=400) @@ -124,12 +124,12 @@ async def test_updated_session_token_hpk(self): feed_ranges_and_session_tokens, True) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) - async def test_updated_session_token_logical_hpk(self): + async def test_latest_session_token_logical_hpk(self): container = await self.database.create_container("test_updated_session_token_from_logical_hpk" + str(uuid.uuid4()), PartitionKey(path=["/state", "/city", "/zipcode"], kind="MultiHash"), offer_throughput=400) @@ -141,7 +141,7 @@ async def test_updated_session_token_logical_hpk(self): previous_session_token, feed_ranges_and_session_tokens, True) - session_token = await container.get_updated_session_token(feed_ranges_and_session_tokens, target_feed_range) + session_token = await container.get_latest_session_token(feed_ranges_and_session_tokens, target_feed_range) assert session_token == target_session_token await self.database.delete_container(container.id) diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index abe491384f91..17be19e7f597 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -95,7 +95,7 @@ def test_get_session_token_update(self, setup): feed_ranges_and_session_tokens = [(feed_range, session_token)] session_token = "0:1#51#3=52" feed_ranges_and_session_tokens.append((feed_range, session_token)) - session_token = setup[COLLECTION].get_updated_session_token(feed_ranges_and_session_tokens, feed_range) + session_token = setup[COLLECTION].get_latest_session_token(feed_ranges_and_session_tokens, feed_range) assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): @@ -106,8 +106,8 @@ def test_many_session_tokens_update_same_range(self, setup): feed_ranges_and_session_tokens.append((feed_range, session_token)) session_token = "0:1#101#3=101" feed_ranges_and_session_tokens.append((feed_range, session_token)) - updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, - feed_range) + updated_session_token = setup["created_collection"].get_latest_session_token(feed_ranges_and_session_tokens, + feed_range) assert updated_session_token == session_token def test_many_session_tokens_update(self, setup): @@ -128,8 +128,8 @@ def test_many_session_tokens_update(self, setup): feed_ranges_and_session_tokens.append((feed_range2, session_token)) session_token = "0:1#101#3=101" feed_ranges_and_session_tokens.append((feed_range, session_token)) - updated_session_token = setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, - feed_range) + updated_session_token = setup["created_collection"].get_latest_session_token(feed_ranges_and_session_tokens, + feed_range) assert updated_session_token == session_token @pytest.mark.parametrize("split_ranges, target_feed_range, expected_session_token", create_split_ranges()) @@ -140,7 +140,7 @@ def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, e True, False)), session_token)) target_feed_range = FeedRangeEpk(Range(target_feed_range[0], target_feed_range[1][1], True, False)) - updated_session_token = setup[COLLECTION].get_updated_session_token(actual_split_ranges, target_feed_range) + updated_session_token = setup[COLLECTION].get_latest_session_token(actual_split_ranges, target_feed_range) assert updated_session_token == expected_session_token def test_invalid_feed_range(self, setup): @@ -148,8 +148,8 @@ def test_invalid_feed_range(self, setup): session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): - setup["created_collection"].get_updated_session_token(feed_ranges_and_session_tokens, - FeedRangeEpk(Range( + setup["created_collection"].get_latest_session_token(feed_ranges_and_session_tokens, + FeedRangeEpk(Range( "CC", "FF", True, From 6914a2037c270eb0ee9954e1f9f7f4e354a577f5 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 16 Oct 2024 17:59:56 -0700 Subject: [PATCH 48/59] reacting to reverted feed range --- .../azure/cosmos/_session_token_helpers.py | 23 +++++++----- .../azure/cosmos/aio/_container.py | 31 +++++++++------- .../azure-cosmos/azure/cosmos/container.py | 35 +++++++++++-------- .../azure-cosmos/test/test_feed_range.py | 18 ++++++---- .../test/test_feed_range_async.py | 12 ++++--- .../test/test_latest_session_token.py | 14 ++++++-- .../test/test_latest_session_token_async.py | 14 ++++++-- .../test/test_session_token_helpers.py | 32 ++++++++++------- 8 files changed, 114 insertions(+), 65 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 5df3c8a93da0..055f8a35eb5b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -21,11 +21,13 @@ """Internal Helper functions for manipulating session tokens. """ -from typing import Tuple, List, cast +import base64 +import json +from typing import Tuple, List from azure.cosmos._routing.routing_range import Range from azure.cosmos._vector_session_token import VectorSessionToken -from ._feed_range import FeedRange, FeedRangeEpk +from ._change_feed.feed_range_internal import FeedRangeInternalEpk # pylint: disable=protected-access @@ -163,16 +165,21 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li overlapping_ranges.remove(overlapping_ranges[0]) return processed_ranges -def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], target_feed_range: FeedRange): - target_feed_range_epk = cast(FeedRangeEpk, target_feed_range) - target_feed_range_normalized = target_feed_range_epk._feed_range_internal.get_normalized_range() +def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[str, str]], target_feed_range: str): + + target_feed_range_str = base64.b64decode(target_feed_range).decode('utf-8') + feed_range_json = json.loads(target_feed_range_str) + target_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + target_feed_range_normalized = target_feed_range_epk.get_normalized_range() # filter out tuples that overlap with target_feed_range and normalizes all the ranges overlapping_ranges = [] for feed_range_to_session_token in feed_ranges_to_session_tokens: - feed_range_epk = cast(FeedRangeEpk, feed_range_to_session_token[0]) + feed_range_str = base64.b64decode(feed_range_to_session_token[0]).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) if Range.overlaps(target_feed_range_normalized, - feed_range_epk._feed_range_internal.get_normalized_range()): - overlapping_ranges.append((feed_range_epk._feed_range_internal.get_normalized_range(), + feed_range_epk.get_normalized_range()): + overlapping_ranges.append((feed_range_epk.get_normalized_range(), feed_range_to_session_token[1])) if len(overlapping_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 81392c6e585f..2653596b8da2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -21,6 +21,8 @@ """Create, read, update and delete items in the Azure Cosmos DB SQL API service. """ +import base64 +import json import warnings from datetime import datetime from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload @@ -1320,8 +1322,8 @@ async def read_feed_ranges( for partitionKeyRange in partition_key_ranges] async def get_latest_session_token(self, - feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], - target_feed_range: FeedRange + feed_ranges_to_session_tokens: List[Tuple[str, str]], + target_feed_range: str ) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a logical partition or by reading the @@ -1337,26 +1339,29 @@ async def get_latest_session_token(self, """ return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) - async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: + async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> str: """Gets the feed range for a given partition key. :param partition_key: partition key to get feed range. :type partition_key: PartitionKey :returns: a feed range - :rtype: FeedRange + :rtype: str """ - return FeedRangeEpk(await self._get_epk_range_for_partition_key(partition_key)) + return FeedRangeInternalEpk(await self._get_epk_range_for_partition_key(partition_key)).__str__() - async def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + async def is_feed_range_subset(self, parent_feed_range: str, child_feed_range: str) -> bool: """Checks if child feed range is a subset of parent feed range. :param parent_feed_range: left feed range - :type parent_feed_range: FeedRange + :type parent_feed_range: str :param child_feed_range: right feed range - :type child_feed_range: FeedRange + :type child_feed_range: str :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - child_feed_range_epk = cast(FeedRangeEpk, child_feed_range) - parent_feed_range_epk = cast(FeedRangeEpk, parent_feed_range) - - return child_feed_range_epk._feed_range_internal.get_normalized_range().is_subset( - parent_feed_range_epk._feed_range_internal.get_normalized_range()) + parent_feed_range_str = base64.b64decode(parent_feed_range).decode('utf-8') + feed_range_json = json.loads(parent_feed_range_str) + parent_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + child_feed_range_str = base64.b64decode(child_feed_range).decode('utf-8') + feed_range_json = json.loads(child_feed_range_str) + child_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + return child_feed_range_epk.get_normalized_range().is_subset( + parent_feed_range_epk.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index dd2f9f6c8ae4..b55b01e8d958 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -21,6 +21,8 @@ """Create, read, update and delete items in the Azure Cosmos DB SQL API service. """ +import base64 +import json import warnings from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload @@ -1390,42 +1392,45 @@ def read_feed_ranges( def get_latest_session_token( self, - feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]], - target_feed_range: FeedRange) -> str: + feed_ranges_to_session_tokens: List[Tuple[str, str]], + target_feed_range: str) -> str: """Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a logical partition or by reading the container feed ranges. This should only be used if maintaining own session token or else the sdk will keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. - :type feed_ranges_to_session_tokens: List[Tuple[FeedRange, str]] + :type feed_ranges_to_session_tokens: List[Tuple[str, str]] :param target_feed_range: feed range to get most up to date session token. - :type target_feed_range: FeedRange + :type target_feed_range: str :returns: a session token :rtype: str """ return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) - def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> FeedRange: + def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> str: """Gets the feed range for a given partition key. :param partition_key: partition key to get feed range. :type partition_key: PartitionKey :returns: a feed range - :rtype: FeedRange + :rtype: str """ - return FeedRangeEpk(self._get_epk_range_for_partition_key(partition_key)) + return FeedRangeInternalEpk(self._get_epk_range_for_partition_key(partition_key)).__str__() - def is_feed_range_subset(self, parent_feed_range: FeedRange, child_feed_range: FeedRange) -> bool: + def is_feed_range_subset(self, parent_feed_range: str, child_feed_range: str) -> bool: """Checks if child feed range is a subset of parent feed range. :param parent_feed_range: left feed range - :type parent_feed_range: FeedRange + :type parent_feed_range: str :param child_feed_range: right feed range - :type child_feed_range: FeedRange + :type child_feed_range: str :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool """ - child_feed_range_epk = cast(FeedRangeEpk, child_feed_range) - parent_feed_range_epk = cast(FeedRangeEpk, parent_feed_range) - - return child_feed_range_epk._feed_range_internal.get_normalized_range().is_subset( - parent_feed_range_epk._feed_range_internal.get_normalized_range()) + parent_feed_range_str = base64.b64decode(parent_feed_range).decode('utf-8') + feed_range_json = json.loads(parent_feed_range_str) + parent_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + child_feed_range_str = base64.b64decode(child_feed_range).decode('utf-8') + feed_range_json = json.loads(child_feed_range_str) + child_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + return child_feed_range_epk.get_normalized_range().is_subset( + parent_feed_range_epk.get_normalized_range()) diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index a4983b41ac19..bf958f3b97e6 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import base64 +import json import unittest import uuid @@ -9,7 +10,8 @@ import azure.cosmos.cosmos_client as cosmos_client import azure.cosmos.partition_key as partition_key import test_config -from azure.cosmos._feed_range import FeedRangeEpk + +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._routing.routing_range import Range @pytest.fixture(scope="class") @@ -92,18 +94,22 @@ def test_partition_key_to_feed_range(self, setup): partition_key=partition_key.PartitionKey(path="/id") ) feed_range = created_container.feed_range_from_partition_key("1") - assert feed_range._feed_range_internal.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", + feed_range_str = base64.b64decode(feed_range).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + assert feed_range_epk.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) setup["created_db"].delete_container(created_container) @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_subset_ranges) def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset): - epk_parent_feed_range = FeedRangeEpk(parent_feed_range) - epk_child_feed_range = FeedRangeEpk(child_feed_range) + epk_parent_feed_range = FeedRangeInternalEpk(parent_feed_range).__str__() + epk_child_feed_range = FeedRangeInternalEpk(child_feed_range).__str__() assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) == is_subset def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) + epk_parent_feed_range = FeedRangeInternalEpk( + Range("", "FF", True, False)).__str__() epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py index fa639b238558..0ace2113a9a6 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range_async.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import base64 +import json import unittest import uuid @@ -9,7 +10,7 @@ import azure.cosmos.partition_key as partition_key import test_config -from azure.cosmos._feed_range import FeedRangeEpk +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._routing.routing_range import Range from azure.cosmos.aio import CosmosClient @@ -48,12 +49,15 @@ async def test_partition_key_to_feed_range(self, setup): partition_key=partition_key.PartitionKey(path="/id") ) feed_range = await created_container.feed_range_from_partition_key("1") - assert feed_range._feed_range_internal.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", + feed_range_str = base64.b64decode(feed_range).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + assert feed_range_epk.get_normalized_range() == Range("3C80B1B7310BB39F29CC4EA05BDD461E", "3c80b1b7310bb39f29cc4ea05bdd461f", True, False) await setup["created_db"].delete_container(created_container) async def test_feed_range_is_subset_from_pk(self, setup): - epk_parent_feed_range = FeedRangeEpk(Range("", "FF", True, False)) + epk_parent_feed_range = FeedRangeInternalEpk(Range("", "FF", True, False)).__str__() epk_child_feed_range = await setup["created_collection"].feed_range_from_partition_key("1") assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range) diff --git a/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py index d1ae160dc7a4..9e61ad928316 100644 --- a/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import base64 +import json import random import time import unittest @@ -10,6 +11,7 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos import DatabaseProxy, PartitionKey +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._session_token_helpers import is_compound_session_token, parse_session_token from azure.cosmos.http_constants import HttpHeaders @@ -174,8 +176,14 @@ def create_items_logical_pk(container, target_pk_range, previous_session_token, session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] pk_range = container.feed_range_from_partition_key(pk) - if (pk_range._feed_range_internal.get_normalized_range() == - target_pk_range._feed_range_internal.get_normalized_range()): + pk_feed_range_str = base64.b64decode(pk_range).decode('utf-8') + feed_range_json = json.loads(pk_feed_range_str) + pk_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + target_feed_range_str = base64.b64decode(target_pk_range).decode('utf-8') + feed_range_json = json.loads(target_feed_range_str) + target_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + if (pk_feed_range_epk.get_normalized_range() == + target_feed_range_epk.get_normalized_range()): target_session_token = session_token previous_session_token = session_token feed_ranges_and_session_tokens.append((pk_range, diff --git a/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py index 39917286ca24..b20bf41983ea 100644 --- a/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. - +import base64 +import json import random import time import unittest @@ -9,6 +10,7 @@ import test_config from azure.cosmos import PartitionKey +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._session_token_helpers import is_compound_session_token, parse_session_token from azure.cosmos.aio import DatabaseProxy from azure.cosmos.aio import CosmosClient @@ -177,8 +179,14 @@ async def create_items_logical_pk(container, target_pk_range, previous_session_t session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] pk_feed_range = await container.feed_range_from_partition_key(pk) - if (pk_feed_range._feed_range_internal.get_normalized_range() == - target_pk_range._feed_range_internal.get_normalized_range()): + pk_feed_range_str = base64.b64decode(pk_feed_range).decode('utf-8') + feed_range_json = json.loads(pk_feed_range_str) + pk_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + target_feed_range_str = base64.b64decode(target_pk_range).decode('utf-8') + feed_range_json = json.loads(target_feed_range_str) + target_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json) + if (pk_feed_range_epk.get_normalized_range() == + target_feed_range_epk.get_normalized_range()): target_session_token = session_token previous_session_token = session_token feed_ranges_and_session_tokens.append((pk_feed_range, diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 17be19e7f597..3b2fe578ee66 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -9,7 +9,7 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos import DatabaseProxy -from azure.cosmos._feed_range import FeedRangeEpk +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk from azure.cosmos._routing.routing_range import Range COLLECTION = "created_collection" @@ -90,7 +90,8 @@ class TestSessionTokenHelpers: TEST_COLLECTION_ID = configs.TEST_SINGLE_PARTITION_CONTAINER_ID def test_get_session_token_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeInternalEpk( + Range("AA", "BB", True, False)).__str__() session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] session_token = "0:1#51#3=52" @@ -99,7 +100,8 @@ def test_get_session_token_update(self, setup): assert session_token == "0:1#54#3=52" def test_many_session_tokens_update_same_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeInternalEpk( + Range("AA", "BB", True, False)).__str__() feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) @@ -111,15 +113,18 @@ def test_many_session_tokens_update_same_range(self, setup): assert updated_session_token == session_token def test_many_session_tokens_update(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeInternalEpk( + Range("AA", "BB", True, False)).__str__() feed_ranges_and_session_tokens = [] for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) feed_ranges_and_session_tokens.append((feed_range, session_token)) # adding irrelevant feed ranges - feed_range1 = FeedRangeEpk(Range("CC", "FF", True, False)) - feed_range2 = FeedRangeEpk(Range("00", "55", True, False)) + feed_range1 = FeedRangeInternalEpk( + Range("CC", "FF", True, False)).__str__() + feed_range2 = FeedRangeInternalEpk( + Range("00", "55", True, False)).__str__() for i in range(1000): session_token = "0:1#" + str(random.randint(1, 100)) + "#3=" + str(random.randint(1, 100)) if i % 2 == 0: @@ -136,24 +141,25 @@ def test_many_session_tokens_update(self, setup): def test_simulated_splits_merges(self, setup, split_ranges, target_feed_range, expected_session_token): actual_split_ranges = [] for feed_range, session_token in split_ranges: - actual_split_ranges.append((FeedRangeEpk(Range(feed_range[0], feed_range[1], - True, False)), session_token)) - target_feed_range = FeedRangeEpk(Range(target_feed_range[0], target_feed_range[1][1], - True, False)) + actual_split_ranges.append((FeedRangeInternalEpk(Range(feed_range[0], feed_range[1], + True, False)).__str__(), session_token)) + target_feed_range = FeedRangeInternalEpk(Range(target_feed_range[0], target_feed_range[1][1], + True, False)).__str__() updated_session_token = setup[COLLECTION].get_latest_session_token(actual_split_ranges, target_feed_range) assert updated_session_token == expected_session_token def test_invalid_feed_range(self, setup): - feed_range = FeedRangeEpk(Range("AA", "BB", True, False)) + feed_range = FeedRangeInternalEpk( + Range("AA", "BB", True, False)).__str__() session_token = "0:1#54#3=50" feed_ranges_and_session_tokens = [(feed_range, session_token)] with pytest.raises(ValueError, match='There were no overlapping feed ranges with the target.'): setup["created_collection"].get_latest_session_token(feed_ranges_and_session_tokens, - FeedRangeEpk(Range( + FeedRangeInternalEpk(Range( "CC", "FF", True, - False))) + False)).__str__()) if __name__ == '__main__': unittest.main() From ab9723a79d4df945247b8d88b29b9852573b8c62 Mon Sep 17 00:00:00 2001 From: annie-mac Date: Wed, 23 Oct 2024 15:37:59 -0700 Subject: [PATCH 49/59] change based on the api review --- .../cosmos/_change_feed/change_feed_state.py | 4 +--- .../azure/cosmos/aio/_container.py | 20 ++++++++++++------- .../azure-cosmos/azure/cosmos/container.py | 19 +++++++++++------- sdk/cosmos/azure-cosmos/samples/examples.py | 3 +-- .../azure-cosmos/samples/examples_async.py | 3 +-- .../azure-cosmos/test/test_change_feed.py | 4 ++-- .../test/test_change_feed_async.py | 4 ++-- 7 files changed, 32 insertions(+), 25 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 1f7d63f96d72..88a05fd7d818 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -383,9 +383,7 @@ def from_initial_state( feed_range: Optional[FeedRangeInternal] = None if change_feed_state_context.get("feedRange"): - feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') - feed_range_json = json.loads(feed_range_str) - feed_range = FeedRangeInternalEpk.from_json(feed_range_json) + feed_range = FeedRangeInternalEpk.from_json(change_feed_state_context["feedRange"]) elif change_feed_state_context.get("partitionKey"): if change_feed_state_context.get("partitionKeyFeedRange"): feed_range =\ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 1434ea68d5ea..f88fbe94850f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -23,7 +23,7 @@ """ import warnings from datetime import datetime -from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload +from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload, Iterable from typing_extensions import Literal from azure.core import MatchConditions @@ -534,7 +534,7 @@ def query_items_change_feed( def query_items_change_feed( self, *, - feed_range: str, + feed_range: Dict[str, Any], max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, @@ -542,7 +542,7 @@ def query_items_change_feed( ) -> AsyncItemPaged[Dict[str, Any]]: """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. @@ -620,7 +620,7 @@ def query_items_change_feed( # pylint: disable=unused-argument """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. :keyword partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] @@ -1296,13 +1296,17 @@ async def read_feed_ranges( *, force_refresh: Optional[bool] = False, **kwargs: Any - ) -> List[str]: + ) -> Iterable[Dict[str, Any]]: """ Obtains a list of feed ranges that can be used to parallelize feed operations. :keyword bool force_refresh: Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. :returns: A list representing the feed ranges in base64 encoded string - :rtype: List[str] + :rtype: Iterable[Dict[str, Any]] + + .. note:: + For each feed range, even through a Dict has been returned, + but in the future, the structure may change. Please just treat it as opaque and do not take any dependent on it. """ if force_refresh is True: @@ -1315,5 +1319,7 @@ async def read_feed_ranges( [Range("", "FF", True, False)], **kwargs) - return [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).__str__() + feed_ranges = [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).to_dict() for partitionKeyRange in partition_key_ranges] + + return (feed_range for feed_range in feed_ranges) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index b1e9c6f947e6..bd03fcf87fe1 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 +from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload, Iterable from typing_extensions import Literal from azure.core import MatchConditions @@ -354,7 +354,7 @@ def query_items_change_feed( def query_items_change_feed( self, *, - feed_range: str, + feed_range: Dict[str, Any], max_item_count: Optional[int] = None, start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None, priority: Optional[Literal["High", "Low"]] = None, @@ -363,7 +363,7 @@ def query_items_change_feed( """Get a sorted list of items that were changed, in the order in which they were modified. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. :keyword int max_item_count: Max number of items to be returned in the enumeration operation. :keyword start_time: The start time to start processing chang feed items. Beginning: Processing the change feed items from the beginning of the change feed. @@ -440,7 +440,7 @@ def query_items_change_feed( """Get a sorted list of items that were changed, in the order in which they were modified. :keyword str continuation: The continuation token retrieved from previous response. - :keyword str feed_range: The feed range that is used to define the scope. + :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope. :keyword partition_key: The partition key that is used to define the scope (logical partition or a subset of a container) :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]] @@ -1365,14 +1365,18 @@ def read_feed_ranges( self, *, force_refresh: Optional[bool] = False, - **kwargs: Any) -> List[str]: + **kwargs: Any) -> Iterable[Dict[str, Any]]: """ Obtains a list of feed ranges that can be used to parallelize feed operations. :keyword bool force_refresh: Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache. :returns: A list representing the feed ranges in base64 encoded string - :rtype: List[str] + :rtype: Iterable[Dict[str, Any]] + + .. note:: + For each feed range, even through a Dict has been returned, + but in the future, the structure may change. Please just treat it as opaque and do not take any dependent on it. """ if force_refresh is True: @@ -1384,5 +1388,6 @@ def read_feed_ranges( [Range("", "FF", True, False)], # default to full range **kwargs) - return [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).__str__() + feed_ranges = [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).to_dict() for partitionKeyRange in partition_key_ranges] + return (feed_range for feed_range in feed_ranges) diff --git a/sdk/cosmos/azure-cosmos/samples/examples.py b/sdk/cosmos/azure-cosmos/samples/examples.py index 958d72c064d1..d6f56ef4a5c4 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples.py +++ b/sdk/cosmos/azure-cosmos/samples/examples.py @@ -259,12 +259,11 @@ # Get the feed ranges list from container. # [START read_feed_ranges] -container.read_feed_ranges() +feed_ranges = list(container.read_feed_ranges()) # [END read_feed_ranges] # Query a sorted list of items that were changed for one feed range # [START query_items_change_feed] -feed_ranges = container.read_feed_ranges() for item in container.query_items_change_feed(feed_range=feed_ranges[0]): print(json.dumps(item, indent=True)) # [END query_items_change_feed] diff --git a/sdk/cosmos/azure-cosmos/samples/examples_async.py b/sdk/cosmos/azure-cosmos/samples/examples_async.py index 33805fc71d7d..c40183a3228b 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples_async.py +++ b/sdk/cosmos/azure-cosmos/samples/examples_async.py @@ -265,14 +265,13 @@ async def examples_async(): # Get the feed ranges list from container. # [START read_feed_ranges] - await container.read_feed_ranges() + feed_ranges = list(await container.read_feed_ranges()) # [END read_feed_ranges] # Query a sorted list of items that were changed for one feed range. # 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_change_feed] - feed_ranges = await container.read_feed_ranges() async for item in container.query_items_change_feed(feed_range=feed_ranges[0]): print(json.dumps(item, indent=True)) # [END query_items_change_feed] diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed.py b/sdk/cosmos/azure-cosmos/test/test_change_feed.py index 01e2dc21ddb6..456f8a7dbd5a 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed.py @@ -38,7 +38,7 @@ class TestChangeFeed: def test_get_feed_ranges(self, setup): created_collection = setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), PartitionKey(path="/pk")) - result = created_collection.read_feed_ranges() + result = list(created_collection.read_feed_ranges()) assert len(result) == 1 @pytest.mark.parametrize("change_feed_filter_param", ["partitionKey", "partitionKeyRangeId", "feedRange"]) @@ -56,7 +56,7 @@ def test_query_change_feed_with_different_filter(self, change_feed_filter_param, elif change_feed_filter_param == "partitionKeyRangeId": filter_param = {"partition_key_range_id": "0"} elif change_feed_filter_param == "feedRange": - feed_ranges = created_collection.read_feed_ranges() + feed_ranges = list(created_collection.read_feed_ranges()) assert len(feed_ranges) == 1 filter_param = {"feed_range": feed_ranges[0]} else: diff --git a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py index 2ef61ee5c8a3..7af753740940 100644 --- a/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_change_feed_async.py @@ -42,7 +42,7 @@ class TestChangeFeedAsync: async def test_get_feed_ranges(self, setup): created_collection = await setup["created_db"].create_container("get_feed_ranges_" + str(uuid.uuid4()), PartitionKey(path="/pk")) - result = await created_collection.read_feed_ranges() + result = list(await created_collection.read_feed_ranges()) assert len(result) == 1 @pytest.mark.parametrize("change_feed_filter_param", ["partitionKey", "partitionKeyRangeId", "feedRange"]) @@ -57,7 +57,7 @@ async def test_query_change_feed_with_different_filter_async(self, change_feed_f elif change_feed_filter_param == "partitionKeyRangeId": filter_param = {"partition_key_range_id": "0"} elif change_feed_filter_param == "feedRange": - feed_ranges = await created_collection.read_feed_ranges() + feed_ranges = list(await created_collection.read_feed_ranges()) assert len(feed_ranges) == 1 filter_param = {"feed_range": feed_ranges[0]} else: From 3a1f1601dee8611ac43907e1a91b3e01b2bf4563 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 25 Oct 2024 13:00:28 -0700 Subject: [PATCH 50/59] Reacting to API review and adding samples. --- .../azure/cosmos/_change_feed/change_feed_state.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py index 88a05fd7d818..1f7d63f96d72 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py @@ -383,7 +383,9 @@ def from_initial_state( feed_range: Optional[FeedRangeInternal] = None if change_feed_state_context.get("feedRange"): - feed_range = FeedRangeInternalEpk.from_json(change_feed_state_context["feedRange"]) + feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8') + feed_range_json = json.loads(feed_range_str) + feed_range = FeedRangeInternalEpk.from_json(feed_range_json) elif change_feed_state_context.get("partitionKey"): if change_feed_state_context.get("partitionKeyFeedRange"): feed_range =\ From 900d001f51e91608300bd0274412dbfd6a8d233f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Fri, 25 Oct 2024 16:35:37 -0700 Subject: [PATCH 51/59] Fixed pylint --- .../azure-cosmos/azure/cosmos/aio/_container.py | 14 +++++++------- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 10 +++++----- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 93a7e05f857f..95298a3cc327 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1306,8 +1306,8 @@ async def read_feed_ranges( :rtype: Iterable[Dict[str, Any]] .. note:: - For each feed range, even through a Dict has been returned, - but in the future, the structure may change. Please just treat it as opaque and do not take any dependent on it. + For each feed range, even through a Dict has been returned, but in the future, the structure may change. + Please just treat it as opaque and do not take any dependent on it. """ if force_refresh is True: @@ -1328,11 +1328,11 @@ async def get_latest_session_token(self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any] ) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples - for a specific target feed range. The feed range can be obtained from a logical partition or by reading the - container feed ranges. This should only be used if maintaining own session token or else the sdk will - keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session - tokens and feed ranges obtained from the same container. + """ **provisional** Gets the the most up to date session token from the list of session token and feed + range tuples for a specific target feed range. The feed range can be obtained from a logical partition + or by reading the container feed ranges. This should only be used if maintaining own session token or else + the sdk willkeep track of session token. Session tokens and feed ranges are scoped to a container. + Only input session tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]] :param target_feed_range: feed range to get most up to date session token. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 3142bf027207..685f8a1972f8 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1397,11 +1397,11 @@ def get_latest_session_token( self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any]) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples - for a specific target feed range. The feed range can be obtained from a logical partition or by reading the - container feed ranges. This should only be used if maintaining own session token or else the sdk will - keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session - tokens and feed ranges obtained from the same container. + """ **provisional** Gets the the most up to date session token from the list of session token and feed + range tuples for a specific target feed range. The feed range can be obtained from a logical partition + or by reading the container feed ranges. This should only be used if maintaining own session token or else + the sdk willkeep track of session token. Session tokens and feed ranges are scoped to a container. + Only input session tokens and feed ranges obtained from the same container. :param feed_ranges_to_session_tokens: List of partition key and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]] :param target_feed_range: feed range to get most up to date session token. From eab1822517f067a8d4ec8dcdd9a9f2e7e90459e6 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 28 Oct 2024 08:21:36 -0400 Subject: [PATCH 52/59] Reacting to comments --- .../azure/cosmos/_session_token_helpers.py | 21 ++++++++++++------- .../azure/cosmos/aio/_container.py | 10 ++++----- .../azure-cosmos/azure/cosmos/container.py | 10 ++++----- .../samples/session_token_management.py | 2 +- .../samples/session_token_management_async.py | 2 +- .../azure-cosmos/test/test_feed_range.py | 6 ++++++ .../test/test_latest_session_token.py | 10 ++++----- .../test/test_latest_session_token_async.py | 10 ++++----- .../test/test_session_token_helpers.py | 17 ++++++++------- 9 files changed, 51 insertions(+), 37 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index bf833c4ced58..8e78b23cd23a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -30,17 +30,20 @@ # pylint: disable=protected-access -# ex inputs: +# ex inputs and outputs: # 1. "1:1#51", "1:1#55" -> "1:1#55" # 2. "0:1#57", "1:1#52" -> "0:1#57" +# 3. "1:1#57#3=54", "2:1#52#3=51" -> "1:1#57#3=54" +# 4. "1:1#57#3=54", "1:1#58#3=53" -> "1:1#58#3=54" def merge_session_tokens_with_same_range(session_token1: str, session_token2: str) -> str: pk_range_id1, vector_session_token1 = parse_session_token(session_token1) pk_range_id2, vector_session_token2 = parse_session_token(session_token2) pk_range_id = pk_range_id1 # The partition key range id could be different in this scenario - # Ex. get_updated_session_token([("AA", "BB"), "1:1#51"], ("AA", "DD")) -> "1:1#51" + # + # Ex. get_updated_session_token([(("AA", "BB"), "1:1#51")], ("AA", "DD")) -> "1:1#51" # Then we input this back into get_updated_session_token after a merge happened - # get_updated_session_token([("AA", "DD"), "1:1#51", ("AA", "DD"), "0:1#55"], ("AA", "DD")) -> "0:1#55" + # get_updated_session_token([(("AA", "DD"), "1:1#51"), (("AA", "DD"), "0:1#55")], ("AA", "DD")) -> "0:1#55" if pk_range_id1 != pk_range_id2: pk_range_id = pk_range_id1 \ if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2 @@ -88,12 +91,14 @@ def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[s return session_tokens # ex inputs: -# 1. [(("AA", "BB"), "1:1#51"), (("BB", "DD"), "2:1#51"), (("AA", "DD"), "0:1#55")] -> -# [("AA", "DD"), "0:1#55"] +# merge scenario +# 1. [(("AA", "BB"), "1:1#51"), (("BB", "DD"), "2:1#51"), (("AA", "DD"), "3:1#55")] -> +# [("AA", "DD"), "3:1#55"] +# split scenario # 2. [(("AA", "BB"), "1:1#57"), (("BB", "DD"), "2:1#58"), (("AA", "DD"), "0:1#55")] -> # [("AA", "DD"), "1:1#57,2:1#58"] -# 3. [(("AA", "BB"), "1:1#57"), (("BB", "DD"), "2:1#52"), (("AA", "DD"), "0:1#55")] -> -# [("AA", "DD"), "1:1#57,2:1#52,0:1#55"] +# 3. [(("AA", "BB"), "4:1#57"), (("BB", "DD"), "1:1#52"), (("AA", "DD"), "3:1#55")] -> +# [("AA", "DD"), "4:1#57,1:1#52,3:1#55"] # goal here is to detect any obvious merges or splits that happened # compound session tokens are not considered will just pass them along def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> List[Tuple[Range, str]]: @@ -188,7 +193,7 @@ def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[Dict[str, session_token = overlapping_ranges[i][1] session_token_1 = overlapping_ranges[j][1] if (not is_compound_session_token(session_token) and - not is_compound_session_token(overlapping_ranges[j][1]) and + not is_compound_session_token(session_token_1) and cur_feed_range == overlapping_ranges[j][0]): session_token = merge_session_tokens_with_same_range(session_token, session_token_1) feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 95298a3cc327..f47d89bb521a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1329,11 +1329,11 @@ async def get_latest_session_token(self, target_feed_range: Dict[str, Any] ) -> str: """ **provisional** Gets the the most up to date session token from the list of session token and feed - range tuples for a specific target feed range. The feed range can be obtained from a logical partition + range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else - the sdk willkeep track of session token. Session tokens and feed ranges are scoped to a container. - Only input session tokens and feed ranges obtained from the same container. - :param feed_ranges_to_session_tokens: List of partition key and session token tuples. + the CosmosClient instance will keep track of session token. Session tokens and feed ranges are + scoped to a container. Only input session tokens and feed ranges obtained from the same container. + :param feed_ranges_to_session_tokens: List of feed range and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]] :param target_feed_range: feed range to get most up to date session token. :type target_feed_range: Dict[str, Any] @@ -1345,7 +1345,7 @@ async def get_latest_session_token(self, async 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: PartitionKey + :type partition_key: PartitionKeyType :returns: a feed range :rtype: Dict[str, Any] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 10cb676b9137..0a6a6e5656b4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1398,11 +1398,11 @@ def get_latest_session_token( feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any]) -> str: """ **provisional** Gets the the most up to date session token from the list of session token and feed - range tuples for a specific target feed range. The feed range can be obtained from a logical partition + range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else - the sdk willkeep track of session token. Session tokens and feed ranges are scoped to a container. - Only input session tokens and feed ranges obtained from the same container. - :param feed_ranges_to_session_tokens: List of partition key and session token tuples. + the CosmosClient instance will keep track of session token. Session tokens and feed ranges are + scoped to a container. Only input session tokens and feed ranges obtained from the same container. + :param feed_ranges_to_session_tokens: List of feed range and session token tuples. :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]] :param target_feed_range: feed range to get most up to date session token. :type target_feed_range: Dict[str, Any] @@ -1414,7 +1414,7 @@ def get_latest_session_token( 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: PartitionKey + :type partition_key: PartitionKeyType :returns: a feed range :rtype: Dict[str, Any] .. note:: diff --git a/sdk/cosmos/azure-cosmos/samples/session_token_management.py b/sdk/cosmos/azure-cosmos/samples/session_token_management.py index 68a1f4b563af..21e303c884f3 100644 --- a/sdk/cosmos/azure-cosmos/samples/session_token_management.py +++ b/sdk/cosmos/azure-cosmos/samples/session_token_management.py @@ -51,7 +51,7 @@ def storing_session_tokens_pk(container): # Everything below is just a simulation of what could be run on different machines and clients # to store session tokens in a cache by feed range from the partition key. - # The cache is a list of tuples here for simplicity but in a real-world scenario, it would be some service. + # The cache is a Dict here for simplicity but in a real-world scenario, it would be some service. feed_ranges_and_session_tokens = [] previous_session_token = "" diff --git a/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py b/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py index ed294944bcc5..f4c953cdeb3d 100644 --- a/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py @@ -52,7 +52,7 @@ async def storing_session_tokens_pk(container): # Everything below is just a simulation of what could be run on different machines and clients # to store session tokens in a cache by feed range from the partition key. - # The cache is a list of tuples here for simplicity but in a real-world scenario, it would be some service. + # The cache is a Dict here for simplicity but in a real-world scenario, it would be some service. feed_ranges_and_session_tokens = [] previous_session_token = "" diff --git a/sdk/cosmos/azure-cosmos/test/test_feed_range.py b/sdk/cosmos/azure-cosmos/test/test_feed_range.py index 1d18290c21cb..f44c39e034a9 100644 --- a/sdk/cosmos/azure-cosmos/test/test_feed_range.py +++ b/sdk/cosmos/azure-cosmos/test/test_feed_range.py @@ -47,6 +47,12 @@ def setup(): False), (Range("3F", "7F", True, False), Range("", "2F", True, False), + False), + (Range("3F", "3F", True, True), + Range("3F", "3F", True, True), + True), + (Range("3F", "3F", True, True), + Range("4F", "4F", True, True), False) ] diff --git a/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py index 0f8d778309b0..7095c22b2279 100644 --- a/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token.py @@ -102,7 +102,7 @@ def test_latest_session_token_from_physical_pk(self): pk_range_id2, session_token2 = parse_session_token(session_tokens[1]) pk_range_ids = [pk_range_id1, pk_range_id2] - assert 320 >= (session_token1.global_lsn + session_token2.global_lsn) + assert 320 == (session_token1.global_lsn + session_token2.global_lsn) assert '1' in pk_range_ids assert '2' in pk_range_ids self.database.delete_container(container.id) @@ -170,8 +170,8 @@ def create_items_logical_pk(container, target_pk_range, previous_session_token, target_session_token = "" for i in range(100): item = create_item(hpk) - container.create_item(item, session_token=previous_session_token) - session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + response = container.create_item(item, session_token=previous_session_token) + session_token = response.get_response_headers()[HttpHeaders.SessionToken] pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] pk_range = container.feed_range_from_partition_key(pk) pk_feed_range_epk = FeedRangeInternalEpk.from_json(pk_range) @@ -196,8 +196,8 @@ def create_items_physical_pk(container, pk_feed_range, previous_session_token, f for i in range(100): item = create_item(hpk) - container.create_item(item, session_token=previous_session_token) - session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + response = container.create_item(item, session_token=previous_session_token) + session_token = response.get_response_headers()[HttpHeaders.SessionToken] if hpk: pk = [item['state'], item['city'], item['zipcode']] curr_feed_range = container.feed_range_from_partition_key(pk) diff --git a/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py index 66a4693529d9..b99751af82e4 100644 --- a/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py +++ b/sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py @@ -105,7 +105,7 @@ async def test_latest_session_token_from_physical_pk(self): pk_range_id2, session_token2 = parse_session_token(session_tokens[1]) pk_range_ids = [pk_range_id1, pk_range_id2] - assert 320 <= (session_token1.global_lsn + session_token2.global_lsn) + assert 320 == (session_token1.global_lsn + session_token2.global_lsn) assert '1' in pk_range_ids assert '2' in pk_range_ids await self.database.delete_container(container.id) @@ -173,8 +173,8 @@ async def create_items_logical_pk(container, target_pk_range, previous_session_t target_session_token = "" for i in range(100): item = create_item(hpk) - await container.create_item(item, session_token=previous_session_token) - session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + response = await container.create_item(item, session_token=previous_session_token) + session_token = response.get_response_headers()[HttpHeaders.SessionToken] pk = item['pk'] if not hpk else [item['state'], item['city'], item['zipcode']] pk_feed_range = await container.feed_range_from_partition_key(pk) pk_feed_range_epk = FeedRangeInternalEpk.from_json(pk_feed_range) @@ -199,8 +199,8 @@ async def create_items_physical_pk(container, pk_feed_range, previous_session_to for i in range(100): item = create_item(hpk) - await container.create_item(item, session_token=previous_session_token) - session_token = container.client_connection.last_response_headers[HttpHeaders.SessionToken] + response = await container.create_item(item, session_token=previous_session_token) + session_token = response.get_response_headers()[HttpHeaders.SessionToken] if hpk: pk = [item['state'], item['city'], item['zipcode']] curr_feed_range = await container.feed_range_from_partition_key(pk) diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 46c735fc457e..1597a00d7902 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -40,6 +40,7 @@ def create_split_ranges(): ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), # several ranges being equal to one range + # Highest GLSN, which is the 55 in this ex. "1:1#55#3=52", is the one that will be used ([(("AA", "DD"), "0:1#42#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "1:1#55#3=52"), @@ -48,15 +49,17 @@ def create_split_ranges(): (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=52"), # several ranges being equal to one range + # AA-DD can be created from the other ranges but the GLSN's are not all larger than the one + # in the AA-DD range so we just compound ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#66#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=52,1:1#66#3=52"), # merge with one child - ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], - ("AA", "DD"), "0:1#55#3=52"), + ([(("AA", "DD"), "3:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52")], + ("AA", "DD"), "3:1#55#3=52"), # merge with two children - ([(("AA", "DD"), "0:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52"), (("BB", "DD"),"2:1#54#3=52")], - ("AA", "DD"), "0:1#55#3=52"), + ([(("AA", "DD"), "3:1#55#3=52"), (("AA", "BB"),"1:1#51#3=52"), (("BB", "DD"),"2:1#54#3=52")], + ("AA", "DD"), "3:1#55#3=52"), # compound session token ([(("AA", "DD"), "2:1#54#3=52,1:1#55#3=52"), (("AA", "BB"),"0:1#51#3=52")], ("AA", "BB"), "2:1#54#3=52,1:1#55#3=52,0:1#51#3=52"), @@ -70,9 +73,9 @@ def create_split_ranges(): ([(("AA", "BB"), "0:1#54#3=52"), (("AA", "BB"),"0:2#57#3=53")], ("AA", "BB"), "0:2#57#3=53"), # mixed scenarios - ([(("AA", "DD"), "0:1#60#3=53"), (("AA", "BB"), "1:1#54#3=52"), (("AA", "BB"), "1:1#52#3=53"), - (("BB", "CC"),"1:1#53#3=52"), (("BB", "CC"),"1:1#70#3=55,3:1#90#3=52"), - (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=53,1:1#70#3=55,3:1#90#3=52") + ([(("AA", "DD"), "3:1#60#3=53"), (("AA", "BB"), "1:1#54#3=52"), (("AA", "BB"), "1:1#52#3=53"), + (("BB", "CC"),"1:1#53#3=52"), (("BB", "CC"),"6:1#70#3=55,4:1#90#3=52"), + (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "3:1#60#3=53,6:1#70#3=55,4:1#90#3=52") ] @pytest.mark.cosmosEmulator From 97ffec7865fef3ab67bd1a035009b9dacb2d8439 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 28 Oct 2024 09:33:57 -0400 Subject: [PATCH 53/59] Reacting to comments --- .../azure/cosmos/_session_token_helpers.py | 35 +++++++++---------- .../test/test_session_token_helpers.py | 18 ++++++---- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 8e78b23cd23a..8c6b1536332f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -69,26 +69,25 @@ def split_compound_session_tokens(compound_session_tokens: List[Tuple[Range, str return session_tokens # ex inputs: -# ["1:1#51", "1:1#55", "1:1#57"] -> ["1:1#57"] +# ["1:1#51", "1:1#55", "1:1#57", "2:1#42", "2:1#45", "2:1#47"] -> ["1:1#57"] def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[str]: - i = 0 - while i < len(session_tokens): - j = i + 1 - while j < len(session_tokens): - pk_range_id1, vector_session_token1 = parse_session_token(session_tokens[i]) - pk_range_id2, vector_session_token2 = parse_session_token(session_tokens[j]) - if pk_range_id1 == pk_range_id2: - vector_session_token = vector_session_token1.merge(vector_session_token2) - session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token) - remove_session_tokens = [session_tokens[i], session_tokens[j]] - for token in remove_session_tokens: - session_tokens.remove(token) - i = -1 - break - j += 1 - i += 1 + pk_session_tokens: Dict[str, List[str]] = {} + for session_token in session_tokens: + pk_range_id, _ = parse_session_token(session_token) + if pk_range_id in pk_session_tokens: + pk_session_tokens[pk_range_id].append(session_token) + else: + pk_session_tokens[pk_range_id] = [session_token] - return session_tokens + processed_session_tokens = [] + for session_tokens_same_pk in pk_session_tokens.values(): + pk_range_id, vector_session_token = parse_session_token(session_tokens_same_pk[0]) + for session_token in session_tokens_same_pk[1:]: + _, vector_session_token_1 = parse_session_token(session_token) + vector_session_token = vector_session_token.merge(vector_session_token_1) + processed_session_tokens.append(pk_range_id + ":" + vector_session_token.session_token) + + return processed_session_tokens # ex inputs: # merge scenario diff --git a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py index 1597a00d7902..d733c213e788 100644 --- a/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/test/test_session_token_helpers.py @@ -39,18 +39,24 @@ def create_split_ranges(): # split with one child ([(("AA", "DD"), "0:1#51#3=52"), (("AA", "BB"),"1:1#55#3=52")], ("AA", "DD"), "0:1#51#3=52,1:1#55#3=52"), - # several ranges being equal to one range - # Highest GLSN, which is the 55 in this ex. "1:1#55#3=52", is the one that will be used + # Highest GLSN, which is 55 in this # cspell:disable-line + # ex. "1:1#55#3=52", is the one that will be returned because + # it is higher than all of the other feed range contained in the same + # range ([(("AA", "DD"), "0:1#42#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "1:1#55#3=52"), - # several ranges being equal to one range + # Highest GLSN, which is 60 in this # cspell:disable-line + # ex. "1:1#60#3=52", is the one that will be returned because + # it is higher than all of the other feed range contained in the same + # range with some of them overlapping with each other ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#53#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=52"), - # several ranges being equal to one range - # AA-DD can be created from the other ranges but the GLSN's are not all larger than the one - # in the AA-DD range so we just compound + # AA-DD can be created from the other ranges + # but the GLSN's are not all larger than the one # cspell:disable-line + # in the AA-DD range so we just compound as cannot make + # conclusions in this case ([(("AA", "DD"), "0:1#60#3=52"), (("AA", "BB"), "1:1#51#3=52"), (("BB", "CC"),"1:1#66#3=52"), (("CC", "DD"),"1:1#55#3=52")], ("AA", "DD"), "0:1#60#3=52,1:1#66#3=52"), From 2264465982cfbfd7b264b01fbaae990d16251fd5 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 29 Oct 2024 08:26:23 -0400 Subject: [PATCH 54/59] Reacting to comments --- sdk/cosmos/azure-cosmos/samples/examples.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/samples/examples.py b/sdk/cosmos/azure-cosmos/samples/examples.py index 7f73ec5b8ad3..5b12a6870574 100644 --- a/sdk/cosmos/azure-cosmos/samples/examples.py +++ b/sdk/cosmos/azure-cosmos/samples/examples.py @@ -271,10 +271,11 @@ # Figure out if a feed range is a subset of another feed range. # This example sees in which feed range from the container a feed range from a partition key is part of. # [START is_feed_range_subset] -parent_feed_range: Dict[str, Any] = next( - (feed_range for feed_range in feed_ranges if container.is_feed_range_subset(feed_range, feed_range_from_pk)), - {} -) +parent_feed_range = {} +for feed_range in feed_ranges: + if container.is_feed_range_subset(feed_range, feed_range_from_pk): + parent_feed_range = feed_range + break # [END is_feed_range_subset] # Query a sorted list of items that were changed for one feed range From 35588fa937f296ef88d9b8a7c4042094e3d63075 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 29 Oct 2024 08:49:19 -0400 Subject: [PATCH 55/59] Reacting to comments --- sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py index 8c6b1536332f..36a31328e20d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py @@ -69,7 +69,7 @@ def split_compound_session_tokens(compound_session_tokens: List[Tuple[Range, str return session_tokens # ex inputs: -# ["1:1#51", "1:1#55", "1:1#57", "2:1#42", "2:1#45", "2:1#47"] -> ["1:1#57"] +# ["1:1#51", "1:1#55", "1:1#57", "2:1#42", "2:1#45", "2:1#47"] -> ["1:1#57", "2:1#47"] def merge_session_tokens_for_same_partition(session_tokens: List[str]) -> List[str]: pk_session_tokens: Dict[str, List[str]] = {} for session_token in session_tokens: From c42966f4b2c1b00cfca7dd2201510f72b27e4ab6 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 30 Oct 2024 16:37:56 -0400 Subject: [PATCH 56/59] Fix pydoc --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.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/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index f47d89bb521a..ff36cb22d8f4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1328,7 +1328,7 @@ async def get_latest_session_token(self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any] ) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed + """ Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 0a6a6e5656b4..df51448b3a13 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1397,7 +1397,7 @@ def get_latest_session_token( self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any]) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed + """ Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are From 786e357a899bd7de99cbac1f7330df598c25d7cf Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 31 Oct 2024 09:38:37 -0400 Subject: [PATCH 57/59] Fix pydoc --- sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py | 4 ++-- sdk/cosmos/azure-cosmos/azure/cosmos/container.py | 3 ++- 2 files changed, 4 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 ff36cb22d8f4..62ae95e6c63d 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1328,7 +1328,7 @@ async def get_latest_session_token(self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any] ) -> str: - """ Gets the the most up to date session token from the list of session token and feed + """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are @@ -1343,7 +1343,7 @@ async def get_latest_session_token(self, return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range) async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict[str, Any]: - """Gets the feed range for a given partition key. + """ Gets the feed range for a given partition key. :param partition_key: partition key to get feed range. :type partition_key: PartitionKeyType :returns: a feed range diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index df51448b3a13..ab221f4693ec 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1397,7 +1397,7 @@ def get_latest_session_token( self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], target_feed_range: Dict[str, Any]) -> str: - """ Gets the the most up to date session token from the list of session token and feed + """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are @@ -1417,6 +1417,7 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict :type partition_key: PartitionKeyType :returns: a feed range :rtype: Dict[str, Any] + .. note:: For the feed range, even through a Dict has been returned, but in the future, the structure may change. Please just treat it as opaque and do not take any dependence on it. From 0de21b49e29d812384d0a62c1f6bdf96a8e10eaa Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 31 Oct 2024 15:26:30 -0400 Subject: [PATCH 58/59] reacting to comments --- .../azure/cosmos/aio/_container.py | 31 ++++++++++--------- .../azure-cosmos/azure/cosmos/container.py | 7 +++-- .../samples/session_token_management.py | 5 +-- .../samples/session_token_management_async.py | 5 +-- 4 files changed, 28 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py index 62ae95e6c63d..b37e3cf17c6f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py @@ -1305,9 +1305,9 @@ async def read_feed_ranges( :returns: A list representing the feed ranges in base64 encoded string :rtype: Iterable[Dict[str, Any]] - .. note:: - For each feed range, even through a Dict has been returned, but in the future, the structure may change. - Please just treat it as opaque and do not take any dependent on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. """ if force_refresh is True: @@ -1324,11 +1324,14 @@ async def read_feed_ranges( for partitionKeyRange in partition_key_ranges] return (feed_range for feed_range in feed_ranges) - async def get_latest_session_token(self, - feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], - target_feed_range: Dict[str, Any] - ) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed + async def get_latest_session_token( + self, + feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], + target_feed_range: Dict[str, Any] + ) -> str: + """ **provisional** This method is still in preview and may be subject to breaking changes. + + Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are @@ -1349,9 +1352,9 @@ async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) - :returns: a feed range :rtype: Dict[str, Any] - .. note:: - For the feed range, even through a Dict has been returned, but in the future, - the structure may change. Please just treat it as opaque and do not take any dependence on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. """ return FeedRangeInternalEpk(await self._get_epk_range_for_partition_key(partition_key)).to_dict() @@ -1366,9 +1369,9 @@ async def is_feed_range_subset(self, parent_feed_range: Dict[str, Any], :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool - .. note:: - For the feed range, even through a Dict has been returned, but in the future, - the structure may change. Please just treat it as opaque and do not take any dependence on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. """ parent_feed_range_epk = FeedRangeInternalEpk.from_json(parent_feed_range) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index ab221f4693ec..5b95c626eb55 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1396,8 +1396,11 @@ def read_feed_ranges( def get_latest_session_token( self, feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]], - target_feed_range: Dict[str, Any]) -> str: - """ **provisional** Gets the the most up to date session token from the list of session token and feed + target_feed_range: Dict[str, Any] + ) -> str: + """ **provisional** This method is still in preview and may be subject to breaking changes. + + Gets the the most up to date session token from the list of session token and feed range tuples for a specific target feed range. The feed range can be obtained from a partition key or by reading the container feed ranges. This should only be used if maintaining own session token or else the CosmosClient instance will keep track of session token. Session tokens and feed ranges are diff --git a/sdk/cosmos/azure-cosmos/samples/session_token_management.py b/sdk/cosmos/azure-cosmos/samples/session_token_management.py index 21e303c884f3..5915f5fa50b7 100644 --- a/sdk/cosmos/azure-cosmos/samples/session_token_management.py +++ b/sdk/cosmos/azure-cosmos/samples/session_token_management.py @@ -12,6 +12,7 @@ import azure.cosmos.exceptions as exceptions import config +from azure.identity import DefaultAzureCredential from azure.cosmos.http_constants import HttpHeaders # ---------------------------------------------------------------------------------------------------------- @@ -39,7 +40,7 @@ # ---------------------------------------------------------------------------------------------------------- HOST = config.settings['host'] -MASTER_KEY = config.settings['master_key'] +CREDENTIAL = DefaultAzureCredential() DATABASE_ID = config.settings['database_id'] CONTAINER_ID = config.settings['container_id'] @@ -117,7 +118,7 @@ def storing_session_tokens_container_feed_ranges(container): def run_sample(): - with CosmosClient(HOST, {'masterKey': MASTER_KEY}) as client: + with CosmosClient(HOST, CREDENTIAL) as client: try: db = client.create_database_if_not_exists(id=DATABASE_ID) container = db.create_container_if_not_exists(id=CONTAINER_ID, partition_key=PartitionKey('/pk')) diff --git a/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py b/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py index f4c953cdeb3d..ff9a17946d69 100644 --- a/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py +++ b/sdk/cosmos/azure-cosmos/samples/session_token_management_async.py @@ -13,6 +13,7 @@ import asyncio import config +from azure.identity import DefaultAzureCredential from azure.cosmos.http_constants import HttpHeaders # ---------------------------------------------------------------------------------------------------------- @@ -40,7 +41,7 @@ # ---------------------------------------------------------------------------------------------------------- HOST = config.settings['host'] -MASTER_KEY = config.settings['master_key'] +CREDENTIAL = DefaultAzureCredential() DATABASE_ID = config.settings['database_id'] CONTAINER_ID = config.settings['container_id'] @@ -119,7 +120,7 @@ async def storing_session_tokens_container_feed_ranges(container): async def run_sample(): - async with CosmosClient(HOST, {'masterKey': MASTER_KEY}) as client: + async with CosmosClient(HOST, CREDENTIAL) as client: try: db = await client.create_database_if_not_exists(id=DATABASE_ID) container = await db.create_container_if_not_exists(id=CONTAINER_ID, partition_key=PartitionKey('/pk')) From d32a6f137806cb3e229d49e9b4de34a6452d6c6f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 31 Oct 2024 15:30:03 -0400 Subject: [PATCH 59/59] reacting to comments --- .../azure-cosmos/azure/cosmos/container.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index 5b95c626eb55..f0c27b449e63 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1375,9 +1375,9 @@ def read_feed_ranges( :returns: A list representing the feed ranges in base64 encoded string :rtype: Iterable[Dict[str, Any]] - .. note:: - For each feed range, even through a Dict has been returned, but in the future, the structure may change. - Please just treat it as opaque and do not take any dependent on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. """ if force_refresh is True: @@ -1421,9 +1421,9 @@ def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict :returns: a feed range :rtype: Dict[str, Any] - .. note:: - For the feed range, even through a Dict has been returned, but in the future, - the structure may change. Please just treat it as opaque and do not take any dependence on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. """ return FeedRangeInternalEpk(self._get_epk_range_for_partition_key(partition_key)).to_dict() @@ -1437,9 +1437,10 @@ def is_feed_range_subset(self, parent_feed_range: Dict[str, Any], child_feed_ran :returns: a boolean indicating if child feed range is a subset of parent feed range :rtype: bool - .. note:: - For the feed range, even through a Dict has been returned, but in the future, - the structure may change. Please just treat it as opaque and do not take any dependence on it. + .. warning:: + The structure of the dict representation of a feed range may vary, including which keys + are present. It therefore should only be treated as an opaque value. + """ parent_feed_range_epk = FeedRangeInternalEpk.from_json(parent_feed_range) child_feed_range_epk = FeedRangeInternalEpk.from_json(child_feed_range)