Fix issues with error handling, refactor tests

Author: tolik0

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -32,7 +32,7 @@ def __init__(self, create_function: Callable[..., ConcurrentCursor]):
         self._create_function = create_function
 
     def create(
-        self, stream_state: Mapping[str, Any], runtime_lookback_window: Any
+        self, stream_state: Mapping[str, Any], runtime_lookback_window: Optional[timedelta]
     ) -> ConcurrentCursor:
         return self._create_function(
             stream_state=stream_state, runtime_lookback_window=runtime_lookback_window
@@ -187,7 +187,7 @@ def _generate_slices_from_partition(self, partition: StreamSlice) -> Iterable[St
         if not cursor:
             cursor = self._create_cursor(
                 self._global_cursor,
-                self._lookback_window if self._global_cursor else self._NO_CURSOR_STATE,
+                self._lookback_window if self._global_cursor else 0,
             )
             self._cursor_per_partition[self._to_partition_key(partition.partition)] = cursor
             self._semaphore_per_partition[self._to_partition_key(partition.partition)] = (
@@ -218,9 +218,6 @@ def _ensure_partition_limit(self) -> None:
                 f"The maximum number of partitions has been reached. Dropping the oldest partition: {oldest_partition}. Over limit: {self._over_limit}."
             )
 
-    def limit_reached(self) -> bool:
-        return self._over_limit > self.DEFAULT_MAX_PARTITIONS_NUMBER
-
     def _set_initial_state(self, stream_state: StreamState) -> None:
         """
         Initialize the cursor's state using the provided `stream_state`.
@@ -286,6 +283,10 @@ def _set_initial_state(self, stream_state: StreamState) -> None:
                 self._global_cursor = deepcopy(stream_state[self._GLOBAL_STATE_KEY])
                 self._new_global_cursor = deepcopy(stream_state[self._GLOBAL_STATE_KEY])
 
+        # Set initial parent state
+        if stream_state.get("parent_state"):
+            self._parent_state = stream_state["parent_state"]
+
         # Set parent state for partition routers based on parent streams
         self._partition_router.set_initial_state(stream_state)
 
@@ -305,12 +306,11 @@ def _to_dict(self, partition_key: str) -> Mapping[str, Any]:
         return self._partition_serializer.to_partition(partition_key)
 
     def _create_cursor(
-        self, cursor_state: Any, runtime_lookback_window: Any = None
+        self, cursor_state: Any, runtime_lookback_window: int = 0
     ) -> ConcurrentCursor:
-        if runtime_lookback_window:
-            runtime_lookback_window = timedelta(seconds=runtime_lookback_window)
         cursor = self._cursor_factory.create(
-            stream_state=deepcopy(cursor_state), runtime_lookback_window=runtime_lookback_window
+            stream_state=deepcopy(cursor_state),
+            runtime_lookback_window=timedelta(seconds=runtime_lookback_window),
         )
         return cursor
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -875,7 +875,7 @@ def create_concurrent_cursor_from_datetime_based_cursor(
         config: Config,
         stream_state: MutableMapping[str, Any],
         message_repository: Optional[MessageRepository] = None,
-        runtime_lookback_window: Optional[int] = None,
+        runtime_lookback_window: Optional[datetime.timedelta] = None,
         **kwargs: Any,
     ) -> ConcurrentCursor:
         component_type = component_definition.get("type")
@@ -933,11 +933,6 @@ def create_concurrent_cursor_from_datetime_based_cursor(
             if evaluated_lookback_window:
                 lookback_window = parse_duration(evaluated_lookback_window)
 
-        if runtime_lookback_window and lookback_window:
-            lookback_window = max(lookback_window, runtime_lookback_window)
-        elif runtime_lookback_window:
-            lookback_window = runtime_lookback_window
-
         connector_state_converter: DateTimeStreamStateConverter
         connector_state_converter = CustomFormatConcurrentStreamStateConverter(
             datetime_format=datetime_format,
@@ -946,6 +941,18 @@ def create_concurrent_cursor_from_datetime_based_cursor(
             cursor_granularity=cursor_granularity,
         )
 
+        # Adjusts the stream state by applying the runtime lookback window.
+        # This is used to ensure correct state handling in case of failed partitions.
+        stream_state_value = stream_state.get(cursor_field.cursor_field_key)
+        if runtime_lookback_window and stream_state_value:
+            new_stream_state = (
+                connector_state_converter.parse_timestamp(stream_state_value)
+                - runtime_lookback_window
+            )
+            stream_state[cursor_field.cursor_field_key] = connector_state_converter.output_format(
+                new_stream_state
+            )
+
         start_date_runtime_value: Union[InterpolatedString, str, MinMaxDatetime]
         if isinstance(datetime_based_cursor_model.start_datetime, MinMaxDatetimeModel):
             start_date_runtime_value = self.create_min_max_datetime(