Introduce FinalStateCursor to emit state messages at the end of full refresh syncs (#35905)

Co-authored-by: brianjlai <[email protected]>

Author: erohmensing

airbyte-cdk/python/airbyte_cdk/sources/file_based/file_based_source.py

@@ -36,7 +36,7 @@
 from airbyte_cdk.sources.file_based.stream.concurrent.cursor import (
     AbstractConcurrentFileBasedCursor,
     FileBasedConcurrentCursor,
-    FileBasedNoopCursor,
+    FileBasedFinalStateCursor,
 )
 from airbyte_cdk.sources.file_based.stream.cursor import AbstractFileBasedCursor
 from airbyte_cdk.sources.message.repository import InMemoryMessageRepository, MessageRepository
@@ -170,7 +170,9 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
                 sync_mode = self._get_sync_mode_from_catalog(stream_config.name)
 
                 if sync_mode == SyncMode.full_refresh and hasattr(self, "_concurrency_level") and self._concurrency_level is not None:
-                    cursor = FileBasedNoopCursor(stream_config)
+                    cursor = FileBasedFinalStateCursor(
+                        stream_config=stream_config, stream_namespace=None, message_repository=self.message_repository
+                    )
                     stream = FileBasedStreamFacade.create_from_stream(
                         self._make_default_stream(stream_config, cursor), self, self.logger, stream_state, cursor
                     )

airbyte-cdk/python/airbyte_cdk/sources/file_based/stream/concurrent/adapters.py

@@ -18,7 +18,7 @@
 from airbyte_cdk.sources.file_based.file_types.file_type_parser import FileTypeParser
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 from airbyte_cdk.sources.file_based.stream import AbstractFileBasedStream
-from airbyte_cdk.sources.file_based.stream.concurrent.cursor import FileBasedNoopCursor
+from airbyte_cdk.sources.file_based.stream.concurrent.cursor import FileBasedFinalStateCursor
 from airbyte_cdk.sources.file_based.stream.cursor import AbstractFileBasedCursor
 from airbyte_cdk.sources.file_based.types import StreamSlice
 from airbyte_cdk.sources.message import MessageRepository
@@ -71,7 +71,7 @@ def create_from_stream(
                 partition_generator=FileBasedStreamPartitionGenerator(
                     stream,
                     message_repository,
-                    SyncMode.full_refresh if isinstance(cursor, FileBasedNoopCursor) else SyncMode.incremental,
+                    SyncMode.full_refresh if isinstance(cursor, FileBasedFinalStateCursor) else SyncMode.incremental,
                     [cursor_field] if cursor_field is not None else None,
                     state,
                     cursor,

airbyte-cdk/python/airbyte_cdk/sources/file_based/stream/concurrent/cursor/__init__.py

@@ -1,5 +1,5 @@
 from .abstract_concurrent_file_based_cursor import AbstractConcurrentFileBasedCursor
 from .file_based_concurrent_cursor import FileBasedConcurrentCursor
-from .file_based_noop_cursor import FileBasedNoopCursor
+from .file_based_final_state_cursor import FileBasedFinalStateCursor
 
-__all__ = ["AbstractConcurrentFileBasedCursor", "FileBasedConcurrentCursor", "FileBasedNoopCursor"]
+__all__ = ["AbstractConcurrentFileBasedCursor", "FileBasedConcurrentCursor", "FileBasedFinalStateCursor"]

airbyte-cdk/python/airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_noop_cursor.py renamed to airbyte-cdk/python/airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py

@@ -4,26 +4,39 @@
 
 import logging
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Iterable, List, MutableMapping
+from typing import TYPE_CHECKING, Any, Iterable, List, MutableMapping, Optional
 
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 from airbyte_cdk.sources.file_based.stream.concurrent.cursor.abstract_concurrent_file_based_cursor import AbstractConcurrentFileBasedCursor
 from airbyte_cdk.sources.file_based.types import StreamState
+from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams import FULL_REFRESH_SENTINEL_STATE_KEY
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamPartition
 
 
-class FileBasedNoopCursor(AbstractConcurrentFileBasedCursor):
-    def __init__(self, stream_config: FileBasedStreamConfig, **kwargs: Any):
-        pass
+class FileBasedFinalStateCursor(AbstractConcurrentFileBasedCursor):
+    """Cursor that is used to guarantee at least one state message is emitted for a concurrent file-based stream."""
+
+    def __init__(
+        self, stream_config: FileBasedStreamConfig, message_repository: MessageRepository, stream_namespace: Optional[str], **kwargs: Any
+    ):
+        self._stream_name = stream_config.name
+        self._stream_namespace = stream_namespace
+        self._message_repository = message_repository
+        # Normally the connector state manager operates at the source-level. However, we only need it to write the sentinel
+        # state message rather than manage overall source state. This is also only temporary as we move to the resumable
+        # full refresh world where every stream uses a FileBasedConcurrentCursor with incremental state.
+        self._connector_state_manager = ConnectorStateManager(stream_instance_map={})
 
     @property
     def state(self) -> MutableMapping[str, Any]:
-        return {}
+        return {FULL_REFRESH_SENTINEL_STATE_KEY: True}
 
     def observe(self, record: Record) -> None:
         pass
@@ -53,4 +66,6 @@ def emit_state_message(self) -> None:
         pass
 
     def ensure_at_least_one_state_emitted(self) -> None:
-        pass
+        self._connector_state_manager.update_state_for_stream(self._stream_name, self._stream_namespace, self.state)
+        state_message = self._connector_state_manager.create_state_message(self._stream_name, self._stream_namespace)
+        self._message_repository.emit_message(state_message)

airbyte-cdk/python/airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -21,7 +21,7 @@
     StreamAvailable,
     StreamUnavailable,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, FinalStateCursor
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
 from airbyte_cdk.sources.streams.concurrent.helpers import get_cursor_field_from_stream, get_primary_key_from_stream
@@ -77,7 +77,7 @@ def create_from_stream(
                 partition_generator=StreamPartitionGenerator(
                     stream,
                     message_repository,
-                    SyncMode.full_refresh if isinstance(cursor, NoopCursor) else SyncMode.incremental,
+                    SyncMode.full_refresh if isinstance(cursor, FinalStateCursor) else SyncMode.incremental,
                     [cursor_field] if cursor_field is not None else None,
                     state,
                     cursor,

airbyte-cdk/python/airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -8,6 +8,7 @@
 
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams import FULL_REFRESH_SENTINEL_STATE_KEY
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import AbstractStreamStateConverter
@@ -65,10 +66,27 @@ def ensure_at_least_one_state_emitted(self) -> None:
         raise NotImplementedError()
 
 
-class NoopCursor(Cursor):
+class FinalStateCursor(Cursor):
+    """Cursor that is used to guarantee at least one state message is emitted for a concurrent stream."""
+
+    def __init__(
+        self,
+        stream_name: str,
+        stream_namespace: Optional[str],
+        message_repository: MessageRepository,
+    ) -> None:
+        self._stream_name = stream_name
+        self._stream_namespace = stream_namespace
+        self._message_repository = message_repository
+        # Normally the connector state manager operates at the source-level. However, we only need it to write the sentinel
+        # state message rather than manage overall source state. This is also only temporary as we move to the resumable
+        # full refresh world where every stream uses a FileBasedConcurrentCursor with incremental state.
+        self._connector_state_manager = ConnectorStateManager(stream_instance_map={})
+        self._has_closed_at_least_one_slice = False
+
     @property
     def state(self) -> MutableMapping[str, Any]:
-        return {}
+        return {FULL_REFRESH_SENTINEL_STATE_KEY: True}
 
     def observe(self, record: Record) -> None:
         pass
@@ -77,7 +95,13 @@ def close_partition(self, partition: Partition) -> None:
         pass
 
     def ensure_at_least_one_state_emitted(self) -> None:
-        pass
+        """
+        Used primarily for full refresh syncs that do not have a valid cursor value to emit at the end of a sync
+        """
+
+        self._connector_state_manager.update_state_for_stream(self._stream_name, self._stream_namespace, self.state)
+        state_message = self._connector_state_manager.create_state_message(self._stream_name, self._stream_namespace)
+        self._message_repository.emit_message(state_message)
 
 
 class ConcurrentCursor(Cursor):

airbyte-cdk/python/airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -9,7 +9,7 @@
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import AbstractAvailabilityStrategy, StreamAvailability
-from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 
@@ -24,7 +24,7 @@ def __init__(
         primary_key: List[str],
         cursor_field: Optional[str],
         logger: Logger,
-        cursor: Optional[Cursor],
+        cursor: Cursor,
         namespace: Optional[str] = None,
     ) -> None:
         self._stream_partition_generator = partition_generator
@@ -34,7 +34,7 @@ def __init__(
         self._primary_key = primary_key
         self._cursor_field = cursor_field
         self._logger = logger
-        self._cursor = cursor or NoopCursor()
+        self._cursor = cursor
         self._namespace = namespace
 
     def generate_partitions(self) -> Iterable[Partition]:
@@ -44,6 +44,10 @@ def generate_partitions(self) -> Iterable[Partition]:
     def name(self) -> str:
         return self._name
 
+    @property
+    def namespace(self) -> Optional[str]:
+        return self._namespace
+
     def check_availability(self) -> StreamAvailability:
         return self._availability_strategy.check_availability(self._logger)
 

airbyte-cdk/python/unit_tests/sources/concurrent_source/test_concurrent_source_adapter.py

@@ -20,7 +20,7 @@
 from airbyte_cdk.sources.message import InMemoryMessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
-from airbyte_cdk.sources.streams.concurrent.cursor import NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
 
 
 class _MockSource(ConcurrentSourceAdapter):
@@ -36,7 +36,7 @@ def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) ->
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         return [
-            StreamFacade.create_from_stream(s, self, self._logger, None, NoopCursor()) if is_concurrent else s
+            StreamFacade.create_from_stream(s, self, self._logger, None, FinalStateCursor(stream_name=s.name, stream_namespace=s.namespace, message_repository=InMemoryMessageRepository())) if is_concurrent else s
             for s, is_concurrent in self._streams_to_is_concurrent.items()
         ]
 

airbyte-cdk/python/unit_tests/sources/file_based/stream/concurrent/test_adapters.py

@@ -23,7 +23,7 @@
     FileBasedStreamPartition,
     FileBasedStreamPartitionGenerator,
 )
-from airbyte_cdk.sources.file_based.stream.concurrent.cursor import FileBasedNoopCursor
+from airbyte_cdk.sources.file_based.stream.concurrent.cursor import FileBasedFinalStateCursor
 from airbyte_cdk.sources.message import InMemoryMessageRepository
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
@@ -36,7 +36,7 @@
 _ANY_STATE = {"state_key": "state_value"}
 _ANY_CURSOR_FIELD = ["a", "cursor", "key"]
 _STREAM_NAME = "stream"
-_ANY_CURSOR = Mock(spec=FileBasedNoopCursor)
+_ANY_CURSOR = Mock(spec=FileBasedFinalStateCursor)
 
 
 @pytest.mark.parametrize(
@@ -165,7 +165,7 @@ def setUp(self):
             supported_sync_modes=[SyncMode.full_refresh],
         )
         self._legacy_stream = DefaultFileBasedStream(
-            cursor=FileBasedNoopCursor(MagicMock()),
+            cursor=FileBasedFinalStateCursor(stream_config=MagicMock(), stream_namespace=None, message_repository=Mock()),
             config=FileBasedStreamConfig(name="stream", format=CsvFormat()),
             catalog_schema={},
             stream_reader=MagicMock(),

airbyte-cdk/python/unit_tests/sources/streams/concurrent/scenarios/stream_facade_builder.py

@@ -21,7 +21,7 @@
 from airbyte_cdk.sources.source import TState
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
-from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField, NoopCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField, FinalStateCursor
 from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import EpochValueConcurrentStreamStateConverter
 from airbyte_protocol.models import ConfiguredAirbyteStream
 from unit_tests.sources.file_based.scenarios.scenario_builder import SourceBuilder
@@ -83,7 +83,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
                     None,
                 )
                 if self._cursor_field
-                else NoopCursor(),
+                else FinalStateCursor(stream_name=stream.name, stream_namespace=stream.namespace, message_repository=self.message_repository),
             )
             for stream, state in zip(self._streams, stream_states)
         ]