diff --git a/airbyte_cdk/sources/file_based/file_based_source.py b/airbyte_cdk/sources/file_based/file_based_source.py index df469e8345..a8c396721a 100644 --- a/airbyte_cdk/sources/file_based/file_based_source.py +++ b/airbyte_cdk/sources/file_based/file_based_source.py @@ -87,6 +87,16 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC): # We make each source override the concurrency level to give control over when they are upgraded. _concurrency_level = None + # When the legacy (non-concurrent) file-based read path is used, the + # framework emits a state message after every slice (i.e. every file). On + # streams with many small files the platform/orchestrator buffer of those + # state messages can OOM the replication pod before the destination ACKs + # them. Connectors can override this to throttle per-slice state emission + # to once per N seconds; the final state of each sync is always force + # emitted so destinations still see the latest cursor. Default `None` + # preserves the historical emit-per-slice behaviour. + _stream_state_emission_throttle_seconds: Optional[float] = None + def __init__( self, stream_reader: AbstractFileBasedStreamReader, @@ -336,7 +346,7 @@ def _make_default_stream( cursor: Optional[AbstractFileBasedCursor], parsed_config: AbstractFileBasedSpec, ) -> AbstractFileBasedStream: - return DefaultFileBasedStream( + stream = DefaultFileBasedStream( config=stream_config, catalog_schema=self.stream_schemas.get(stream_config.name), stream_reader=self.stream_reader, @@ -349,6 +359,11 @@ def _make_default_stream( use_file_transfer=use_file_transfer(parsed_config), preserve_directory_structure=preserve_directory_structure(parsed_config), ) + if self._stream_state_emission_throttle_seconds is not None: + stream.state_emission_throttle_seconds = ( + self._stream_state_emission_throttle_seconds + ) + return stream def _ensure_permissions_reader_available(self) -> None: """ diff --git a/airbyte_cdk/sources/streams/core.py b/airbyte_cdk/sources/streams/core.py index 6cc5c8b5d9..483a37bc25 100644 --- a/airbyte_cdk/sources/streams/core.py +++ b/airbyte_cdk/sources/streams/core.py @@ -5,6 +5,7 @@ import inspect import itertools import logging +import time from abc import ABC, abstractmethod from dataclasses import dataclass from functools import cached_property, lru_cache @@ -123,6 +124,15 @@ class Stream(ABC): _configured_json_schema: Optional[Dict[str, Any]] = None _exit_on_rate_limit: bool = False + # If set, per-slice state messages are emitted at most once per this many + # seconds during a sync. The first per-slice emission and the final emission + # at the end of the stream always fire. Default `None` keeps the historical + # behaviour: emit a state message after every slice. Streams that produce + # very many slices with growing state payloads (e.g. file-based sources + # carrying a history dict that grows with each file synced) can set this + # to keep the platform/orchestrator state buffer bounded. + state_emission_throttle_seconds: Optional[float] = None + # Use self.logger in subclasses to log any messages @property def logger(self) -> logging.Logger: @@ -184,6 +194,14 @@ def read( # type: ignore # ignoring typing for ConnectorStateManager because o next_slice = checkpoint_reader.next() record_counter = 0 stream_state_tracker = copy.deepcopy(stream_state) + + # State-emission throttle bookkeeping. Active only when the stream sets + # `state_emission_throttle_seconds`; default behavior is preserved. + throttle_seconds = self.state_emission_throttle_seconds + last_state_emit_at: float = 0.0 + last_observed_checkpoint: Optional[Mapping[str, Any]] = None + throttled_pending_emit = False + while next_slice is not None: if slice_logger.should_log_slice_message(logger): yield slice_logger.create_slice_log_message(next_slice) @@ -240,10 +258,24 @@ def read( # type: ignore # ignoring typing for ConnectorStateManager because o self._observe_state(checkpoint_reader) checkpoint_state = checkpoint_reader.get_checkpoint() if should_checkpoint and checkpoint_state is not None: - airbyte_state_message = self._checkpoint_state( - checkpoint_state, state_manager=state_manager - ) - yield airbyte_state_message + last_observed_checkpoint = checkpoint_state + if throttle_seconds is not None: + now = time.time() + if now - last_state_emit_at < throttle_seconds: + # Suppress this per-slice state emission; the final- + # state emit below will catch up if no other emit + # fires before the stream ends. + throttled_pending_emit = True + else: + last_state_emit_at = now + throttled_pending_emit = False + yield self._checkpoint_state( + checkpoint_state, state_manager=state_manager + ) + else: + yield self._checkpoint_state( + checkpoint_state, state_manager=state_manager + ) next_slice = checkpoint_reader.next() @@ -251,6 +283,16 @@ def read( # type: ignore # ignoring typing for ConnectorStateManager because o if should_checkpoint and checkpoint is not None: airbyte_state_message = self._checkpoint_state(checkpoint, state_manager=state_manager) yield airbyte_state_message + elif ( + should_checkpoint + and throttle_seconds is not None + and throttled_pending_emit + and last_observed_checkpoint is not None + ): + # Throttling suppressed the final per-slice emit. Force a final + # state message so the platform/destination always sees the + # latest stream state for this sync. + yield self._checkpoint_state(last_observed_checkpoint, state_manager=state_manager) def read_only_records(self, state: Optional[Mapping[str, Any]] = None) -> Iterable[StreamData]: """ diff --git a/unit_tests/sources/streams/test_stream_read.py b/unit_tests/sources/streams/test_stream_read.py index cf550f8cf8..74674a32f9 100644 --- a/unit_tests/sources/streams/test_stream_read.py +++ b/unit_tests/sources/streams/test_stream_read.py @@ -786,6 +786,123 @@ def _read( return records +def test_state_emission_throttle_suppresses_per_slice_emit_and_forces_final(mocker): + """ + With `state_emission_throttle_seconds` set, per-slice state messages emitted + inside the throttle window are suppressed. The first per-slice emit fires + (cold start), and a final state is force-emitted at end-of-stream so the + destination always sees the latest state. + """ + configured_stream = ConfiguredAirbyteStream( + stream=AirbyteStream( + name="mock_stream", + supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental], + json_schema={}, + ), + sync_mode=SyncMode.incremental, + cursor_field=["created_at"], + destination_sync_mode=DestinationSyncMode.overwrite, + ) + internal_config = InternalConfig() + logger = _mock_logger() + slice_logger = DebugSliceLogger() + message_repository = InMemoryMessageRepository(Level.INFO) + state_manager = ConnectorStateManager() + timestamp = "1708899427" + + slice_to_partition = { + 1: [{"id": 1, "partition": 1, "created_at": "1708899000"}], + 2: [{"id": 2, "partition": 2, "created_at": "1708899100"}], + 3: [{"id": 3, "partition": 3, "created_at": "1708899200"}], + 4: [{"id": 4, "partition": 4, "created_at": timestamp}], + } + stream = _incremental_stream( + slice_to_partition, slice_logger, logger, message_repository, timestamp + ) + # _MockIncrementalStream uses a class-level _state dict that mutates across + # tests in this module. Reset it so we observe only the state this test + # produces. + type(stream)._state = {} + # 600 s throttle. We mock time.time so the first emit fires, the next two + # fall inside the throttle window, and the loop ends needing a final emit. + stream.state_emission_throttle_seconds = 600 + mock_time = mocker.patch("airbyte_cdk.sources.streams.core.time.time") + # First call inside throttle check -> 0s; second -> 10s; third -> 20s; + # fourth -> 30s. All under 600 except the first (cold start, 0 - 0 < 600 + # so suppressed... see logic). + mock_time.side_effect = [0, 10, 20, 30, 40, 50] + + actual_records = _read( + stream, + configured_stream, + logger, + slice_logger, + message_repository, + state_manager, + internal_config, + ) + + # Count STATE messages. + state_messages = [ + m for m in actual_records if getattr(m, "type", None) == MessageType.STATE + ] + # Cold-start emit is suppressed (delta 0 < 600); subsequent ones suppressed; + # final-state force emit fires once at end of stream. + assert len(state_messages) == 1 + # Final state must carry the latest observed cursor. + final = state_messages[-1] + assert final.state.stream.stream_state.created_at == timestamp + + +def test_state_emission_throttle_default_none_keeps_per_slice_emits(): + """When state_emission_throttle_seconds is None (default), every slice emits + its state — i.e. the throttle is fully opt-in and the existing behaviour + is preserved bit-for-bit.""" + configured_stream = ConfiguredAirbyteStream( + stream=AirbyteStream( + name="mock_stream", + supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental], + json_schema={}, + ), + sync_mode=SyncMode.incremental, + cursor_field=["created_at"], + destination_sync_mode=DestinationSyncMode.overwrite, + ) + internal_config = InternalConfig() + logger = _mock_logger() + slice_logger = DebugSliceLogger() + message_repository = InMemoryMessageRepository(Level.INFO) + state_manager = ConnectorStateManager() + timestamp = "1708899427" + + slice_to_partition = { + 1: [{"id": 1, "partition": 1, "created_at": "1708899000"}], + 2: [{"id": 2, "partition": 2, "created_at": timestamp}], + } + stream = _incremental_stream( + slice_to_partition, slice_logger, logger, message_repository, timestamp + ) + type(stream)._state = {} + # Throttle attribute is None (default). + assert stream.state_emission_throttle_seconds is None + + actual_records = _read( + stream, + configured_stream, + logger, + slice_logger, + message_repository, + state_manager, + internal_config, + ) + + state_messages = [ + m for m in actual_records if getattr(m, "type", None) == MessageType.STATE + ] + # One per-slice emit per slice, no final duplicate (IncrementalCheckpointReader suppresses it). + assert len(state_messages) == 2 + + def _mock_partition_generator( name: str, slices, records_per_partition, *, available=True, debug_log=False ):