airbyte-cdk 0.67.0__py3-none-any.whl → 0.67.2__py3-none-any.whl

unit_tests/sources/streams/test_stream_read.py

@@ -1,17 +1,34 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+
 import logging
-from typing import Any, Iterable, List, Mapping, Optional, Union
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 from unittest.mock import Mock
 
 import pytest
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level, SyncMode
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateType,
+    AirbyteStream,
+    AirbyteStreamState,
+    ConfiguredAirbyteStream,
+    DestinationSyncMode,
+    Level,
+    StreamDescriptor,
+    SyncMode,
+)
 from airbyte_cdk.models import Type as MessageType
-from airbyte_cdk.sources.message import InMemoryMessageRepository
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
+from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
 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 Cursor, NoopCursor
+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.core import StreamData
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
@@ -49,20 +66,66 @@ class _MockStream(Stream):
         return {}
 
 
+class MockConcurrentCursor(Cursor):
+    _state: MutableMapping[str, Any]
+    _message_repository: MessageRepository
+
+    def __init__(self, message_repository: MessageRepository):
+        self._message_repository = message_repository
+        self._state = {}
+
+    @property
+    def state(self) -> MutableMapping[str, Any]:
+        return self._state
+
+    def observe(self, record: Record) -> None:
+        partition = str(record.data.get("partition"))
+        timestamp = record.data.get("created_at")
+        self._state[partition] = {"created_at": timestamp}
+
+    def close_partition(self, partition: Partition) -> None:
+        self._message_repository.emit_message(
+            AirbyteMessage(
+                type=MessageType.STATE,
+                state=AirbyteStateMessage(
+                    type=AirbyteStateType.STREAM,
+                    stream=AirbyteStreamState(
+                        stream_descriptor=StreamDescriptor(name='__mock_stream', namespace=None),
+                        stream_state=AirbyteStateBlob(**self._state),
+                    )
+                ),
+            )
+        )
+
+    def ensure_at_least_one_state_emitted(self) -> None:
+        pass
+
+
 def _stream(slice_to_partition_mapping, slice_logger, logger, message_repository):
     return _MockStream(slice_to_partition_mapping)
 
 
-def _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message_repository):
+def _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message_repository, cursor: Cursor = NoopCursor()):
     stream = _stream(slice_to_partition_mapping, slice_logger, logger, message_repository)
     source = Mock()
     source._slice_logger = slice_logger
     source.message_repository = message_repository
-    stream = StreamFacade.create_from_stream(stream, source, logger, _NO_STATE, NoopCursor())
+    stream = StreamFacade.create_from_stream(stream, source, logger, _NO_STATE, cursor)
     stream.logger.setLevel(logger.level)
     return stream
 
 
+def _incremental_stream(slice_to_partition_mapping, slice_logger, logger, message_repository, timestamp):
+    stream = _stream(slice_to_partition_mapping, slice_logger, logger, message_repository)
+    stream.state = {"created_at": timestamp}
+    return stream
+
+
+def _incremental_concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message_repository, cursor):
+    stream = _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message_repository, cursor)
+    return stream
+
+
 @pytest.mark.parametrize(
     "constructor",
     [
@@ -73,6 +136,8 @@ def _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message
 def test_full_refresh_read_a_single_slice_with_debug(constructor):
     # This test verifies that a concurrent stream adapted from a Stream behaves the same as the Stream object.
     # It is done by running the same test cases on both streams
+    configured_stream = ConfiguredAirbyteStream(stream=AirbyteStream(name="mock_stream", supported_sync_modes=[SyncMode.full_refresh], json_schema={}), sync_mode=SyncMode.full_refresh,destination_sync_mode=DestinationSyncMode.overwrite)
+    internal_config = InternalConfig()
     records = [
         {"id": 1, "partition": 1},
         {"id": 2, "partition": 1},
@@ -82,6 +147,7 @@ def test_full_refresh_read_a_single_slice_with_debug(constructor):
     logger = _mock_logger(True)
     message_repository = InMemoryMessageRepository(Level.DEBUG)
     stream = constructor(slice_to_partition, slice_logger, logger, message_repository)
+    state_manager = ConnectorStateManager(stream_instance_map={})
 
     expected_records = [
         AirbyteMessage(
@@ -94,7 +160,22 @@ def test_full_refresh_read_a_single_slice_with_debug(constructor):
         *records,
     ]
 
-    actual_records = _read(stream, logger, slice_logger, message_repository)
+    # Temporary check to only validate the final state message for synchronous sources since it has not been implemented for concurrent yet
+    if constructor == _stream:
+        expected_records.append(
+            AirbyteMessage(
+                type=MessageType.STATE,
+                state=AirbyteStateMessage(
+                    type=AirbyteStateType.STREAM,
+                    stream=AirbyteStreamState(
+                        stream_descriptor=StreamDescriptor(name='__mock_stream', namespace=None),
+                        stream_state=AirbyteStateBlob(__ab_full_refresh_state_message=True),
+                    )
+                ),
+            ),
+        )
+
+    actual_records = _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config)
 
     assert expected_records == actual_records
 
@@ -109,9 +190,12 @@ def test_full_refresh_read_a_single_slice_with_debug(constructor):
 def test_full_refresh_read_a_single_slice(constructor):
     # This test verifies that a concurrent stream adapted from a Stream behaves the same as the Stream object.
     # It is done by running the same test cases on both streams
+    configured_stream = ConfiguredAirbyteStream(stream=AirbyteStream(name="mock_stream", supported_sync_modes=[SyncMode.full_refresh], json_schema={}), sync_mode=SyncMode.full_refresh,destination_sync_mode=DestinationSyncMode.overwrite)
+    internal_config = InternalConfig()
     logger = _mock_logger()
     slice_logger = DebugSliceLogger()
     message_repository = InMemoryMessageRepository(Level.INFO)
+    state_manager = ConnectorStateManager(stream_instance_map={})
 
     records = [
         {"id": 1, "partition": 1},
@@ -122,7 +206,22 @@ def test_full_refresh_read_a_single_slice(constructor):
 
     expected_records = [*records]
 
-    actual_records = _read(stream, logger, slice_logger, message_repository)
+    # Temporary check to only validate the final state message for synchronous sources since it has not been implemented for concurrent yet
+    if constructor == _stream:
+        expected_records.append(
+            AirbyteMessage(
+                type=MessageType.STATE,
+                state=AirbyteStateMessage(
+                    type=AirbyteStateType.STREAM,
+                    stream=AirbyteStreamState(
+                        stream_descriptor=StreamDescriptor(name='__mock_stream', namespace=None),
+                        stream_state=AirbyteStateBlob(__ab_full_refresh_state_message=True),
+                    )
+                ),
+            ),
+        )
+
+    actual_records = _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config)
 
     assert expected_records == actual_records
 
@@ -137,9 +236,12 @@ def test_full_refresh_read_a_single_slice(constructor):
 def test_full_refresh_read_a_two_slices(constructor):
     # This test verifies that a concurrent stream adapted from a Stream behaves the same as the Stream object
     # It is done by running the same test cases on both streams
+    configured_stream = ConfiguredAirbyteStream(stream=AirbyteStream(name="mock_stream", supported_sync_modes=[SyncMode.full_refresh], json_schema={}), sync_mode=SyncMode.full_refresh,destination_sync_mode=DestinationSyncMode.overwrite)
+    internal_config = InternalConfig()
     logger = _mock_logger()
     slice_logger = DebugSliceLogger()
     message_repository = InMemoryMessageRepository(Level.INFO)
+    state_manager = ConnectorStateManager(stream_instance_map={})
 
     records_partition_1 = [
         {"id": 1, "partition": 1},
@@ -157,16 +259,111 @@ def test_full_refresh_read_a_two_slices(constructor):
         *records_partition_2,
     ]
 
-    actual_records = _read(stream, logger, slice_logger, message_repository)
+    # Temporary check to only validate the final state message for synchronous sources since it has not been implemented for concurrent yet
+    if constructor == _stream:
+        expected_records.append(
+            AirbyteMessage(
+                type=MessageType.STATE,
+                state=AirbyteStateMessage(
+                    type=AirbyteStateType.STREAM,
+                    stream=AirbyteStreamState(
+                        stream_descriptor=StreamDescriptor(name='__mock_stream', namespace=None),
+                        stream_state=AirbyteStateBlob(__ab_full_refresh_state_message=True),
+                    )
+                ),
+            ),
+        )
+
+    actual_records = _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config)
 
     for record in expected_records:
         assert record in actual_records
     assert len(expected_records) == len(actual_records)
 
 
-def _read(stream, logger, slice_logger, message_repository):
+def test_incremental_read_two_slices():
+    # This test verifies that a stream running in incremental mode emits state messages correctly
+    configured_stream = ConfiguredAirbyteStream(stream=AirbyteStream(name="mock_stream", supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental], json_schema={}), sync_mode=SyncMode.incremental,destination_sync_mode=DestinationSyncMode.overwrite)
+    internal_config = InternalConfig()
+    logger = _mock_logger()
+    slice_logger = DebugSliceLogger()
+    message_repository = InMemoryMessageRepository(Level.INFO)
+    state_manager = ConnectorStateManager(stream_instance_map={})
+    timestamp = "1708899427"
+
+    records_partition_1 = [
+        {"id": 1, "partition": 1},
+        {"id": 2, "partition": 1},
+    ]
+    records_partition_2 = [
+        {"id": 3, "partition": 2},
+        {"id": 4, "partition": 2},
+    ]
+    slice_to_partition = {1: records_partition_1, 2: records_partition_2}
+    stream = _incremental_stream(slice_to_partition, slice_logger, logger, message_repository, timestamp)
+
+    expected_records = [
+        *records_partition_1,
+        _create_state_message("__mock_stream", {"created_at": timestamp}),
+        *records_partition_2,
+        _create_state_message("__mock_stream", {"created_at": timestamp})
+    ]
+
+    actual_records = _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config)
+
+    for record in expected_records:
+        assert record in actual_records
+    assert len(expected_records) == len(actual_records)
+
+
+def test_concurrent_incremental_read_two_slices():
+    # This test verifies that an incremental concurrent stream manages state correctly for multiple slices syncing concurrently
+    configured_stream = ConfiguredAirbyteStream(stream=AirbyteStream(name="mock_stream", supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental], json_schema={}), sync_mode=SyncMode.incremental,destination_sync_mode=DestinationSyncMode.overwrite)
+    internal_config = InternalConfig()
+    logger = _mock_logger()
+    slice_logger = DebugSliceLogger()
+    message_repository = InMemoryMessageRepository(Level.INFO)
+    state_manager = ConnectorStateManager(stream_instance_map={})
+    slice_timestamp_1 = "1708850000"
+    slice_timestamp_2 = "1708950000"
+    cursor = MockConcurrentCursor(message_repository)
+
+    records_partition_1 = [
+        {"id": 1, "partition": 1, "created_at": "1708800000"},
+        {"id": 2, "partition": 1, "created_at": slice_timestamp_1},
+    ]
+    records_partition_2 = [
+        {"id": 3, "partition": 2, "created_at": "1708900000"},
+        {"id": 4, "partition": 2, "created_at": slice_timestamp_2},
+    ]
+    slice_to_partition = {1: records_partition_1, 2: records_partition_2}
+    stream = _incremental_concurrent_stream(slice_to_partition, slice_logger, logger, message_repository, cursor)
+
+    expected_records = [
+        *records_partition_1,
+        *records_partition_2,
+    ]
+
+    expected_state = _create_state_message("__mock_stream", {"1": {"created_at": slice_timestamp_1}, "2": {"created_at": slice_timestamp_2}})
+
+    actual_records = _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config)
+
+    for record in expected_records:
+        assert record in actual_records
+    assert len(expected_records) == len(actual_records)
+
+    # We don't have a real source that reads from the message_repository for state, so we read from the queue directly to verify
+    # the cursor observed records correctly and updated partition states
+    mock_partition = Mock()
+    cursor.close_partition(mock_partition)
+    actual_state = [state for state in message_repository.consume_queue()]
+    assert len(actual_state) == 1
+    assert actual_state[0] == expected_state
+
+
+def _read(stream, configured_stream, logger, slice_logger, message_repository, state_manager, internal_config):
     records = []
-    for record in stream.read_full_refresh(_A_CURSOR_FIELD, logger, slice_logger):
+    for record in stream.read(configured_stream, logger, slice_logger, {}, state_manager, internal_config):
         for message in message_repository.consume_queue():
             records.append(message)
         records.append(record)
@@ -192,3 +389,16 @@ def _mock_logger(enabled_for_debug=False):
     logger.isEnabledFor.return_value = enabled_for_debug
     logger.level = logging.DEBUG if enabled_for_debug else logging.INFO
     return logger
+
+
+def _create_state_message(stream: str, state: Mapping[str, Any]) -> AirbyteMessage:
+    return AirbyteMessage(
+        type=MessageType.STATE,
+        state=AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(name=stream, namespace=None),
+                stream_state=AirbyteStateBlob(**state),
+            )
+        ),
+    )