airbyte-cdk 0.54.0__py3-none-any.whl → 0.55.0__py3-none-any.whl

unit_tests/sources/streams/concurrent/{test_thread_based_concurrent_stream.py → test_default_stream.py}

@@ -1,49 +1,33 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import unittest
-from unittest.mock import Mock, call
+from unittest.mock import Mock
 
-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
-
-_MAX_CONCURRENT_TASKS = 2
+from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 
 
 class ThreadBasedConcurrentStreamTest(unittest.TestCase):
     def setUp(self):
         self._partition_generator = Mock()
-        self._max_workers = 1
         self._name = "name"
         self._json_schema = {}
         self._availability_strategy = Mock()
         self._primary_key = []
         self._cursor_field = None
-        self._slice_logger = Mock()
         self._logger = Mock()
-        self._message_repository = Mock()
         self._cursor = Mock(spec=Cursor)
-        self._stream = ThreadBasedConcurrentStream(
+        self._stream = DefaultStream(
             self._partition_generator,
-            self._max_workers,
             self._name,
             self._json_schema,
             self._availability_strategy,
             self._primary_key,
             self._cursor_field,
-            self._slice_logger,
             self._logger,
-            self._message_repository,
-            1,
-            _MAX_CONCURRENT_TASKS,
-            0,
-            cursor=self._cursor,
         )
 
     def test_get_json_schema(self):
@@ -65,13 +49,6 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         with self.assertRaises(Exception):
             self._stream._check_for_errors(futures)
 
-    def test_check_for_error_raises_no_exception_if_all_futures_succeeded(self):
-        futures = [Mock() for _ in range(3)]
-        for f in futures:
-            f.exception.return_value = None
-
-        self._stream._check_for_errors(futures)
-
     def test_check_for_error_raises_an_exception_if_any_of_the_futures_raised_an_exception(self):
         futures = [Mock() for _ in range(3)]
         for f in futures:
@@ -81,96 +58,6 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         with self.assertRaises(Exception):
             self._stream._check_for_errors(futures)
 
-    def test_read_observe_records_and_close_partition(self):
-        partition = Mock(spec=Partition)
-        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]
-        actual_records = list(self._stream.read())
-
-        assert expected_records == actual_records
-
-        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)
-        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]
-        actual_records = list(self._stream.read())
-
-        assert expected_records == actual_records
-
-        self._message_repository.emit_message.assert_not_called()
-
-    def test_read_log_slice_message(self):
-        partition = Mock(spec=Partition)
-        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 = True
-        slice_log_message = Mock()
-        self._slice_logger.create_slice_log_message.return_value = slice_log_message
-
-        self._partition_generator.generate.return_value = [partition]
-        list(self._stream.read())
-
-        self._message_repository.emit_message.assert_called_once_with(slice_log_message)
-
-    def test_wait_while_task_queue_is_full(self):
-        f1 = Mock()
-        f2 = Mock()
-
-        # Verify that the done() method will be called until only one future is still running
-        f1.done.side_effect = [False, False]
-        f1.exception.return_value = None
-        f2.done.side_effect = [False, True]
-        f2.exception.return_value = None
-        futures = [f1, f2]
-        self._stream._wait_while_too_many_pending_futures(futures)
-
-        f1.done.assert_has_calls([call(), call()])
-        f2.done.assert_has_calls([call(), call()])
-
-    def test_given_exception_then_fail_immediately(self):
-        f1 = Mock()
-        f2 = Mock()
-
-        # Verify that the done() method will be called until only one future is still running
-        f1.done.return_value = True
-        f1.exception.return_value = None
-        f2.done.return_value = True
-        f2.exception.return_value = ValueError("ERROR")
-        futures = [f1, f2]
-
-        with pytest.raises(RuntimeError):
-            self._stream._wait_while_too_many_pending_futures(futures)
-
-    def test_given_removing_multiple_elements_when_pruning_then_fail_immediately(self):
-        # Verify that the done() method will be called until only one future is still running
-        futures = []
-        for _ in range(_MAX_CONCURRENT_TASKS + 1):
-            future = Mock()
-            future.done.return_value = True
-            future.exception.return_value = None
-            futures.append(future)
-
-        pending_future = Mock()
-        pending_future.done.return_value = False
-        pending_future.exception.return_value = None
-        futures.append(pending_future)
-
-        self._stream._wait_while_too_many_pending_futures(futures)
-
-        assert futures == [pending_future]
-
     def test_as_airbyte_stream(self):
         expected_airbyte_stream = AirbyteStream(
             name=self._name,
@@ -193,20 +80,14 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
                 "id_b": {"type": ["null", "string"]},
             },
         }
-        stream = ThreadBasedConcurrentStream(
+        stream = DefaultStream(
             self._partition_generator,
-            self._max_workers,
             self._name,
             json_schema,
             self._availability_strategy,
             ["id"],
             self._cursor_field,
-            self._slice_logger,
             self._logger,
-            self._message_repository,
-            1,
-            2,
-            0,
         )
 
         expected_airbyte_stream = AirbyteStream(
@@ -230,20 +111,14 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
                 "id_b": {"type": ["null", "string"]},
             },
         }
-        stream = ThreadBasedConcurrentStream(
+        stream = DefaultStream(
             self._partition_generator,
-            self._max_workers,
             self._name,
             json_schema,
             self._availability_strategy,
             ["id_a", "id_b"],
             self._cursor_field,
-            self._slice_logger,
             self._logger,
-            self._message_repository,
-            1,
-            2,
-            0,
         )
 
         expected_airbyte_stream = AirbyteStream(
@@ -267,20 +142,14 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
                 "date": {"type": ["null", "string"]},
             },
         }
-        stream = ThreadBasedConcurrentStream(
+        stream = DefaultStream(
             self._partition_generator,
-            self._max_workers,
             self._name,
             json_schema,
             self._availability_strategy,
             self._primary_key,
             "date",
-            self._slice_logger,
             self._logger,
-            self._message_repository,
-            1,
-            2,
-            0,
         )
 
         expected_airbyte_stream = AirbyteStream(
@@ -297,20 +166,14 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         assert expected_airbyte_stream == airbyte_stream
 
     def test_as_airbyte_stream_with_namespace(self):
-        stream = ThreadBasedConcurrentStream(
+        stream = DefaultStream(
             self._partition_generator,
-            self._max_workers,
             self._name,
             self._json_schema,
             self._availability_strategy,
             self._primary_key,
             self._cursor_field,
-            self._slice_logger,
             self._logger,
-            self._message_repository,
-            1,
-            2,
-            0,
             namespace="test",
         )
         expected_airbyte_stream = AirbyteStream(

unit_tests/sources/streams/concurrent/test_partition_reader.py

@@ -16,8 +16,8 @@ def test_partition_reader():
 
     stream_partition = Mock()
     records = [
-        Record({"id": 1, "name": "Jack"}),
-        Record({"id": 2, "name": "John"}),
+        Record({"id": 1, "name": "Jack"}, "stream"),
+        Record({"id": 2, "name": "John"}, "stream"),
     ]
     stream_partition.read.return_value = iter(records)
 

unit_tests/sources/streams/concurrent/test_thread_pool_manager.py

@@ -0,0 +1,98 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+from concurrent.futures import Future, ThreadPoolExecutor
+from unittest import TestCase
+from unittest.mock import Mock, patch
+
+from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+
+_SLEEP_TIME = 2
+
+
+class ThreadPoolManagerTest(TestCase):
+    def setUp(self):
+        self._threadpool = Mock(spec=ThreadPoolExecutor)
+        self._thread_pool_manager = ThreadPoolManager(self._threadpool, Mock(), max_concurrent_tasks=1, sleep_time=_SLEEP_TIME)
+        self._fn = lambda x: x
+        self._arg = "arg"
+
+    def test_submit_calls_underlying_thread_pool(self):
+        self._thread_pool_manager.submit(self._fn, self._arg)
+        self._threadpool.submit.assert_called_with(self._fn, self._arg)
+
+        assert len(self._thread_pool_manager._futures) == 1
+
+    def test_submit_too_many_concurrent_tasks(self):
+        future = Mock(spec=Future)
+        future.exception.return_value = None
+        future.done.side_effect = [False, True]
+
+        with patch("time.sleep") as sleep_mock:
+            self._thread_pool_manager._futures = [future]
+            self._thread_pool_manager.submit(self._fn, self._arg)
+            self._threadpool.submit.assert_called_with(self._fn, self._arg)
+            sleep_mock.assert_called_with(_SLEEP_TIME)
+
+            assert len(self._thread_pool_manager._futures) == 1
+
+    def test_submit_task_previous_task_failed(self):
+        future = Mock(spec=Future)
+        future.exception.return_value = RuntimeError
+        future.done.side_effect = [False, True]
+
+        self._thread_pool_manager._futures = [future]
+
+        with self.assertRaises(RuntimeError):
+            self._thread_pool_manager.submit(self._fn, self._arg)
+
+    def test_shutdown(self):
+        self._thread_pool_manager.shutdown()
+        self._threadpool.shutdown.assert_called_with(wait=False, cancel_futures=True)
+
+    def test_is_done_is_false_if_not_all_futures_are_done(self):
+        future = Mock(spec=Future)
+        future.done.return_value = False
+
+        self._thread_pool_manager._futures = [future]
+
+        assert not self._thread_pool_manager.is_done()
+
+    def test_is_done_is_true_if_all_futures_are_done(self):
+        future = Mock(spec=Future)
+        future.done.return_value = True
+
+        self._thread_pool_manager._futures = [future]
+
+        assert self._thread_pool_manager.is_done()
+
+    def test_threadpool_shutdown_if_errors(self):
+        future = Mock(spec=Future)
+        future.exception.return_value = RuntimeError
+
+        self._thread_pool_manager._futures = [future]
+
+        with self.assertRaises(RuntimeError):
+            self._thread_pool_manager.check_for_errors_and_shutdown()
+        self._threadpool.shutdown.assert_called_with(wait=False, cancel_futures=True)
+
+    def test_check_for_errors_and_shutdown_raises_error_if_futures_are_not_done(self):
+        future = Mock(spec=Future)
+        future.exception.return_value = None
+        future.done.return_value = False
+
+        self._thread_pool_manager._futures = [future]
+
+        with self.assertRaises(RuntimeError):
+            self._thread_pool_manager.check_for_errors_and_shutdown()
+        self._threadpool.shutdown.assert_called_with(wait=False, cancel_futures=True)
+
+    def test_check_for_errors_and_shutdown_does_not_raise_error_if_futures_are_done(self):
+        future = Mock(spec=Future)
+        future.exception.return_value = None
+        future.done.return_value = True
+
+        self._thread_pool_manager._futures = [future]
+
+        self._thread_pool_manager.check_for_errors_and_shutdown()
+        self._threadpool.shutdown.assert_called_with(wait=False, cancel_futures=True)

unit_tests/sources/streams/test_stream_read.py

@@ -1,7 +1,6 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import logging
 from typing import Any, Iterable, List, Mapping, Optional, Union
 from unittest.mock import Mock
@@ -59,7 +58,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, _NO_STATE, NoopCursor())
+    stream = StreamFacade.create_from_stream(stream, source, logger, _NO_STATE, NoopCursor())
     stream.logger.setLevel(logger.level)
     return stream
 

unit_tests/sources/test_concurrent_source.py

@@ -0,0 +1,105 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import concurrent
+import logging
+from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Tuple
+from unittest.mock import Mock
+
+from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
+from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
+from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability, StreamAvailable, StreamUnavailable
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_protocol.models import AirbyteStream
+
+logger = logging.getLogger("airbyte")
+
+
+class _MockSource(ConcurrentSource):
+    def __init__(
+        self,
+        check_lambda: Callable[[], Tuple[bool, Optional[Any]]] = None,
+        per_stream: bool = True,
+        message_repository: MessageRepository = InMemoryMessageRepository(),
+        threadpool: ThreadPoolManager = ThreadPoolManager(
+            concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="workerpool"), logger
+        ),
+        exception_on_missing_stream: bool = True,
+    ):
+        super().__init__(threadpool, Mock(), Mock(), message_repository)
+        self.check_lambda = check_lambda
+        self.per_stream = per_stream
+        self.exception_on_missing_stream = exception_on_missing_stream
+        self._message_repository = message_repository
+
+
+MESSAGE_FROM_REPOSITORY = Mock()
+
+
+class _MockStream(AbstractStream):
+    def __init__(self, name: str, available: bool = True, json_schema: Dict[str, Any] = {}):
+        self._name = name
+        self._available = available
+        self._json_schema = json_schema
+
+    def generate_partitions(self) -> Iterable[Partition]:
+        yield _MockPartition(self._name)
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def cursor_field(self) -> Optional[str]:
+        raise NotImplementedError
+
+    def check_availability(self) -> StreamAvailability:
+        if self._available:
+            return StreamAvailable()
+        else:
+            return StreamUnavailable("stream is unavailable")
+
+    def get_json_schema(self) -> Mapping[str, Any]:
+        return self._json_schema
+
+    def as_airbyte_stream(self) -> AirbyteStream:
+        return AirbyteStream(name=self.name, json_schema=self.get_json_schema(), supported_sync_modes=[SyncMode.full_refresh])
+
+    def log_stream_sync_configuration(self) -> None:
+        raise NotImplementedError
+
+
+class _MockPartition(Partition):
+    def __init__(self, name: str):
+        self._name = name
+        self._closed = False
+
+    def read(self) -> Iterable[Record]:
+        yield from [Record({"key": "value"}, self._name)]
+
+    def to_slice(self) -> Optional[Mapping[str, Any]]:
+        return {}
+
+    def stream_name(self) -> str:
+        return self._name
+
+    def close(self) -> None:
+        self._closed = True
+
+    def is_closed(self) -> bool:
+        return self._closed
+
+    def __hash__(self) -> int:
+        return hash(self._name)
+
+
+def test_concurrent_source_reading_from_no_streams():
+    stream = _MockStream("my_stream", False, {})
+    source = _MockSource()
+    messages = []
+    for m in source.read([stream]):
+        messages.append(m)