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

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

@@ -1,418 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-import logging
-
-from airbyte_cdk.sources.message import InMemoryMessageRepository
-from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
-from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenarioBuilder
-from unit_tests.sources.streams.concurrent.scenarios.thread_based_concurrent_stream_source_builder import (
-    AlwaysAvailableAvailabilityStrategy,
-    ConcurrentSourceBuilder,
-    InMemoryPartition,
-    InMemoryPartitionGenerator,
-)
-
-_message_repository = InMemoryMessageRepository()
-
-_id_only_stream = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=[],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-_id_only_stream_with_slice_logger = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=[],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-_id_only_stream_with_primary_key = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=["id"],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-_id_only_stream_multiple_partitions = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [
-            InMemoryPartition("partition1", "stream1", {"p": "1"}, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")]),
-            InMemoryPartition("partition2", "stream1", {"p": "2"}, [Record({"id": "3"}, "stream1"), Record({"id": "4"}, "stream1")]),
-        ]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=[],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-_id_only_stream_multiple_partitions_concurrency_level_two = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [
-            InMemoryPartition("partition1", "stream1", {"p": "1"}, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")]),
-            InMemoryPartition("partition2", "stream1", {"p": "2"}, [Record({"id": "3"}, "stream1"), Record({"id": "4"}, "stream1")]),
-        ]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=[],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-_stream_raising_exception = DefaultStream(
-    partition_generator=InMemoryPartitionGenerator(
-        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), ValueError("test exception")])]
-    ),
-    name="stream1",
-    json_schema={
-        "type": "object",
-        "properties": {
-            "id": {"type": ["null", "string"]},
-        },
-    },
-    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-    primary_key=[],
-    cursor_field=None,
-    logger=logging.getLogger("test_logger"),
-    cursor=FinalStateCursor(stream_name="stream1", stream_namespace=None, message_repository=_message_repository),
-)
-
-test_concurrent_cdk_single_stream = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_cdk_single_stream")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _id_only_stream,
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-        ]
-    )
-    .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"})
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_concurrent_cdk_single_stream_with_primary_key = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_cdk_single_stream_with_primary_key")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _id_only_stream_with_primary_key,
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .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"],
-                    "source_defined_primary_key": [["id"]],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_concurrent_cdk_multiple_streams = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_cdk_multiple_streams")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _id_only_stream,
-                DefaultStream(
-                    partition_generator=InMemoryPartitionGenerator(
-                        [
-                            InMemoryPartition(
-                                "partition1",
-                                "stream2",
-                                None,
-                                [Record({"id": "10", "key": "v1"}, "stream2"), Record({"id": "20", "key": "v2"}, "stream2")],
-                            )
-                        ]
-                    ),
-                    name="stream2",
-                    json_schema={
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                            "key": {"type": ["null", "string"]},
-                        },
-                    },
-                    availability_strategy=AlwaysAvailableAvailabilityStrategy(),
-                    primary_key=[],
-                    cursor_field=None,
-                    logger=logging.getLogger("test_logger"),
-                    cursor=FinalStateCursor(stream_name="stream2", stream_namespace=None, message_repository=_message_repository),
-                ),
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-            {"data": {"id": "2"}, "stream": "stream1"},
-            {"data": {"id": "10", "key": "v1"}, "stream": "stream2"},
-            {"data": {"id": "20", "key": "v2"}, "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"]},
-                            "key": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream2",
-                    "supported_sync_modes": ["full_refresh"],
-                },
-            ]
-        }
-    )
-    .build()
-)
-
-test_concurrent_cdk_partition_raises_exception = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_partition_raises_exception")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _stream_raising_exception,
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .set_expected_records(
-        [
-            {"data": {"id": "1"}, "stream": "stream1"},
-        ]
-    )
-    .set_expected_read_error(ValueError, "test exception")
-    .set_expected_catalog(
-        {
-            "streams": [
-                {
-                    "json_schema": {
-                        "type": "object",
-                        "properties": {
-                            "id": {"type": ["null", "string"]},
-                        },
-                    },
-                    "name": "stream1",
-                    "supported_sync_modes": ["full_refresh"],
-                }
-            ]
-        }
-    )
-    .build()
-)
-
-test_concurrent_cdk_single_stream_multiple_partitions = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_cdk_single_stream_multiple_partitions")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _id_only_stream_multiple_partitions,
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .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_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_two = (
-    TestScenarioBuilder()
-    .set_name("test_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_2")
-    .set_config({})
-    .set_source_builder(
-        ConcurrentSourceBuilder()
-        .set_streams(
-            [
-                _id_only_stream_multiple_partitions_concurrency_level_two,
-            ]
-        )
-        .set_message_repository(_message_repository)
-    )
-    .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()
-)

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

@@ -1,142 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-import json
-import logging
-from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
-
-from airbyte_cdk.models import 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.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.availability_strategy import AbstractAvailabilityStrategy, StreamAvailability, StreamAvailable
-from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
-from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
-from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from airbyte_cdk.sources.streams.core import StreamData
-from airbyte_cdk.sources.utils.slice_logger import SliceLogger
-from airbyte_protocol.models import ConfiguredAirbyteStream
-from unit_tests.sources.file_based.scenarios.scenario_builder import SourceBuilder
-
-
-class LegacyStream(Stream):
-    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
-        return None
-
-    def read_records(
-        self,
-        sync_mode: SyncMode,
-        cursor_field: Optional[List[str]] = None,
-        stream_slice: Optional[Mapping[str, Any]] = None,
-        stream_state: Optional[Mapping[str, Any]] = None,
-    ) -> Iterable[StreamData]:
-        yield from []
-
-
-class ConcurrentCdkSource(ConcurrentSourceAdapter):
-    def __init__(self, streams: List[DefaultStream], message_repository: Optional[MessageRepository], max_workers, timeout_in_seconds):
-        concurrent_source = ConcurrentSource.create(1, 1, streams[0]._logger, NeverLogSliceLogger(), message_repository)
-        super().__init__(concurrent_source)
-        self._streams = streams
-        self._message_repository = message_repository
-
-    def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
-        # Check is not verified because it is up to the source to implement this method
-        return True, None
-
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        return [StreamFacade(s, LegacyStream(), FinalStateCursor(stream_name=s.name, stream_namespace=s.namespace, message_repository=self.message_repository), NeverLogSliceLogger(), s._logger) for s in self._streams]
-
-    def spec(self, *args: Any, **kwargs: Any) -> ConnectorSpecification:
-        return ConnectorSpecification(connectionSpecification={})
-
-    def read_catalog(self, catalog_path: str) -> ConfiguredAirbyteCatalog:
-        return ConfiguredAirbyteCatalog(
-            streams=[
-                ConfiguredAirbyteStream(
-                    stream=StreamFacade(s, LegacyStream(), FinalStateCursor(stream_name=s.name, stream_namespace=s.namespace, message_repository=InMemoryMessageRepository()), NeverLogSliceLogger(), s._logger).as_airbyte_stream(),
-                    sync_mode=SyncMode.full_refresh,
-                    destination_sync_mode=DestinationSyncMode.overwrite,
-                )
-                for s in self._streams
-            ]
-        )
-
-    @property
-    def message_repository(self) -> Union[None, MessageRepository]:
-        return self._message_repository
-
-
-class InMemoryPartitionGenerator(PartitionGenerator):
-    def __init__(self, partitions: List[Partition]):
-        self._partitions = partitions
-
-    def generate(self) -> Iterable[Partition]:
-        yield from self._partitions
-
-
-class InMemoryPartition(Partition):
-    def stream_name(self) -> str:
-        return self._stream_name
-
-    def __init__(self, name, stream_name, _slice, records):
-        self._name = name
-        self._stream_name = stream_name
-        self._slice = _slice
-        self._records = records
-        self._is_closed = False
-
-    def read(self) -> Iterable[Record]:
-        for record_or_exception in self._records:
-            if isinstance(record_or_exception, Exception):
-                raise record_or_exception
-            else:
-                yield record_or_exception
-
-    def to_slice(self) -> Optional[Mapping[str, Any]]:
-        return self._slice
-
-    def __hash__(self) -> int:
-        if self._slice:
-            # Convert the slice to a string so that it can be hashed
-            s = json.dumps(self._slice, sort_keys=True)
-            return hash((self._name, s))
-        else:
-            return hash(self._name)
-
-    def close(self) -> None:
-        self._is_closed = True
-
-    def is_closed(self) -> bool:
-        return self._is_closed
-
-
-class ConcurrentSourceBuilder(SourceBuilder[ConcurrentCdkSource]):
-    def __init__(self):
-        self._streams: List[DefaultStream] = []
-        self._message_repository = None
-
-    def build(self, configured_catalog: Optional[Mapping[str, Any]], _, __) -> ConcurrentCdkSource:
-        return ConcurrentCdkSource(self._streams, self._message_repository, 1, 1)
-
-    def set_streams(self, streams: List[DefaultStream]) -> "ConcurrentSourceBuilder":
-        self._streams = streams
-        return self
-
-    def set_message_repository(self, message_repository: MessageRepository) -> "ConcurrentSourceBuilder":
-        self._message_repository = message_repository
-        return self
-
-
-class AlwaysAvailableAvailabilityStrategy(AbstractAvailabilityStrategy):
-    def check_availability(self, logger: logging.Logger) -> StreamAvailability:
-        return StreamAvailable()
-
-
-class NeverLogSliceLogger(SliceLogger):
-    def should_log_slice_message(self, logger: logging.Logger) -> bool:
-        return False

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

@@ -1,55 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
-
-from airbyte_cdk.models import SyncMode
-from airbyte_cdk.sources.streams import Stream
-from airbyte_cdk.sources.streams.core import StreamData
-
-
-class MockStream(Stream):
-    def __init__(
-        self,
-        slices_and_records_or_exception: Iterable[Tuple[Optional[Mapping[str, Any]], Iterable[Union[Exception, Mapping[str, Any]]]]],
-        name,
-        json_schema,
-        primary_key=None,
-    ):
-        self._slices_and_records_or_exception = slices_and_records_or_exception
-        self._name = name
-        self._json_schema = json_schema
-        self._primary_key = primary_key
-
-    def read_records(
-        self,
-        sync_mode: SyncMode,
-        cursor_field: Optional[List[str]] = None,
-        stream_slice: Optional[Mapping[str, Any]] = None,
-        stream_state: Optional[Mapping[str, Any]] = None,
-    ) -> Iterable[StreamData]:
-        for _slice, records_or_exception in self._slices_and_records_or_exception:
-            if stream_slice == _slice:
-                for item in records_or_exception:
-                    if isinstance(item, Exception):
-                        raise item
-                    yield item
-
-    @property
-    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
-        return self._primary_key
-
-    @property
-    def name(self) -> str:
-        return self._name
-
-    def get_json_schema(self) -> Mapping[str, Any]:
-        return self._json_schema
-
-    def stream_slices(
-        self, *, sync_mode: SyncMode, cursor_field: Optional[List[str]] = None, stream_state: Optional[Mapping[str, Any]] = None
-    ) -> Iterable[Optional[Mapping[str, Any]]]:
-        if self._slices_and_records_or_exception:
-            yield from [_slice for _slice, records_or_exception in self._slices_and_records_or_exception]
-        else:
-            yield None