airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev1__py3-none-any.whl

unit_tests/sources/streams/concurrent/scenarios/stream_facade_builder.py

@@ -1,140 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-import concurrent
-import logging
-from typing import Any, List, Mapping, Optional, Tuple, Union
-
-from airbyte_cdk.models import (
-    AirbyteStateMessage,
-    AirbyteStream,
-    ConfiguredAirbyteCatalog,
-    ConnectorSpecification,
-    DestinationSyncMode,
-    SyncMode,
-)
-from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
-from airbyte_cdk.sources.concurrent_source.concurrent_source_adapter import ConcurrentSourceAdapter
-from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
-from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
-from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
-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, 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
-from unit_tests.sources.streams.concurrent.scenarios.thread_based_concurrent_stream_source_builder import NeverLogSliceLogger
-
-_CURSOR_FIELD = "cursor_field"
-_NO_STATE = None
-
-
-class StreamFacadeConcurrentConnectorStateConverter(EpochValueConcurrentStreamStateConverter):
-    pass
-
-
-class StreamFacadeSource(ConcurrentSourceAdapter):
-    def __init__(
-        self,
-        streams: List[Stream],
-        threadpool: concurrent.futures.ThreadPoolExecutor,
-        cursor_field: Optional[CursorField] = None,
-        cursor_boundaries: Optional[Tuple[str, str]] = None,
-        input_state: Optional[List[Mapping[str, Any]]] = _NO_STATE,
-    ):
-        self._message_repository = InMemoryMessageRepository()
-        threadpool_manager = ThreadPoolManager(threadpool, streams[0].logger)
-        concurrent_source = ConcurrentSource(threadpool_manager, streams[0].logger, NeverLogSliceLogger(), self._message_repository)
-        super().__init__(concurrent_source)
-        self._streams = streams
-        self._threadpool = threadpool_manager
-        self._cursor_field = cursor_field
-        self._cursor_boundaries = cursor_boundaries
-        self._state = [AirbyteStateMessage.parse_obj(s) for s in input_state] if input_state else None
-
-    def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
-        return True, None
-
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        state_manager = ConnectorStateManager(
-            stream_instance_map={s.name: AirbyteStream(name=s.name, namespace=None, json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental]) for s in self._streams},
-            state=self._state,
-        )  # The input values into the AirbyteStream are dummy values; the connector state manager only uses `name` and `namespace`
-
-        state_converter = StreamFacadeConcurrentConnectorStateConverter()
-        stream_states = [state_manager.get_stream_state(stream.name, stream.namespace) for stream in self._streams]
-        return [
-            StreamFacade.create_from_stream(
-                stream,
-                self,
-                stream.logger,
-                state,
-                ConcurrentCursor(
-                    stream.name,
-                    stream.namespace,
-                    state,
-                    self.message_repository,  # type: ignore  # for this source specifically, we always return `InMemoryMessageRepository`
-                    state_manager,
-                    state_converter,
-                    self._cursor_field,
-                    self._cursor_boundaries,
-                    None,
-                )
-                if self._cursor_field
-                else FinalStateCursor(stream_name=stream.name, stream_namespace=stream.namespace, message_repository=self.message_repository),
-            )
-            for stream, state in zip(self._streams, stream_states)
-        ]
-
-    @property
-    def message_repository(self) -> Union[None, MessageRepository]:
-        return self._message_repository
-
-    def spec(self, logger: logging.Logger) -> ConnectorSpecification:
-        return ConnectorSpecification(connectionSpecification={})
-
-    def read_catalog(self, catalog_path: str) -> ConfiguredAirbyteCatalog:
-        return ConfiguredAirbyteCatalog(
-            streams=[
-                ConfiguredAirbyteStream(
-                    stream=s.as_airbyte_stream(),
-                    sync_mode=SyncMode.full_refresh,
-                    destination_sync_mode=DestinationSyncMode.overwrite,
-                )
-                for s in self._streams
-            ]
-        )
-
-
-class StreamFacadeSourceBuilder(SourceBuilder[StreamFacadeSource]):
-    def __init__(self):
-        self._source = None
-        self._streams = []
-        self._max_workers = 1
-        self._cursor_field = None
-        self._cursor_boundaries = None
-        self._input_state = None
-        self._raw_input_state = None
-
-    def set_streams(self, streams: List[Stream]) -> "StreamFacadeSourceBuilder":
-        self._streams = streams
-        return self
-
-    def set_max_workers(self, max_workers: int) -> "StreamFacadeSourceBuilder":
-        self._max_workers = max_workers
-        return self
-
-    def set_incremental(self, cursor_field: CursorField, cursor_boundaries: Optional[Tuple[str, str]]) -> "StreamFacadeSourceBuilder":
-        self._cursor_field = cursor_field
-        self._cursor_boundaries = cursor_boundaries
-        return self
-
-    def set_input_state(self, state: List[Mapping[str, Any]]) -> "StreamFacadeSourceBuilder":
-        self._input_state = state
-        return self
-
-    def build(self, configured_catalog: Optional[Mapping[str, Any]], config: Optional[Mapping[str, Any]], state: Optional[TState]) -> StreamFacadeSource:
-        threadpool = concurrent.futures.ThreadPoolExecutor(max_workers=self._max_workers, thread_name_prefix="workerpool")
-        return StreamFacadeSource(self._streams, threadpool, self._cursor_field, self._cursor_boundaries, state)

unit_tests/sources/streams/concurrent/scenarios/stream_facade_scenarios.py

@@ -1,452 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
-from unit_tests.sources.file_based.scenarios.scenario_builder import IncrementalScenarioConfig, TestScenarioBuilder
-from unit_tests.sources.streams.concurrent.scenarios.stream_facade_builder import StreamFacadeSourceBuilder
-from unit_tests.sources.streams.concurrent.scenarios.utils import MockStream
-
-_stream1 = MockStream(
-    [
-        (None, [{"id": "1"}, {"id": "2"}]),
-    ],
-    "stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-)
-
-_stream_raising_exception = MockStream(
-    [
-        (None, [{"id": "1"}, ValueError("test exception")]),
-    ],
-    "stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-)
-
-_stream_with_primary_key = MockStream(
-    [
-        (None, [{"id": "1"}, {"id": "2"}]),
-    ],
-    "stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    primary_key="id",
-)
-
-_stream2 = MockStream(
-    [
-        (None, [{"id": "A"}, {"id": "B"}]),
-    ],
-    "stream2",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-)
-
-_stream_with_single_slice = MockStream(
-    [
-        ({"slice_key": "s1"}, [{"id": "1"}, {"id": "2"}]),
-    ],
-    "stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-)
-
-_stream_with_multiple_slices = MockStream(
-    [
-        ({"slice_key": "s1"}, [{"id": "1"}, {"id": "2"}]),
-        ({"slice_key": "s2"}, [{"id": "3"}, {"id": "4"}]),
-    ],
-    "stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-)
-
-test_stream_facade_single_stream = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_single_stream")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream1]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .set_expected_logs(
-        {
-            "read": [
-                {"level": "INFO", "message": "Starting syncing"},
-                {"level": "INFO", "message": "Marking stream stream1 as STARTED"},
-                {"level": "INFO", "message": "Syncing stream: stream1"},
-                {"level": "INFO", "message": "Marking stream stream1 as RUNNING"},
-                {"level": "INFO", "message": "Read 2 records from stream1 stream"},
-                {"level": "INFO", "message": "Marking stream stream1 as STOPPED"},
-                {"level": "INFO", "message": "Finished syncing stream1"},
-                {"level": "INFO", "message": "Finished syncing"},
-            ]
-        }
-    )
-    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
-    .build()
-)
-
-test_stream_facade_raises_exception = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_raises_exception")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream_raising_exception]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .set_expected_read_error(ValueError, "test exception")
-    .build()
-)
-
-test_stream_facade_single_stream_with_primary_key = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_stream_with_primary_key")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream1]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_stream_facade_multiple_streams = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_multiple_streams")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream1, _stream2]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-            {"data": {"id": "A"}, "stream": "stream2"},
-            {"data": {"id": "B"}, "stream": "stream2"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                },
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream2",
-                    "supported_sync_modes": ["full_refresh"],
-                },
-            ]
-        }
-    )
-    .build()
-)
-
-test_stream_facade_single_stream_with_single_slice = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_single_stream_with_single_slice")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream1]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_stream_facade_single_stream_with_multiple_slices = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_single_stream_with_multiple_slice")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream_with_multiple_slices]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-            {"data": {"id": "3"}, "stream": "stream1"},
-            {"data": {"id": "4"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_stream_facade_single_stream_with_multiple_slices_with_concurrency_level_two = (
-    TestScenarioBuilder()
-    .set_name("test_stream_facade_single_stream_with_multiple_slice_with_concurrency_level_two")
-    .set_config({})
-    .set_source_builder(StreamFacadeSourceBuilder().set_streams([_stream_with_multiple_slices]))
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-            {"data": {"id": "3"}, "stream": "stream1"},
-            {"data": {"id": "4"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-
-test_incremental_stream_with_slice_boundaries = (
-    TestScenarioBuilder()
-    .set_name("test_incremental_stream_with_slice_boundaries")
-    .set_config({})
-    .set_source_builder(
-        StreamFacadeSourceBuilder()
-        .set_streams(
-            [
-                MockStream(
-                    [
-                        ({"from": 0, "to": 1}, [{"id": "1", "cursor_field": 0}, {"id": "2", "cursor_field": 1}]),
-                        ({"from": 1, "to": 2}, [{"id": "3", "cursor_field": 2}, {"id": "4", "cursor_field": 3}]),
-                    ],
-                    "stream1",
-                    json_schema={
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                )
-            ]
-        )
-        .set_incremental(CursorField("cursor_field"), ("from", "to"))
-    )
-    .set_expected_records(
-        [
-            {"data": {"id": "1", "cursor_field": 0}, "stream": "stream1"},
-            {"data": {"id": "2", "cursor_field": 1}, "stream": "stream1"},
-            {"cursor_field": 1},
-            {"data": {"id": "3", "cursor_field": 2}, "stream": "stream1"},
-            {"data": {"id": "4", "cursor_field": 3}, "stream": "stream1"},
-            {"cursor_field": 2},
-            {"cursor_field": 2},  # see Cursor.ensure_at_least_one_state_emitted
-        ]
-    )
-    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
-    .set_incremental_scenario_config(
-        IncrementalScenarioConfig(
-            input_state=[],
-        )
-    )
-    .build()
-)
-
-
-_NO_SLICE_BOUNDARIES = None
-test_incremental_stream_without_slice_boundaries = (
-    TestScenarioBuilder()
-    .set_name("test_incremental_stream_without_slice_boundaries")
-    .set_config({})
-    .set_source_builder(
-        StreamFacadeSourceBuilder()
-        .set_streams(
-            [
-                MockStream(
-                    [
-                        (None, [{"id": "1", "cursor_field": 0}, {"id": "2", "cursor_field": 3}]),
-                    ],
-                    "stream1",
-                    json_schema={
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                )
-            ]
-        )
-        .set_incremental(CursorField("cursor_field"), _NO_SLICE_BOUNDARIES)
-    )
-    .set_expected_records(
-        [
-            {"data": {"id": "1", "cursor_field": 0}, "stream": "stream1"},
-            {"data": {"id": "2", "cursor_field": 3}, "stream": "stream1"},
-            {"cursor_field": 3},
-            {"cursor_field": 3},  # see Cursor.ensure_at_least_one_state_emitted
-        ]
-    )
-    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
-    .set_incremental_scenario_config(
-        IncrementalScenarioConfig(
-            input_state=[],
-        )
-    )
-    .build()
-)
-
-test_incremental_stream_with_many_slices_but_without_slice_boundaries = (
-    TestScenarioBuilder()
-    .set_name("test_incremental_stream_with_many_slices_byt_without_slice_boundaries")
-    .set_config({})
-    .set_source_builder(
-        StreamFacadeSourceBuilder()
-        .set_streams(
-            [
-                MockStream(
-                    [
-                        ({"parent_id": 1}, [{"id": "1", "cursor_field": 0}]),
-                        ({"parent_id": 309}, [{"id": "3", "cursor_field": 0}]),
-                    ],
-                    "stream1",
-                    json_schema={
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                )
-            ]
-        )
-        .set_incremental(CursorField("cursor_field"), _NO_SLICE_BOUNDARIES)
-    )
-    .set_expected_read_error(ValueError, "test exception")
-    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
-    .set_incremental_scenario_config(
-        IncrementalScenarioConfig(
-            input_state=[],
-        )
-    )
-    .build()
-)

unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py

@@ -1,76 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-from pathlib import PosixPath
-
-import pytest
-from _pytest.capture import CaptureFixture
-from freezegun import freeze_time
-from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenario
-from unit_tests.sources.file_based.test_scenarios import verify_discover, verify_read
-from unit_tests.sources.streams.concurrent.scenarios.incremental_scenarios import (
-    test_incremental_stream_with_slice_boundaries_no_input_state,
-    test_incremental_stream_with_slice_boundaries_with_concurrent_state,
-    test_incremental_stream_with_slice_boundaries_with_legacy_state,
-    test_incremental_stream_without_slice_boundaries_no_input_state,
-    test_incremental_stream_without_slice_boundaries_with_concurrent_state,
-    test_incremental_stream_without_slice_boundaries_with_legacy_state,
-)
-from unit_tests.sources.streams.concurrent.scenarios.stream_facade_scenarios import (
-    test_incremental_stream_with_many_slices_but_without_slice_boundaries,
-    test_incremental_stream_with_slice_boundaries,
-    test_incremental_stream_without_slice_boundaries,
-    test_stream_facade_multiple_streams,
-    test_stream_facade_raises_exception,
-    test_stream_facade_single_stream,
-    test_stream_facade_single_stream_with_multiple_slices,
-    test_stream_facade_single_stream_with_multiple_slices_with_concurrency_level_two,
-    test_stream_facade_single_stream_with_primary_key,
-    test_stream_facade_single_stream_with_single_slice,
-)
-from unit_tests.sources.streams.concurrent.scenarios.thread_based_concurrent_stream_scenarios import (
-    test_concurrent_cdk_multiple_streams,
-    test_concurrent_cdk_partition_raises_exception,
-    test_concurrent_cdk_single_stream,
-    test_concurrent_cdk_single_stream_multiple_partitions,
-    test_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_two,
-    test_concurrent_cdk_single_stream_with_primary_key,
-)
-
-scenarios = [
-    test_concurrent_cdk_single_stream,
-    test_concurrent_cdk_multiple_streams,
-    test_concurrent_cdk_single_stream_multiple_partitions,
-    test_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_two,
-    test_concurrent_cdk_single_stream_with_primary_key,
-    test_concurrent_cdk_partition_raises_exception,
-    # test streams built using the facade
-    test_stream_facade_single_stream,
-    test_stream_facade_multiple_streams,
-    test_stream_facade_single_stream_with_primary_key,
-    test_stream_facade_single_stream_with_single_slice,
-    test_stream_facade_single_stream_with_multiple_slices,
-    test_stream_facade_single_stream_with_multiple_slices_with_concurrency_level_two,
-    test_stream_facade_raises_exception,
-    test_incremental_stream_with_slice_boundaries,
-    test_incremental_stream_without_slice_boundaries,
-    test_incremental_stream_with_many_slices_but_without_slice_boundaries,
-    test_incremental_stream_with_slice_boundaries_no_input_state,
-    test_incremental_stream_with_slice_boundaries_with_concurrent_state,
-    test_incremental_stream_with_slice_boundaries_with_legacy_state,
-    test_incremental_stream_without_slice_boundaries_no_input_state,
-    test_incremental_stream_without_slice_boundaries_with_concurrent_state,
-    test_incremental_stream_without_slice_boundaries_with_legacy_state,
-]
-
-
-@pytest.mark.parametrize("scenario", scenarios, ids=[s.name for s in scenarios])
-@freeze_time("2023-06-09T00:00:00Z")
-def test_concurrent_read(scenario: TestScenario) -> None:
-    verify_read(scenario)
-
-
-@pytest.mark.parametrize("scenario", scenarios, ids=[s.name for s in scenarios])
-def test_concurrent_discover(capsys: CaptureFixture[str], tmp_path: PosixPath, scenario: TestScenario) -> None:
-    verify_discover(capsys, tmp_path, scenario)