diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 1c56386883..ed140f976d 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1651,6 +1651,7 @@ def create_concurrent_cursor_from_datetime_based_cursor( slice_range=step_length, cursor_granularity=cursor_granularity, clamping_strategy=clamping_strategy, + is_compare_strictly=datetime_based_cursor_model.is_compare_strictly or False, ) def create_concurrent_cursor_from_incrementing_count_cursor( diff --git a/airbyte_cdk/sources/streams/concurrent/clamping.py b/airbyte_cdk/sources/streams/concurrent/clamping.py index 022534bc7f..59fe63db87 100644 --- a/airbyte_cdk/sources/streams/concurrent/clamping.py +++ b/airbyte_cdk/sources/streams/concurrent/clamping.py @@ -37,7 +37,7 @@ def __init__(self, is_ceiling: bool = True) -> None: def clamp(self, value: datetime) -> datetime: # type: ignore # datetime implements method from CursorValueType return_value = value.replace(hour=0, minute=0, second=0, microsecond=0) - if self._is_ceiling: + if self._is_ceiling and return_value != value: return return_value + timedelta(days=1) return return_value @@ -93,7 +93,7 @@ def clamp(self, value: datetime) -> datetime: # type: ignore # datetime implem ) delta = ( timedelta(days_diff_to_ceiling) - if self._is_ceiling + if self._is_ceiling or days_diff_to_ceiling == 0 else timedelta(days_diff_to_ceiling - 7) ) return value.replace(hour=0, minute=0, second=0, microsecond=0) + delta diff --git a/airbyte_cdk/sources/streams/concurrent/cursor.py b/airbyte_cdk/sources/streams/concurrent/cursor.py index d6dff17e26..3cf3a723b9 100644 --- a/airbyte_cdk/sources/streams/concurrent/cursor.py +++ b/airbyte_cdk/sources/streams/concurrent/cursor.py @@ -195,6 +195,7 @@ def copy_without_state(self) -> "ConcurrentCursor": slice_range=self._slice_range, cursor_granularity=self._cursor_granularity, clamping_strategy=self._clamping_strategy, + is_compare_strictly=self._is_compare_strictly, ) def __init__( @@ -213,6 +214,7 @@ def __init__( slice_range: Optional[GapType] = None, cursor_granularity: Optional[GapType] = None, clamping_strategy: ClampingStrategy = NoClamping(), + is_compare_strictly: bool = False, ) -> None: self._stream_name = stream_name self._stream_namespace = stream_namespace @@ -235,6 +237,7 @@ def __init__( # Flag to track if the logger has been triggered (per stream) self._should_be_synced_logger_triggered = False self._clamping_strategy = clamping_strategy + self._is_compare_strictly = is_compare_strictly self._is_ascending_order = True # A lock is required when closing a partition because updating the cursor's concurrent_state is @@ -483,20 +486,30 @@ def _calculate_lower_boundary_of_last_slice( def _split_per_slice_range( self, lower: CursorValueType, upper: CursorValueType, upper_is_end: bool ) -> Iterable[StreamSlice]: - if lower >= upper: + if lower > upper or (lower == upper and not self._cursor_granularity): return if self._start and upper < self._start: return lower = max(lower, self._start) if self._start else lower + # Clamp first, else a step shorter than the clamping period collapses the first window + lower = self._clamping_strategy.clamp(lower) + # With a granularity the last window includes upper, so lower == upper is one unit to read, + # unless the API rejects start == end. In a gap, the next slice already reads upper. + if lower > upper or (lower == upper and (self._is_compare_strictly or not upper_is_end)): + return if not self._slice_range or self._evaluate_upper_safely(lower, self._slice_range) >= upper: - clamped_lower = self._clamping_strategy.clamp(lower) - clamped_upper = self._clamping_strategy.clamp(upper) + # With a granularity upper is the last unit to read: clamping it would read past the end + clamped_upper = ( + upper + if upper_is_end and self._cursor_granularity + else self._clamping_strategy.clamp(upper) + ) start_value, end_value = ( - (clamped_lower, clamped_upper - self._cursor_granularity) + (lower, clamped_upper - self._cursor_granularity) if self._cursor_granularity and not upper_is_end - else (clamped_lower, clamped_upper) + else (lower, clamped_upper) ) yield StreamSlice( partition={}, @@ -516,13 +529,15 @@ def _split_per_slice_range( current_upper_boundary = min( self._evaluate_upper_safely(current_lower_boundary, self._slice_range), upper ) - has_reached_upper_boundary = current_upper_boundary >= upper - clamped_upper = ( self._clamping_strategy.clamp(current_upper_boundary) if current_upper_boundary != upper else current_upper_boundary ) + # A boundary clamped onto upper makes this the last window, else upper is dropped + has_reached_upper_boundary = ( + current_upper_boundary >= upper or clamped_upper == upper + ) clamped_lower = self._clamping_strategy.clamp(current_lower_boundary) if clamped_lower >= clamped_upper: # clamping collapsed both values which means that it is time to stop processing @@ -546,7 +561,7 @@ def _split_per_slice_range( }, ) current_lower_boundary = clamped_upper - if current_upper_boundary >= upper: + if has_reached_upper_boundary: stop_processing = True def _evaluate_upper_safely(self, lower: CursorValueType, step: GapType) -> CursorValueType: diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index cc16d3a975..522087d014 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -5403,6 +5403,41 @@ def test_incrementing_count_cursor_with_partition_router_raises_error(): ) +@pytest.mark.parametrize( + "is_compare_strictly, expected_slices", + [ + pytest.param( + False, [{"start_time": "2024-09-01", "end_time": "2024-09-01"}], id="reads_the_one_day" + ), + pytest.param(True, [], id="skipped_when_compare_strictly"), + ], +) +def test_create_concurrent_cursor_from_datetime_based_cursor_given_start_equals_end( + is_compare_strictly, expected_slices +): + concurrent_cursor = ( + ModelToComponentFactory().create_concurrent_cursor_from_datetime_based_cursor( + model_type=DatetimeBasedCursorModel, + component_definition={ + "type": "DatetimeBasedCursor", + "cursor_field": "updated_at", + "datetime_format": "%Y-%m-%d", + "start_datetime": "2024-09-01", + "end_datetime": "2024-09-01", + "step": "P30D", + "cursor_granularity": "P1D", + "is_compare_strictly": is_compare_strictly, + }, + stream_name="test", + stream_namespace=None, + stream_state={}, + config={}, + ) + ) + + assert list(concurrent_cursor.stream_slices()) == expected_slices + + def test_create_concurrent_cursor_uses_min_max_datetime_format_if_defined(): """ Validates a special case for when the start_time.datetime_format and end_time.datetime_format are defined, the date to diff --git a/unit_tests/sources/streams/concurrent/test_clamping.py b/unit_tests/sources/streams/concurrent/test_clamping.py index f66f4fe5d9..153e0f02f1 100644 --- a/unit_tests/sources/streams/concurrent/test_clamping.py +++ b/unit_tests/sources/streams/concurrent/test_clamping.py @@ -24,9 +24,14 @@ def test_when_clamp_then_remove_every_unit_smaller_than_days(self) -> None: assert result.microsecond == 0 def test_given_last_day_of_month_when_clamp_then_result_is_next_month(self) -> None: - result = self._strategy.clamp(datetime(2024, 1, 31)) + result = self._strategy.clamp(datetime(2024, 1, 31, 20, 23, 3, 2039)) assert result == datetime(2024, 2, 1) + def test_given_midnight_when_clamp_then_return_same_date(self) -> None: + midnight = datetime(2024, 1, 1) + result = self._strategy.clamp(midnight) + assert result == midnight + def test_given_is_not_ceiling_when_clamp_then_just_remove_unit_smaller_than_days(self) -> None: strategy = DayClampingStrategy(is_ceiling=False) result = strategy.clamp(datetime(2024, 1, 1, 20, 23, 3, 1029)) @@ -101,6 +106,11 @@ def test_given_not_weekday_after_target_when_clamp_then_return_next_occurrence_o _DATETIME_ON_WEDNESDAY.day + 2, ) + def test_given_not_ceiling_and_same_weekday_when_clamp_then_return_same_date(self) -> None: + strategy = WeekClampingStrategy(Weekday.TUESDAY, is_ceiling=False) + result = strategy.clamp(_DATETIME_ON_TUESDAY) + assert result == _DATETIME_ON_TUESDAY + def test_given_not_ceiling_when_clamp_then_round_down(self) -> None: strategy = WeekClampingStrategy(Weekday.FRIDAY, is_ceiling=False) result = strategy.clamp(_DATETIME_ON_WEDNESDAY) diff --git a/unit_tests/sources/streams/concurrent/test_cursor.py b/unit_tests/sources/streams/concurrent/test_cursor.py index 14cfcd0c07..771adeab9a 100644 --- a/unit_tests/sources/streams/concurrent/test_cursor.py +++ b/unit_tests/sources/streams/concurrent/test_cursor.py @@ -18,6 +18,7 @@ from airbyte_cdk.sources.streams.concurrent.clamping import ( ClampingEndProvider, ClampingStrategy, + DayClampingStrategy, MonthClampingStrategy, WeekClampingStrategy, Weekday, @@ -564,7 +565,7 @@ def test_given_small_slice_range_with_granularity_when_generate_slices_then_crea ] @freezegun.freeze_time(time_to_freeze=datetime.fromtimestamp(50, timezone.utc)) - def test_given_difference_between_slices_match_slice_range_and_cursor_granularity_when_generate_slices_then_create_one_slice( + def test_given_difference_between_slices_match_slice_range_and_cursor_granularity_when_generate_slices_then_create_one_slice_for_the_gap_and_read_the_end_unit( self, ): start = datetime.fromtimestamp(1, timezone.utc) @@ -608,6 +609,14 @@ def test_given_difference_between_slices_match_slice_range_and_cursor_granularit _SLICE_BOUNDARY_FIELDS[1]: 40, }, ), + # the last state slice ends at the end, which the last window includes + StreamSlice( + partition={}, + cursor_slice={ + _SLICE_BOUNDARY_FIELDS[0]: 50, + _SLICE_BOUNDARY_FIELDS[1]: 50, + }, + ), ] @freezegun.freeze_time(time_to_freeze=datetime.fromtimestamp(50, timezone.utc)) @@ -1124,6 +1133,255 @@ def test_given_weekly_clamp_and_granularity_when_stream_slices_then_slice_per_we {"lower_boundary": "2024-01-23T00:00:00Z", "upper_boundary": "2024-01-29T00:00:00Z"}, ] + @freezegun.freeze_time(time_to_freeze=datetime(2024, 1, 30, 10, tzinfo=timezone.utc)) + def test_given_weekly_clamp_and_now_on_target_weekday_when_stream_slices_then_include_last_complete_week( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 1, 2, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + WeekClampingStrategy(Weekday.TUESDAY, is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(days=7), + granularity=timedelta(days=1), + clamping_strategy=WeekClampingStrategy(Weekday.TUESDAY), + ) + stream_slices = list(cursor.stream_slices()) + assert stream_slices[-1] == { + "lower_boundary": "2024-01-23T00:00:00Z", + "upper_boundary": "2024-01-29T00:00:00Z", + } + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 1, 5, 12, tzinfo=timezone.utc)) + def test_given_daily_clamp_and_midnight_boundaries_when_stream_slices_then_do_not_skip_days( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 1, 1, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + DayClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(seconds=1), + ), + slice_range=timedelta(days=1), + granularity=timedelta(seconds=1), + clamping_strategy=DayClampingStrategy(), + ) + stream_slices = list(cursor.stream_slices()) + assert stream_slices == [ + {"lower_boundary": "2024-01-01T00:00:00Z", "upper_boundary": "2024-01-01T23:59:59Z"}, + {"lower_boundary": "2024-01-02T00:00:00Z", "upper_boundary": "2024-01-02T23:59:59Z"}, + {"lower_boundary": "2024-01-03T00:00:00Z", "upper_boundary": "2024-01-03T23:59:59Z"}, + {"lower_boundary": "2024-01-04T00:00:00Z", "upper_boundary": "2024-01-04T23:59:59Z"}, + ] + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 1, 5, 12, tzinfo=timezone.utc)) + def test_given_daily_clamp_and_one_window_when_stream_slices_then_end_at_the_clamped_end( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 1, 1, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + DayClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(seconds=1), + ), + slice_range=timedelta(days=30), + granularity=timedelta(seconds=1), + clamping_strategy=DayClampingStrategy(), + ) + stream_slices = list(cursor.stream_slices()) + assert stream_slices == [ + {"lower_boundary": "2024-01-01T00:00:00Z", "upper_boundary": "2024-01-04T23:59:59Z"}, + ] + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 1, 4, 12, tzinfo=timezone.utc)) + def test_given_daily_clamp_and_step_not_in_whole_days_when_stream_slices_then_read_the_last_day( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 1, 1, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + DayClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(hours=36), + granularity=timedelta(days=1), + clamping_strategy=DayClampingStrategy(), + ) + stream_slices = list(cursor.stream_slices()) + assert stream_slices == [ + {"lower_boundary": "2024-01-01T00:00:00Z", "upper_boundary": "2024-01-03T00:00:00Z"}, + ] + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 6, 15, tzinfo=timezone.utc)) + def test_given_monthly_clamp_and_step_shorter_than_distance_to_next_month_when_stream_slices_then_slice_per_month( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 3, 2, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + MonthClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(days=27), + granularity=timedelta(days=1), + clamping_strategy=MonthClampingStrategy(), + ) + stream_slices = list(cursor.stream_slices()) + assert stream_slices == [ + {"lower_boundary": "2024-04-01T00:00:00Z", "upper_boundary": "2024-04-30T00:00:00Z"}, + {"lower_boundary": "2024-05-01T00:00:00Z", "upper_boundary": "2024-05-31T00:00:00Z"}, + ] + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 3, 15, tzinfo=timezone.utc)) + def test_given_monthly_clamp_and_no_complete_month_left_when_stream_slices_then_do_not_read_past_the_end( + self, + ) -> None: + cursor = self._cursor( + start=datetime(2024, 2, 10, tzinfo=timezone.utc), + end_provider=ClampingEndProvider( + MonthClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(days=27), + granularity=timedelta(days=1), + clamping_strategy=MonthClampingStrategy(), + ) + assert list(cursor.stream_slices()) == [] + + @freezegun.freeze_time(time_to_freeze=datetime(2024, 10, 2, tzinfo=timezone.utc)) + def test_given_state_at_clamped_end_when_stream_slices_then_do_not_read_past_the_end( + self, + ) -> None: + cursor = ConcurrentCursor( + _A_STREAM_NAME, + _A_STREAM_NAMESPACE, + { + "state_type": ConcurrencyCompatibleStateType.date_range.value, + "slices": [{"start": "2024-01-01T00:00:00Z", "end": "2024-09-30T00:00:00Z"}], + }, + self._message_repository, + self._state_manager, + CustomFormatConcurrentStreamStateConverter( + "%Y-%m-%dT%H:%M:%SZ", is_sequential_state=_NOT_SEQUENTIAL + ), + CursorField(_A_CURSOR_FIELD_KEY), + _SLICE_BOUNDARY_FIELDS, + datetime(2024, 1, 1, tzinfo=timezone.utc), + ClampingEndProvider( + MonthClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(days=27), + cursor_granularity=timedelta(days=1), + clamping_strategy=MonthClampingStrategy(), + ) + + # the end is 2024-09-30: clamping [end, end] would move it into October, which is not over + assert list(cursor.stream_slices()) == [] + + +@freezegun.freeze_time(time_to_freeze=datetime.fromtimestamp(50, timezone.utc)) +@pytest.mark.parametrize( + "cursor_granularity, is_compare_strictly, expected_slices", + [ + pytest.param( + timedelta(seconds=1), + False, + [{_LOWER_SLICE_BOUNDARY_FIELD: 50, _UPPER_SLICE_BOUNDARY_FIELD: 50}], + id="one_unit_window_with_granularity", + ), + pytest.param(timedelta(seconds=1), True, [], id="skipped_when_compare_strictly"), + pytest.param(None, False, [], id="skipped_without_granularity"), + ], +) +def test_given_start_equals_end_when_stream_slices_then_read_the_end_unit_if_the_window_includes_it( + cursor_granularity: Optional[timedelta], is_compare_strictly: bool, expected_slices +) -> None: + cursor = ConcurrentCursor( + _A_STREAM_NAME, + _A_STREAM_NAMESPACE, + deepcopy(_NO_STATE), + Mock(spec=MessageRepository), + Mock(spec=ConnectorStateManager), + EpochValueConcurrentStreamStateConverter(is_sequential_state=True), + CursorField(_A_CURSOR_FIELD_KEY), + _SLICE_BOUNDARY_FIELDS, + datetime.fromtimestamp(50, timezone.utc), + EpochValueConcurrentStreamStateConverter.get_end_provider(), + slice_range=timedelta(seconds=10) if cursor_granularity else None, + cursor_granularity=cursor_granularity, + is_compare_strictly=is_compare_strictly, + ) + + assert list(cursor.stream_slices()) == expected_slices + assert list(cursor.copy_without_state().stream_slices()) == expected_slices + + +@freezegun.freeze_time(time_to_freeze=datetime(2024, 3, 5, 12, tzinfo=timezone.utc)) +@pytest.mark.parametrize( + "state_slices, is_compare_strictly, expected_slices", + [ + pytest.param( + [ + {"start": "2024-01-01T00:00:00Z", "end": "2024-01-30T10:00:00Z"}, + {"start": "2024-02-01T00:00:00Z", "end": "2024-03-01T00:00:00Z"}, + ], + False, + [ + { + _LOWER_SLICE_BOUNDARY_FIELD: "2024-03-01T00:00:00Z", + _UPPER_SLICE_BOUNDARY_FIELD: "2024-03-04T00:00:00Z", + } + ], + id="gap_left_empty_by_the_clamp", + ), + pytest.param( + [{"start": "2024-01-01T00:00:00Z", "end": "2024-03-03T10:00:00Z"}], + True, + [], + id="end_unit_with_compare_strictly", + ), + ], +) +def test_given_daily_clamp_moves_lower_onto_upper_when_stream_slices_then_skip_the_window( + state_slices, is_compare_strictly: bool, expected_slices +) -> None: + cursor = ConcurrentCursor( + _A_STREAM_NAME, + _A_STREAM_NAMESPACE, + { + "state_type": ConcurrencyCompatibleStateType.date_range.value, + "slices": deepcopy(state_slices), + }, + Mock(spec=MessageRepository), + Mock(spec=ConnectorStateManager), + CustomFormatConcurrentStreamStateConverter( + "%Y-%m-%dT%H:%M:%SZ", is_sequential_state=_NOT_SEQUENTIAL + ), + CursorField(_A_CURSOR_FIELD_KEY), + _SLICE_BOUNDARY_FIELDS, + datetime(2024, 1, 1, tzinfo=timezone.utc), + ClampingEndProvider( + DayClampingStrategy(is_ceiling=False), + CustomFormatConcurrentStreamStateConverter.get_end_provider(), + granularity=timedelta(days=1), + ), + slice_range=timedelta(days=10), + cursor_granularity=timedelta(days=1), + clamping_strategy=DayClampingStrategy(), + is_compare_strictly=is_compare_strictly, + ) + + assert list(cursor.stream_slices()) == expected_slices + _SHOULD_BE_SYNCED_START = 10