airbyte-cdk 0.52.6__py3-none-any.whl → 0.52.8__py3-none-any.whl

unit_tests/sources/streams/concurrent/test_cursor.py

@@ -0,0 +1,130 @@
+from typing import Any, Mapping, Optional
+from unittest import TestCase
+from unittest.mock import Mock
+
+import pytest
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
+from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams.concurrent.cursor import Comparable, ConcurrentCursor, CursorField
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+
+_A_STREAM_NAME = "a stream name"
+_A_STREAM_NAMESPACE = "a stream namespace"
+_ANY_STATE = None
+_A_CURSOR_FIELD_KEY = "a_cursor_field_key"
+_NO_PARTITION_IDENTIFIER = None
+_NO_SLICE = None
+_NO_SLICE_BOUNDARIES = None
+_LOWER_SLICE_BOUNDARY_FIELD = "lower_boundary"
+_UPPER_SLICE_BOUNDARY_FIELD = "upper_boundary"
+_SLICE_BOUNDARY_FIELDS = (_LOWER_SLICE_BOUNDARY_FIELD, _UPPER_SLICE_BOUNDARY_FIELD)
+_A_VERY_HIGH_CURSOR_VALUE = 1000000000
+
+
+def _partition(_slice: Optional[Mapping[str, Any]]) -> Partition:
+    partition = Mock(spec=Partition)
+    partition.to_slice.return_value = _slice
+    return partition
+
+
+def _record(cursor_value: Comparable) -> Record:
+    return Record(data={_A_CURSOR_FIELD_KEY: cursor_value})
+
+
+class ConcurrentCursorTest(TestCase):
+    def setUp(self) -> None:
+        self._message_repository = Mock(spec=MessageRepository)
+        self._state_manager = Mock(spec=ConnectorStateManager)
+
+    def _cursor_with_slice_boundary_fields(self) -> ConcurrentCursor:
+        return ConcurrentCursor(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            _ANY_STATE,
+            self._message_repository,
+            self._state_manager,
+            CursorField(_A_CURSOR_FIELD_KEY),
+            _SLICE_BOUNDARY_FIELDS,
+        )
+
+    def _cursor_without_slice_boundary_fields(self) -> ConcurrentCursor:
+        return ConcurrentCursor(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            _ANY_STATE,
+            self._message_repository,
+            self._state_manager,
+            CursorField(_A_CURSOR_FIELD_KEY),
+            None,
+        )
+
+    def test_given_boundary_fields_when_close_partition_then_emit_state(self) -> None:
+        self._cursor_with_slice_boundary_fields().close_partition(
+            _partition(
+                {_LOWER_SLICE_BOUNDARY_FIELD: 12, _UPPER_SLICE_BOUNDARY_FIELD: 30},
+            )
+        )
+
+        self._message_repository.emit_message.assert_called_once_with(self._state_manager.create_state_message.return_value)
+        self._state_manager.update_state_for_stream.assert_called_once_with(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            {
+                "slices": [
+                    {
+                        "start": 12,
+                        "end": 30,
+                    },
+                ]
+            },
+        )
+
+    def test_given_boundary_fields_and_record_observed_when_close_partition_then_ignore_records(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        cursor.observe(_record(_A_VERY_HIGH_CURSOR_VALUE))
+
+        cursor.close_partition(_partition({_LOWER_SLICE_BOUNDARY_FIELD: 12, _UPPER_SLICE_BOUNDARY_FIELD: 30}))
+
+        assert self._state_manager.update_state_for_stream.call_args_list[0].args[2]["slices"][0]["end"] != _A_VERY_HIGH_CURSOR_VALUE
+
+    def test_given_no_boundary_fields_when_close_partition_then_emit_state(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.observe(_record(10))
+        cursor.close_partition(_partition(_NO_SLICE))
+
+        self._state_manager.update_state_for_stream.assert_called_once_with(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            {
+                "slices": [
+                    {
+                        "start": 0,
+                        "end": 10,
+                    },
+                ]
+            },
+        )
+
+    def test_given_no_boundary_fields_when_close_multiple_partitions_then_raise_exception(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.observe(_record(10))
+        cursor.close_partition(_partition(_NO_SLICE))
+
+        with pytest.raises(ValueError):
+            cursor.close_partition(_partition(_NO_SLICE))
+
+    def test_given_no_records_observed_when_close_partition_then_do_not_emit_state(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.close_partition(_partition(_NO_SLICE))
+        assert self._message_repository.emit_message.call_count == 0
+
+    def test_given_slice_boundaries_and_no_slice_when_close_partition_then_raise_error(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        with pytest.raises(KeyError):
+            cursor.close_partition(_partition(_NO_SLICE))
+
+    def test_given_slice_boundaries_not_matching_slice_when_close_partition_then_raise_error(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        with pytest.raises(KeyError):
+            cursor.close_partition(_partition({"not_matching_key": "value"}))

unit_tests/sources/streams/concurrent/test_thread_based_concurrent_stream.py

@@ -5,9 +5,9 @@
 import unittest
 from unittest.mock import Mock, call
 
-import pytest
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import STREAM_AVAILABLE
+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.record import Record
 from airbyte_cdk.sources.streams.concurrent.thread_based_concurrent_stream import ThreadBasedConcurrentStream
@@ -25,6 +25,7 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         self._slice_logger = Mock()
         self._logger = Mock()
         self._message_repository = Mock()
+        self._cursor = Mock(spec=Cursor)
         self._stream = ThreadBasedConcurrentStream(
             self._partition_generator,
             self._max_workers,
@@ -39,6 +40,7 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
             1,
             2,
             0,
+            cursor=self._cursor,
         )
 
     def test_get_json_schema(self):
@@ -76,17 +78,20 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         with self.assertRaises(Exception):
             self._stream._check_for_errors(futures)
 
-    def test_read_raises_an_exception_if_a_partition_raises_an_exception(self):
+    def test_read_observe_records_and_close_partition(self):
         partition = Mock(spec=Partition)
-        partition.read.side_effect = RuntimeError("error")
+        expected_records = [Record({"id": 1}), Record({"id": "2"})]
+        partition.read.return_value = expected_records
+        partition.to_slice.return_value = {"slice": "slice"}
+        self._slice_logger.should_log_slice_message.return_value = False
+
         self._partition_generator.generate.return_value = [partition]
-        with pytest.raises(RuntimeError):
-            list(self._stream.read())
+        actual_records = list(self._stream.read())
+
+        assert expected_records == actual_records
 
-    def test_read_raises_an_exception_if_partition_generator_raises_an_exception(self):
-        self._partition_generator.generate.side_effect = RuntimeError("error")
-        with pytest.raises(RuntimeError):
-            list(self._stream.read())
+        self._cursor.observe.has_calls([call(record) for record in expected_records])
+        self._cursor.close_partition.assert_called_once_with(partition)
 
     def test_read_no_slice_message(self):
         partition = Mock(spec=Partition)
@@ -218,7 +223,6 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         assert expected_airbyte_stream == airbyte_stream
 
     def test_as_airbyte_stream_with_a_cursor(self):
-
         json_schema = {
             "type": "object",
             "properties": {

unit_tests/sources/streams/test_stream_read.py

@@ -12,6 +12,7 @@ from airbyte_cdk.models import Type as MessageType
 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.core import StreamData
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
@@ -19,6 +20,7 @@ from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
 _A_CURSOR_FIELD = ["NESTED", "CURSOR"]
 _DEFAULT_INTERNAL_CONFIG = InternalConfig()
 _STREAM_NAME = "STREAM"
+_NO_STATE = None
 
 
 class _MockStream(Stream):
@@ -57,7 +59,7 @@ def _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message
     source = Mock()
     source._slice_logger = slice_logger
     source.message_repository = message_repository
-    stream = StreamFacade.create_from_stream(stream, source, logger, 1)
+    stream = StreamFacade.create_from_stream(stream, source, logger, 1, _NO_STATE, NoopCursor())
     stream.logger.setLevel(logger.level)
     return stream
 

unit_tests/sources/test_abstract_source.py

@@ -996,10 +996,11 @@ class TestIncrementalRead:
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 # stream 1 slice 2
                 _as_record("s1", stream_output[0]),
-                _as_record("s1", stream_output[1]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_record("s1", stream_output[1]),
                 _as_record("s1", stream_output[2]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
                 # stream 2 slice 1
                 _as_stream_status("s2", AirbyteStreamStatus.STARTED),
@@ -1011,17 +1012,18 @@ class TestIncrementalRead:
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 # stream 2 slice 2
                 _as_record("s2", stream_output[0]),
-                _as_record("s2", stream_output[1]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_record("s2", stream_output[1]),
                 _as_record("s2", stream_output[2]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 _as_stream_status("s2", AirbyteStreamStatus.COMPLETE),
             ]
         )
 
         messages = _fix_emitted_at(list(src.read(logger, {}, catalog, state=input_state)))
 
-        assert expected == messages
+        assert messages == expected
 
     @pytest.mark.parametrize(
         "per_stream_enabled",
@@ -1108,11 +1110,12 @@ class TestIncrementalRead:
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 # stream 1 slice 2
                 stream_data_to_airbyte_message("s1", stream_output[0]),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 stream_data_to_airbyte_message("s1", stream_output[1]),
                 stream_data_to_airbyte_message("s1", stream_output[2]),
-                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 stream_data_to_airbyte_message("s1", stream_output[3]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
                 # stream 2 slice 1
                 _as_stream_status("s2", AirbyteStreamStatus.STARTED),
@@ -1125,33 +1128,33 @@ class TestIncrementalRead:
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 # stream 2 slice 2
                 stream_data_to_airbyte_message("s2", stream_output[0]),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 stream_data_to_airbyte_message("s2", stream_output[1]),
                 stream_data_to_airbyte_message("s2", stream_output[2]),
-                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 stream_data_to_airbyte_message("s2", stream_output[3]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 _as_stream_status("s2", AirbyteStreamStatus.COMPLETE),
             ]
         )
 
         messages = _fix_emitted_at(list(src.read(logger, {}, catalog, state=input_state)))
 
-        assert expected == messages
+        assert messages == expected
 
 
 def test_checkpoint_state_from_stream_instance():
     teams_stream = MockStreamOverridesStateMethod()
     managers_stream = StreamNoStateMethod()
-    src = MockSource(streams=[teams_stream, managers_stream])
     state_manager = ConnectorStateManager({"teams": teams_stream, "managers": managers_stream}, [])
 
     # The stream_state passed to checkpoint_state() should be ignored since stream implements state function
     teams_stream.state = {"updated_at": "2022-09-11"}
-    actual_message = src._checkpoint_state(teams_stream, {"ignored": "state"}, state_manager)
+    actual_message = teams_stream._checkpoint_state({"ignored": "state"}, state_manager, True)
     assert actual_message == _as_state({"teams": {"updated_at": "2022-09-11"}}, "teams", {"updated_at": "2022-09-11"})
 
     # The stream_state passed to checkpoint_state() should be used since the stream does not implement state function
-    actual_message = src._checkpoint_state(managers_stream, {"updated": "expected_here"}, state_manager)
+    actual_message = managers_stream._checkpoint_state({"updated": "expected_here"}, state_manager, True)
     assert actual_message == _as_state(
         {"teams": {"updated_at": "2022-09-11"}, "managers": {"updated": "expected_here"}}, "managers", {"updated": "expected_here"}
     )