airbyte-cdk 0.61.2__py3-none-any.whl → 0.62.1__py3-none-any.whl

unit_tests/sources/file_based/scenarios/incremental_scenarios.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from airbyte_cdk.sources.file_based.stream.cursor import DefaultFileBasedCursor
 from airbyte_cdk.test.state_builder import StateBuilder
 from unit_tests.sources.file_based.helpers import LowHistoryLimitCursor
 from unit_tests.sources.file_based.scenarios.file_based_source_builder import FileBasedSourceBuilder
@@ -37,6 +38,7 @@ single_csv_input_state_is_earlier_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_incremental_scenario_config(
         IncrementalScenarioConfig(
@@ -130,6 +132,7 @@ single_csv_file_is_skipped_if_same_modified_at_as_in_history = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_incremental_scenario_config(
         IncrementalScenarioConfig(
@@ -205,6 +208,7 @@ single_csv_file_is_synced_if_modified_at_is_more_recent_than_in_history = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_incremental_scenario_config(
         IncrementalScenarioConfig(
@@ -298,6 +302,7 @@ single_csv_no_input_state_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -397,6 +402,7 @@ multi_csv_same_timestamp_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -511,6 +517,7 @@ single_csv_input_state_is_later_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -615,6 +622,7 @@ multi_csv_different_timestamps_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -753,6 +761,7 @@ multi_csv_per_timestamp_scenario = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -913,6 +922,7 @@ multi_csv_skip_file_if_already_in_history = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {
@@ -1059,6 +1069,7 @@ multi_csv_include_missing_files_within_history_range = (
             }
         )
         .set_file_type("csv")
+        .set_cursor_cls(DefaultFileBasedCursor)
     )
     .set_expected_catalog(
         {

unit_tests/sources/file_based/scenarios/scenario_builder.py

@@ -6,7 +6,7 @@ from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any, Generic, List, Mapping, Optional, Set, Tuple, Type, TypeVar
 
-from airbyte_cdk.models import AirbyteAnalyticsTraceMessage, SyncMode
+from airbyte_cdk.models import AirbyteAnalyticsTraceMessage, AirbyteStateMessage, SyncMode
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.source import TState
 from airbyte_protocol.models import ConfiguredAirbyteCatalog
@@ -191,10 +191,14 @@ class TestScenarioBuilder(Generic[SourceType]):
     def build(self) -> "TestScenario[SourceType]":
         if self.source_builder is None:
             raise ValueError("source_builder is not set")
+        if self._incremental_scenario_config and self._incremental_scenario_config.input_state:
+            state = [AirbyteStateMessage.parse_obj(s) for s in self._incremental_scenario_config.input_state]
+        else:
+            state = None
         source = self.source_builder.build(
             self._configured_catalog(SyncMode.incremental if self._incremental_scenario_config else SyncMode.full_refresh),
             self._config,
-            self._incremental_scenario_config.input_state if self._incremental_scenario_config else None,
+            state,
         )
         return TestScenario(
             self._name,

unit_tests/sources/file_based/stream/concurrent/test_adapters.py

@@ -0,0 +1,365 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import logging
+import unittest
+from datetime import datetime
+from unittest.mock import MagicMock, Mock
+
+import pytest
+from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteStream, Level, SyncMode
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.sources.file_based.availability_strategy import DefaultFileBasedAvailabilityStrategy
+from airbyte_cdk.sources.file_based.config.csv_format import CsvFormat
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import FileBasedStreamConfig
+from airbyte_cdk.sources.file_based.discovery_policy import DefaultDiscoveryPolicy
+from airbyte_cdk.sources.file_based.exceptions import FileBasedErrorsCollector
+from airbyte_cdk.sources.file_based.file_types import default_parsers
+from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.file_based.schema_validation_policies import EmitRecordPolicy
+from airbyte_cdk.sources.file_based.stream import DefaultFileBasedStream
+from airbyte_cdk.sources.file_based.stream.concurrent.adapters import (
+    FileBasedStreamFacade,
+    FileBasedStreamPartition,
+    FileBasedStreamPartitionGenerator,
+)
+from airbyte_cdk.sources.file_based.stream.concurrent.cursor import FileBasedNoopCursor
+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
+from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.utils.slice_logger import SliceLogger
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
+from freezegun import freeze_time
+
+_ANY_SYNC_MODE = SyncMode.full_refresh
+_ANY_STATE = {"state_key": "state_value"}
+_ANY_CURSOR_FIELD = ["a", "cursor", "key"]
+_STREAM_NAME = "stream"
+_ANY_CURSOR = Mock(spec=FileBasedNoopCursor)
+
+
+@pytest.mark.parametrize(
+    "sync_mode",
+    [
+        pytest.param(SyncMode.full_refresh, id="test_full_refresh"),
+        pytest.param(SyncMode.incremental, id="test_incremental"),
+    ],
+)
+def test_file_based_stream_partition_generator(sync_mode):
+    stream = Mock()
+    message_repository = Mock()
+    stream_slices = [{"files": [RemoteFile(uri="1", last_modified=datetime.now())]},
+                     {"files": [RemoteFile(uri="2", last_modified=datetime.now())]}]
+    stream.stream_slices.return_value = stream_slices
+
+    partition_generator = FileBasedStreamPartitionGenerator(stream, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
+
+    partitions = list(partition_generator.generate())
+    slices = [partition.to_slice() for partition in partitions]
+    assert slices == stream_slices
+    stream.stream_slices.assert_called_once_with(sync_mode=_ANY_SYNC_MODE, cursor_field=_ANY_CURSOR_FIELD, stream_state=_ANY_STATE)
+
+
+@pytest.mark.parametrize(
+    "transformer, expected_records",
+    [
+        pytest.param(
+            TypeTransformer(TransformConfig.NoTransform),
+            [Record({"data": "1"}, _STREAM_NAME), Record({"data": "2"}, _STREAM_NAME)],
+            id="test_no_transform",
+        ),
+        pytest.param(
+            TypeTransformer(TransformConfig.DefaultSchemaNormalization),
+            [Record({"data": 1}, _STREAM_NAME), Record({"data": 2}, _STREAM_NAME)],
+            id="test_default_transform",
+        ),
+    ],
+)
+def test_file_based_stream_partition(transformer, expected_records):
+    stream = Mock()
+    stream.name = _STREAM_NAME
+    stream.get_json_schema.return_value = {"type": "object", "properties": {"data": {"type": ["integer"]}}}
+    stream.transformer = transformer
+    message_repository = InMemoryMessageRepository()
+    _slice = None
+    sync_mode = SyncMode.full_refresh
+    cursor_field = None
+    state = None
+    partition = FileBasedStreamPartition(stream, _slice, message_repository, sync_mode, cursor_field, state, _ANY_CURSOR)
+
+    a_log_message = AirbyteMessage(
+        type=MessageType.LOG,
+        log=AirbyteLogMessage(
+            level=Level.INFO,
+            message='slice:{"partition": 1}',
+        ),
+    )
+
+    stream_data = [a_log_message, {"data": "1"}, {"data": "2"}]
+    stream.read_records.return_value = stream_data
+
+    records = list(partition.read())
+    messages = list(message_repository.consume_queue())
+
+    assert records == expected_records
+    assert messages == [a_log_message]
+
+
+@pytest.mark.parametrize(
+    "exception_type, expected_display_message",
+    [
+        pytest.param(Exception, None, id="test_exception_no_display_message"),
+        pytest.param(ExceptionWithDisplayMessage, "display_message", id="test_exception_no_display_message"),
+    ],
+)
+def test_file_based_stream_partition_raising_exception(exception_type, expected_display_message):
+    stream = Mock()
+    stream.get_error_display_message.return_value = expected_display_message
+
+    message_repository = InMemoryMessageRepository()
+    _slice = None
+
+    partition = FileBasedStreamPartition(stream, _slice, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
+
+    stream.read_records.side_effect = Exception()
+
+    with pytest.raises(exception_type) as e:
+        list(partition.read())
+        if isinstance(e, ExceptionWithDisplayMessage):
+            assert e.display_message == "display message"
+
+
+@freeze_time("2023-06-09T00:00:00Z")
+@pytest.mark.parametrize(
+    "_slice, expected_hash",
+    [
+        pytest.param({"files": [RemoteFile(uri="1", last_modified=datetime.strptime("2023-06-09T00:00:00Z", "%Y-%m-%dT%H:%M:%SZ"))]}, hash(("stream", "2023-06-09T00:00:00.000000Z_1")), id="test_hash_with_slice"),
+        pytest.param(None, hash("stream"), id="test_hash_no_slice"),
+    ],
+)
+def test_file_based_stream_partition_hash(_slice, expected_hash):
+    stream = Mock()
+    stream.name = "stream"
+    partition = FileBasedStreamPartition(stream, _slice, Mock(), _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
+
+    _hash = partition.__hash__()
+    assert _hash == expected_hash
+
+
+class StreamFacadeTest(unittest.TestCase):
+    def setUp(self):
+        self._abstract_stream = Mock()
+        self._abstract_stream.name = "stream"
+        self._abstract_stream.as_airbyte_stream.return_value = AirbyteStream(
+            name="stream",
+            json_schema={"type": "object"},
+            supported_sync_modes=[SyncMode.full_refresh],
+        )
+        self._legacy_stream = DefaultFileBasedStream(
+            cursor=FileBasedNoopCursor(MagicMock()),
+            config=FileBasedStreamConfig(name="stream", format=CsvFormat()),
+            catalog_schema={},
+            stream_reader=MagicMock(),
+            availability_strategy=DefaultFileBasedAvailabilityStrategy(MagicMock()),
+            discovery_policy=DefaultDiscoveryPolicy(),
+            parsers=default_parsers,
+            validation_policy=EmitRecordPolicy(),
+            errors_collector=FileBasedErrorsCollector(),
+        )
+        self._cursor = Mock(spec=Cursor)
+        self._logger = Mock()
+        self._slice_logger = Mock()
+        self._slice_logger.should_log_slice_message.return_value = False
+        self._facade = FileBasedStreamFacade(self._abstract_stream, self._legacy_stream, self._cursor, self._slice_logger, self._logger)
+        self._source = Mock()
+
+        self._stream = Mock()
+        self._stream.primary_key = "id"
+
+    def test_name_is_delegated_to_wrapped_stream(self):
+        assert self._facade.name == self._abstract_stream.name
+
+    def test_cursor_field_is_a_string(self):
+        self._abstract_stream.cursor_field = "cursor_field"
+        assert self._facade.cursor_field == "cursor_field"
+
+    def test_source_defined_cursor_is_true(self):
+        assert self._facade.source_defined_cursor
+
+    def test_json_schema_is_delegated_to_wrapped_stream(self):
+        json_schema = {"type": "object"}
+        self._abstract_stream.get_json_schema.return_value = json_schema
+        assert self._facade.get_json_schema() == json_schema
+        self._abstract_stream.get_json_schema.assert_called_once_with()
+
+    def test_given_cursor_is_noop_when_supports_incremental_then_return_legacy_stream_response(self):
+        assert (
+            FileBasedStreamFacade(
+                self._abstract_stream, self._legacy_stream, _ANY_CURSOR, Mock(spec=SliceLogger), Mock(spec=logging.Logger)
+            ).supports_incremental
+            == self._legacy_stream.supports_incremental
+        )
+
+    def test_given_cursor_is_not_noop_when_supports_incremental_then_return_true(self):
+        assert FileBasedStreamFacade(
+            self._abstract_stream, self._legacy_stream, Mock(spec=Cursor), Mock(spec=SliceLogger), Mock(spec=logging.Logger)
+        ).supports_incremental
+
+    def test_full_refresh(self):
+        expected_stream_data = [{"data": 1}, {"data": 2}]
+        records = [Record(data, "stream") for data in expected_stream_data]
+
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
+
+        actual_stream_data = list(self._facade.read_records(SyncMode.full_refresh, None, {}, None))
+
+        assert actual_stream_data == expected_stream_data
+
+    def test_read_records_full_refresh(self):
+        expected_stream_data = [{"data": 1}, {"data": 2}]
+        records = [Record(data, "stream") for data in expected_stream_data]
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
+
+        actual_stream_data = list(self._facade.read_full_refresh(None, None, None))
+
+        assert actual_stream_data == expected_stream_data
+
+    def test_read_records_incremental(self):
+        expected_stream_data = [{"data": 1}, {"data": 2}]
+        records = [Record(data, "stream") for data in expected_stream_data]
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
+
+        actual_stream_data = list(self._facade.read_incremental(None, None, None, None, None, None, None))
+
+        assert actual_stream_data == expected_stream_data
+
+    def test_create_from_stream_stream(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = "id"
+        stream.cursor_field = "cursor"
+
+        facade = FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+        assert facade.name == "stream"
+        assert facade.cursor_field == "cursor"
+        assert facade._abstract_stream._primary_key == ["id"]
+
+    def test_create_from_stream_stream_with_none_primary_key(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = None
+        stream.cursor_field = []
+
+        facade = FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+        assert facade._abstract_stream._primary_key == []
+
+    def test_create_from_stream_with_composite_primary_key(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = ["id", "name"]
+        stream.cursor_field = []
+
+        facade = FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+        assert facade._abstract_stream._primary_key == ["id", "name"]
+
+    def test_create_from_stream_with_empty_list_cursor(self):
+        stream = Mock()
+        stream.primary_key = "id"
+        stream.cursor_field = []
+
+        facade = FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+        assert facade.cursor_field == []
+
+    def test_create_from_stream_raises_exception_if_primary_key_is_nested(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = [["field", "id"]]
+
+        with self.assertRaises(ValueError):
+            FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+    def test_create_from_stream_raises_exception_if_primary_key_has_invalid_type(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = 123
+
+        with self.assertRaises(ValueError):
+            FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+    def test_create_from_stream_raises_exception_if_cursor_field_is_nested(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = "id"
+        stream.cursor_field = ["field", "cursor"]
+
+        with self.assertRaises(ValueError):
+            FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+    def test_create_from_stream_with_cursor_field_as_list(self):
+        stream = Mock()
+        stream.name = "stream"
+        stream.primary_key = "id"
+        stream.cursor_field = ["cursor"]
+
+        facade = FileBasedStreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+        assert facade.cursor_field == "cursor"
+
+    def test_create_from_stream_none_message_repository(self):
+        self._stream.name = "stream"
+        self._stream.primary_key = "id"
+        self._stream.cursor_field = "cursor"
+        self._source.message_repository = None
+
+        with self.assertRaises(ValueError):
+            FileBasedStreamFacade.create_from_stream(self._stream, self._source, self._logger, {}, self._cursor)
+
+    def test_get_error_display_message_no_display_message(self):
+        self._stream.get_error_display_message.return_value = "display_message"
+
+        facade = FileBasedStreamFacade.create_from_stream(self._stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+        expected_display_message = None
+        e = Exception()
+
+        display_message = facade.get_error_display_message(e)
+
+        assert expected_display_message == display_message
+
+    def test_get_error_display_message_with_display_message(self):
+        self._stream.get_error_display_message.return_value = "display_message"
+
+        facade = FileBasedStreamFacade.create_from_stream(self._stream, self._source, self._logger, _ANY_STATE, self._cursor)
+
+        expected_display_message = "display_message"
+        e = ExceptionWithDisplayMessage("display_message")
+
+        display_message = facade.get_error_display_message(e)
+
+        assert expected_display_message == display_message
+
+
+@pytest.mark.parametrize(
+    "exception, expected_display_message",
+    [
+        pytest.param(Exception("message"), None, id="test_no_display_message"),
+        pytest.param(ExceptionWithDisplayMessage("message"), "message", id="test_no_display_message"),
+    ],
+)
+def test_get_error_display_message(exception, expected_display_message):
+    stream = Mock()
+    legacy_stream = Mock()
+    cursor = Mock(spec=Cursor)
+    facade = FileBasedStreamFacade(stream, legacy_stream, cursor, Mock().Mock(), Mock())
+
+    display_message = facade.get_error_display_message(exception)
+
+    assert display_message == expected_display_message