refactor

Author: annie-mac

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")
-
-

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,15 +269,17 @@ 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 = 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)

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

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

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")
+
+
+