ExposeFeedRangeAsClassType (#37444)

* expose feedRange as a class type

---------

Co-authored-by: annie-mac <xinlian@microsoft.com>

sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py

@@ -42,6 +42,7 @@
 )
 from .partition_key import PartitionKey
 from .permission import Permission
+from ._feed_range import FeedRange
 
 __all__ = (
     "CosmosClient",
@@ -64,5 +65,6 @@
     "TriggerType",
     "ConnectionRetryPolicy",
     "ThroughputProperties",
+    "FeedRange"
 )
 __version__ = VERSION

sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py

@@ -34,7 +34,8 @@
 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_internal import (FeedRangeInternal, FeedRangeInternalEpk,
+                                                           FeedRangeInternalPartitionKey)
 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
@@ -79,7 +80,7 @@ def apply_server_response_continuation(self, continuation: str, has_modified_res
     def from_json(
             container_link: str,
             container_rid: str,
-            change_feed_state_context: Dict[str, Any]):
+            change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedState':
 
         if (change_feed_state_context.get("partitionKeyRangeId")
                 or change_feed_state_context.get("continuationPkRangeId")):
@@ -184,7 +185,7 @@ def __init__(
             self,
             container_link: str,
             container_rid: str,
-            feed_range: FeedRange,
+            feed_range: FeedRangeInternal,
             change_feed_start_from: ChangeFeedStartFromInternal,
             continuation: Optional[FeedRangeCompositeContinuation]
     ) -> None:
@@ -380,22 +381,20 @@ def from_initial_state(
             collection_rid: str,
             change_feed_state_context: Dict[str, Any]) -> 'ChangeFeedStateV2':
 
-        feed_range: Optional[FeedRange] = None
+        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 = FeedRangeEpk(Range.ParseFromDict(feed_range_json))
+            feed_range = change_feed_state_context.get("feedRange")
         elif change_feed_state_context.get("partitionKey"):
             if change_feed_state_context.get("partitionKeyFeedRange"):
                 feed_range =\
-                    FeedRangePartitionKey(
+                    FeedRangeInternalPartitionKey(
                         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 = FeedRangeEpk(
+            feed_range = FeedRangeInternalEpk(
                 Range(
                 "",
                 "FF",
@@ -405,9 +404,12 @@ def from_initial_state(
 
         change_feed_start_from = (
             ChangeFeedStartFromInternal.from_start_time(change_feed_state_context.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)
+
+        if feed_range is not None:
+            return cls(
+                container_link=container_link,
+                container_rid=collection_rid,
+                feed_range=feed_range,
+                change_feed_start_from=change_feed_start_from,
+                continuation=None)
+        raise RuntimeError("feed_range is empty")

sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_composite_continuation_token.py

@@ -26,7 +26,8 @@
 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
+from azure.cosmos._change_feed.feed_range_internal import (FeedRangeInternal, FeedRangeInternalEpk,
+                                                           FeedRangeInternalPartitionKey)
 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
@@ -39,7 +40,7 @@ class FeedRangeCompositeContinuation:
     def __init__(
             self,
             container_rid: str,
-            feed_range: FeedRange,
+            feed_range: FeedRangeInternal,
             continuation: Deque[CompositeContinuationToken]) -> None:
         if container_rid is None:
             raise ValueError("container_rid is missing")
@@ -87,11 +88,11 @@ def from_json(cls, data) -> 'FeedRangeCompositeContinuation':
                         for child_range_continuation_token in continuation_data]
 
         # parsing feed range
-        feed_range: Optional[FeedRange] = None
-        if data.get(FeedRangeEpk.type_property_name):
-            feed_range = FeedRangeEpk.from_json(data)
-        elif data.get(FeedRangePartitionKey.type_property_name):
-            feed_range = FeedRangePartitionKey.from_json(data, continuation[0].feed_range)
+        feed_range: Optional[FeedRangeInternal] = None
+        if data.get(FeedRangeInternalEpk.type_property_name):
+            feed_range = FeedRangeInternalEpk.from_json(data)
+        elif data.get(FeedRangeInternalPartitionKey.type_property_name):
+            feed_range = FeedRangeInternalPartitionKey.from_json(data, continuation[0].feed_range)
         else:
             raise ValueError("Invalid feed range composite continuation token [Missing feed range scope]")
 
@@ -171,5 +172,5 @@ def apply_not_modified_response(self) -> None:
             self._initial_no_result_range = self._current_token.feed_range
 
     @property
-    def feed_range(self) -> FeedRange:
+    def feed_range(self) -> FeedRangeInternal:
         return self._feed_range

sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range.py → sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/feed_range_internal.py

@@ -22,14 +22,16 @@
 """Internal class for feed range implementation in the Azure Cosmos
 database service.
 """
+import base64
+import json
 from abc import ABC, abstractmethod
 from typing import Union, List, Dict, Any
 
 from azure.cosmos._routing.routing_range import Range
 from azure.cosmos.partition_key import _Undefined, _Empty
 
 
-class FeedRange(ABC):
+class FeedRangeInternal(ABC):
 
     @abstractmethod
     def get_normalized_range(self) -> Range:
@@ -39,7 +41,15 @@ def get_normalized_range(self) -> Range:
     def to_dict(self) -> Dict[str, Any]:
         pass
 
-class FeedRangePartitionKey(FeedRange):
+    def _to_base64_encoded_string(self) -> str:
+        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')
+
+class FeedRangeInternalPartitionKey(FeedRangeInternal):
     type_property_name = "PK"
 
     def __init__(
@@ -69,7 +79,7 @@ def to_dict(self) -> Dict[str, Any]:
         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) -> 'FeedRangeInternalPartitionKey':
         if data.get(cls.type_property_name):
             pk_value = data.get(cls.type_property_name)
             if not pk_value:
@@ -80,18 +90,19 @@ def from_json(cls, data: Dict[str, Any], feed_range: Range) -> 'FeedRangePartiti
                 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,"
+        raise ValueError(f"Can not parse FeedRangeInternalPartitionKey from the json,"
                          f" there is no property {cls.type_property_name}")
 
 
-class FeedRangeEpk(FeedRange):
+class FeedRangeInternalEpk(FeedRangeInternal):
     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
+        self._base64_encoded_string = None
 
     def get_normalized_range(self) -> Range:
         return self._range.to_normalized_range()
@@ -102,8 +113,20 @@ def to_dict(self) -> Dict[str, Any]:
         }
 
     @classmethod
-    def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeEpk':
+    def from_json(cls, data: Dict[str, Any]) -> 'FeedRangeInternalEpk':
         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}")
+        raise ValueError(f"Can not parse FeedRangeInternalEPK from the json,"
+                         f" there is no property {cls.type_property_name}")
+
+    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.
+        """
+        if self._base64_encoded_string is None:
+            self._base64_encoded_string = self._to_base64_encoded_string()
+
+        return self._base64_encoded_string

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

@@ -1162,7 +1162,6 @@ def _QueryChangeFeed(
             options = {}
         else:
             options = dict(options)
-        options["changeFeed"] = True
 
         resource_key_map = {"Documents": "docs"}
 

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

@@ -44,7 +44,6 @@ def __init__(self, client, options):
         """
         self._client = client
         self._options = options
-        self._is_change_feed = "changeFeed" in options and options["changeFeed"] is True
         self._continuation = self._get_initial_continuation()
         self._has_started = False
         self._has_finished = False
@@ -117,10 +116,6 @@ async def _fetch_items_helper_no_retries(self, fetch_function):
         fetched_items = []
         new_options = copy.deepcopy(self._options)
         while self._continuation or not self._has_started:
-            # Check if this is first fetch for read from specific time change feed.
-            # For read specific time the first fetch will return empty even if we have more pages.
-            is_s_time_first_fetch = self._is_change_feed and self._options.get("startTime") and not self._has_started
-
             new_options["continuation"] = self._continuation
 
             response_headers = {}
@@ -129,16 +124,8 @@ async def _fetch_items_helper_no_retries(self, fetch_function):
                 self._has_started = True
 
             continuation_key = http_constants.HttpHeaders.Continuation
-            # Use Etag as continuation token for change feed queries.
-            if self._is_change_feed:
-                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.
-            # No initial fetch for start time change feed, so we need to pass continuation token for first fetch
-            if not self._is_change_feed or fetched_items or is_s_time_first_fetch:
-                self._continuation = response_headers.get(continuation_key)
-            else:
-                self._continuation = None
+            self._continuation = response_headers.get(continuation_key)
+
             if fetched_items:
                 break
         return fetched_items

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

@@ -42,7 +42,6 @@ def __init__(self, client, options):
         """
         self._client = client
         self._options = options
-        self._is_change_feed = "changeFeed" in options and options["changeFeed"] is True
         self._continuation = self._get_initial_continuation()
         self._has_started = False
         self._has_finished = False
@@ -115,9 +114,6 @@ def _fetch_items_helper_no_retries(self, fetch_function):
         fetched_items = []
         new_options = copy.deepcopy(self._options)
         while self._continuation or not self._has_started:
-            # Check if this is first fetch for read from specific time change feed.
-            # For read specific time the first fetch will return empty even if we have more pages.
-            is_s_time_first_fetch = self._is_change_feed and self._options.get("startTime") and not self._has_started
             if not self._has_started:
                 self._has_started = True
             new_options["continuation"] = self._continuation
@@ -126,16 +122,8 @@ def _fetch_items_helper_no_retries(self, fetch_function):
             (fetched_items, response_headers) = fetch_function(new_options)
 
             continuation_key = http_constants.HttpHeaders.Continuation
-            # Use Etag as continuation token for change feed queries.
-            if self._is_change_feed:
-                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
-            if not self._is_change_feed or fetched_items or is_s_time_first_fetch:
-                self._continuation = response_headers.get(continuation_key)
-            else:
-                self._continuation = None
+            self._continuation = response_headers.get(continuation_key)
+
             if fetched_items:
                 break
         return fetched_items

sdk/cosmos/azure-cosmos/azure/cosmos/_feed_range.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.
+
+import base64
+import json
+from abc import ABC
+from typing import Any, Dict
+
+from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk
+from azure.cosmos._routing.routing_range import Range
+
+# pylint: disable=protected-access
+class FeedRange(ABC):
+    """Represents a single feed range in an Azure Cosmos DB SQL API container.
+
+    """
+    @staticmethod
+    def from_string(json_str: str) -> 'FeedRange':
+        """
+        Create a feed range from previously obtained string representation.
+
+        :param str json_str: A string representation of a feed range.
+        :return: A feed range.
+        :rtype: ~azure.cosmos.FeedRange
+        """
+        feed_range_json_str = base64.b64decode(json_str).decode('utf-8')
+        feed_range_json = json.loads(feed_range_json_str)
+        if feed_range_json.get(FeedRangeEpk.type_property_name):
+            return FeedRangeEpk._from_json(feed_range_json)
+
+        raise ValueError("Invalid feed range base64 encoded string [Wrong feed range type]")
+
+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._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.
+        """
+        return self._feed_range_internal.__str__()
+
+    @classmethod
+    def _from_json(cls, data: Dict[str, Any]) -> 'FeedRange':
+        return cls(FeedRangeInternalEpk.from_json(data)._range)

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

@@ -25,12 +25,8 @@
 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"""