update tests and migrate stop condition too

Author: maxime.c

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -516,6 +516,13 @@ def _get_cursor(self, record: Record) -> ConcurrentCursor:
             raise ValueError(
                 "Invalid state as stream slices that are emitted should refer to an existing cursor"
             )
+
+        if self._use_global_cursor:
+            return self._create_cursor(
+                self._global_cursor,
+                self._lookback_window if self._global_cursor else 0,
+            )
+
         partition_key = self._to_partition_key(record.associated_slice.partition)
         if (
             partition_key not in self._cursor_per_partition

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -603,7 +603,7 @@
     WeekClampingStrategy,
     Weekday,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField
+from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField, Cursor
 from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import (
     CustomFormatConcurrentStreamStateConverter,
     DateTimeStreamStateConverter,
@@ -1932,17 +1932,17 @@ def create_declarative_stream(
             and hasattr(model.incremental_sync, "is_data_feed")
             and model.incremental_sync.is_data_feed
         )
-        client_side_incremental_sync = None
-        if (
+        client_side_filetering_enabled = (
             model.incremental_sync
             and hasattr(model.incremental_sync, "is_client_side_incremental")
             and model.incremental_sync.is_client_side_incremental
-        ):
+        )
+        concurrent_cursor = None
+        if stop_condition_on_cursor or client_side_filetering_enabled:
             stream_slicer = self._build_stream_slicer_from_partition_router(
                 model.retriever, config, stream_name=model.name
             )
-            cursor = self._build_concurrent_cursor(model, stream_slicer, config)
-            client_side_incremental_sync = {"cursor": cursor}
+            concurrent_cursor = self._build_concurrent_cursor(model, stream_slicer, config)
 
         if model.incremental_sync and isinstance(model.incremental_sync, DatetimeBasedCursorModel):
             cursor_model = model.incremental_sync
@@ -2017,8 +2017,10 @@ def create_declarative_stream(
             primary_key=primary_key,
             stream_slicer=combined_slicers,
             request_options_provider=request_options_provider,
-            stop_condition_on_cursor=stop_condition_on_cursor,
-            client_side_incremental_sync=client_side_incremental_sync,
+            stop_condition_cursor=concurrent_cursor,
+            client_side_incremental_sync={"cursor": concurrent_cursor}
+            if client_side_filetering_enabled
+            else None,
             transformations=transformations,
             file_uploader=file_uploader,
             incremental_sync=model.incremental_sync,
@@ -3194,7 +3196,7 @@ def create_simple_retriever(
         primary_key: Optional[Union[str, List[str], List[List[str]]]],
         stream_slicer: Optional[StreamSlicer],
         request_options_provider: Optional[RequestOptionsProvider] = None,
-        stop_condition_on_cursor: bool = False,
+        stop_condition_cursor: Optional[Cursor] = None,
         client_side_incremental_sync: Optional[Dict[str, Any]] = None,
         transformations: List[RecordTransformation],
         file_uploader: Optional[DefaultFileUploader] = None,
@@ -3325,7 +3327,7 @@ def _get_url() -> str:
                 ),
             )
 
-        cursor_used_for_stop_condition = cursor if stop_condition_on_cursor else None
+        cursor_used_for_stop_condition = cursor if stop_condition_cursor else None
         paginator = (
             self._create_component_from_model(
                 model=model.paginator,

airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py

@@ -7,11 +7,10 @@
 
 import requests
 
-from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
     PaginationStrategy,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.types import Record
 
 
@@ -29,8 +28,7 @@ def is_met(self, record: Record) -> bool:
 class CursorStopCondition(PaginationStopCondition):
     def __init__(
         self,
-        cursor: DeclarativeCursor
-        | ConcurrentCursor,  # migrate to use both old and concurrent versions
+        cursor: Cursor,
     ):
         self._cursor = cursor
 

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -492,7 +492,7 @@ def should_be_synced(self, record: Record) -> bool:
         except ValueError:
             self._log_for_record_without_cursor_value()
             return True
-        return self.start <= record_cursor_value
+        return self.start <= record_cursor_value <= self._end_provider()
 
     def _log_for_record_without_cursor_value(self) -> None:
         if not self._should_be_synced_logger_triggered:

unit_tests/sources/declarative/extractors/test_record_filter.py

@@ -1,7 +1,9 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+from datetime import datetime, timedelta, timezone
 from typing import List, Mapping, Optional
+from unittest.mock import Mock
 
 import pytest
 
@@ -11,11 +13,10 @@
     RecordFilter,
 )
 from airbyte_cdk.sources.declarative.incremental import (
-    CursorFactory,
-    DatetimeBasedCursor,
-    GlobalSubstreamCursor,
-    PerPartitionWithGlobalCursor,
+    ConcurrentPerPartitionCursor,
+    ConcurrentCursorFactory,
 )
+from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.models import (
     CustomRetriever,
@@ -24,7 +25,11 @@
 )
 from airbyte_cdk.sources.declarative.partition_routers import SubstreamPartitionRouter
 from airbyte_cdk.sources.declarative.types import StreamSlice
+from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import (
+    CustomFormatConcurrentStreamStateConverter,
+)
 from airbyte_cdk.sources.types import Record
+from airbyte_cdk.utils.datetime_helpers import ab_datetime_parse, ab_datetime_now
 
 DATE_FORMAT = "%Y-%m-%d"
 RECORDS_TO_FILTER_DATE_FORMAT = [
@@ -272,25 +277,27 @@ def test_client_side_record_filter_decorator_no_parent_stream(
     records_to_filter: List[Mapping],
     expected_record_ids: List[int],
 ):
-    date_time_based_cursor = DatetimeBasedCursor(
-        start_datetime=MinMaxDatetime(
-            datetime="2021-01-01", datetime_format=DATE_FORMAT, parameters={}
+    datetime_based_cursor = ConcurrentCursor(
+        stream_name="any_stream",
+        stream_namespace=None,
+        stream_state=stream_state,
+        message_repository=Mock(),
+        connector_state_manager=Mock(),
+        connector_state_converter=CustomFormatConcurrentStreamStateConverter(
+            datetime_format=datetime_format
         ),
-        end_datetime=MinMaxDatetime(datetime=end_datetime, parameters={}) if end_datetime else None,
-        step="P10Y",
-        cursor_field=InterpolatedString.create("created_at", parameters={}),
-        datetime_format=datetime_format,
-        cursor_granularity="P1D",
-        config={},
-        parameters={},
+        cursor_field=CursorField("created_at"),
+        slice_boundary_fields=("start", "end"),
+        start=datetime(2021, 1, 1, tzinfo=timezone.utc),
+        end_provider=lambda: ab_datetime_parse(end_datetime) if end_datetime else ab_datetime_now(),
+        slice_range=timedelta(days=365 * 10),
     )
-    date_time_based_cursor.set_initial_state(stream_state)
 
     record_filter_decorator = ClientSideIncrementalRecordFilterDecorator(
         config={},
         condition=record_filter_expression,
         parameters={},
-        cursor=date_time_based_cursor,
+        cursor=datetime_based_cursor,
     )
 
     filtered_records = list(
@@ -341,7 +348,7 @@ def test_client_side_record_filter_decorator_no_parent_stream(
                     }
                 ],
             },
-            "per_partition_with_global",
+            "global_substream",
             [2, 3],
         ),
         # Use PerPartitionWithGlobalCursor with partition state missing, global cursor used
@@ -363,23 +370,26 @@ def test_client_side_record_filter_decorator_no_parent_stream(
 def test_client_side_record_filter_decorator_with_cursor_types(
     stream_state: Optional[Mapping], cursor_type: str, expected_record_ids: List[int]
 ):
-    def date_time_based_cursor_factory() -> DatetimeBasedCursor:
-        return DatetimeBasedCursor(
-            start_datetime=MinMaxDatetime(
-                datetime="2021-01-01", datetime_format=DATE_FORMAT, parameters={}
-            ),
-            end_datetime=MinMaxDatetime(
-                datetime="2021-01-05", datetime_format=DATE_FORMAT, parameters={}
+    def date_time_based_cursor_factory(stream_state, runtime_lookback_window) -> ConcurrentCursor:
+        return ConcurrentCursor(
+            stream_name="any_stream",
+            stream_namespace=None,
+            stream_state=stream_state,
+            message_repository=Mock(),
+            connector_state_manager=Mock(),
+            connector_state_converter=CustomFormatConcurrentStreamStateConverter(
+                datetime_format=DATE_FORMAT
             ),
-            step="P10Y",
-            cursor_field=InterpolatedString.create("created_at", parameters={}),
-            datetime_format=DATE_FORMAT,
-            cursor_granularity="P1D",
-            config={},
-            parameters={},
+            cursor_field=CursorField("created_at"),
+            slice_boundary_fields=("start", "end"),
+            start=datetime(2021, 1, 1, tzinfo=timezone.utc),
+            end_provider=lambda: datetime(2021, 1, 5, tzinfo=timezone.utc),
+            slice_range=timedelta(days=365 * 10),
+            cursor_granularity=timedelta(days=1),
+            lookback_window=runtime_lookback_window,
         )
 
-    date_time_based_cursor = date_time_based_cursor_factory()
+    date_time_based_cursor = date_time_based_cursor_factory(stream_state, timedelta(0))
 
     substream_cursor = None
     partition_router = SubstreamPartitionRouter(
@@ -401,29 +411,26 @@ def date_time_based_cursor_factory() -> DatetimeBasedCursor:
     if cursor_type == "datetime":
         # Use only DatetimeBasedCursor
         pass  # No additional cursor needed
-    elif cursor_type == "global_substream":
+    elif cursor_type in ["global_substream", "per_partition_with_global"]:
         # Create GlobalSubstreamCursor instance
-        substream_cursor = GlobalSubstreamCursor(
-            stream_cursor=date_time_based_cursor,
+        substream_cursor = ConcurrentPerPartitionCursor(
+            cursor_factory=ConcurrentCursorFactory(date_time_based_cursor_factory),
             partition_router=partition_router,
-        )
-        if stream_state:
-            substream_cursor.set_initial_state(stream_state)
-    elif cursor_type == "per_partition_with_global":
-        # Create PerPartitionWithGlobalCursor instance
-        substream_cursor = PerPartitionWithGlobalCursor(
-            cursor_factory=CursorFactory(date_time_based_cursor_factory),
-            partition_router=partition_router,
-            stream_cursor=date_time_based_cursor,
+            stream_name="a_stream",
+            stream_namespace=None,
+            stream_state=stream_state,
+            message_repository=Mock(),
+            connector_state_manager=Mock(),
+            connector_state_converter=CustomFormatConcurrentStreamStateConverter(
+                datetime_format=DATE_FORMAT
+            ),
+            cursor_field=CursorField("created_at"),
+            use_global_cursor=cursor_type == "global_substream",
+            attempt_to_create_cursor_if_not_provided=True,
         )
     else:
         raise ValueError(f"Unsupported cursor type: {cursor_type}")
 
-    if substream_cursor and stream_state:
-        substream_cursor.set_initial_state(stream_state)
-    elif stream_state:
-        date_time_based_cursor.set_initial_state(stream_state)
-
     # Create the record_filter_decorator with appropriate cursor
     record_filter_decorator = ClientSideIncrementalRecordFilterDecorator(
         config={},

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -188,7 +188,12 @@
 
 transformer = ManifestComponentTransformer()
 
-input_config = {"apikey": "verysecrettoken", "repos": ["airbyte", "airbyte-cloud"]}
+input_config = {
+    "apikey": "verysecrettoken",
+    "repos": ["airbyte", "airbyte-cloud"],
+    "start_time": "2024-01-01T00:00:00.000+00:00",
+    "end_time": "2025-01-01T00:00:00.000+00:00",
+}
 
 
 def get_factory_with_parameters(
@@ -1997,7 +2002,7 @@ def test_create_default_paginator():
             "subcomponent_field_with_hint",
             DpathExtractor(
                 field_path=[],
-                config={"apikey": "verysecrettoken", "repos": ["airbyte", "airbyte-cloud"]},
+                config=input_config,
                 decoder=JsonDecoder(parameters={}),
                 parameters={},
             ),
@@ -2013,7 +2018,7 @@ def test_create_default_paginator():
             "subcomponent_field_with_hint",
             DpathExtractor(
                 field_path=[],
-                config={"apikey": "verysecrettoken", "repos": ["airbyte", "airbyte-cloud"]},
+                config=input_config,
                 parameters={},
             ),
             None,
@@ -2102,11 +2107,11 @@ def test_create_default_paginator():
                 pagination_strategy=OffsetIncrement(
                     page_size=10,
                     extractor=None,
-                    config={"apikey": "verysecrettoken", "repos": ["airbyte", "airbyte-cloud"]},
+                    config=input_config,
                     parameters={},
                 ),
                 url_base="https://physical_100.com",
-                config={"apikey": "verysecrettoken", "repos": ["airbyte", "airbyte-cloud"]},
+                config=input_config,
                 parameters={"decoder": {"type": "JsonDecoder"}},
             ),
             None,