refactor

Author: annie-mac

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,16 +52,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")
         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
 

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