airbyte_cdk.sources.concurrent_source.concurrent_source_adapter

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5import logging
  6from abc import ABC
  7from datetime import timedelta
  8from typing import Any, Callable, Iterator, List, Mapping, MutableMapping, Optional, Tuple
  9
 10from airbyte_cdk.models import AirbyteMessage, AirbyteStateMessage, ConfiguredAirbyteCatalog
 11from airbyte_cdk.sources import AbstractSource
 12from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
 13from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 14from airbyte_cdk.sources.streams import Stream
 15from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 16from airbyte_cdk.sources.streams.concurrent.abstract_stream_facade import AbstractStreamFacade
 17from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
 18from airbyte_cdk.sources.streams.concurrent.cursor import (
 19    ConcurrentCursor,
 20    Cursor,
 21    CursorField,
 22    CursorValueType,
 23    FinalStateCursor,
 24    GapType,
 25)
 26from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
 27    AbstractStreamStateConverter,
 28)
 29
 30DEFAULT_LOOKBACK_SECONDS = 0
 31
 32
 33class ConcurrentSourceAdapter(AbstractSource, ABC):
 34    def __init__(self, concurrent_source: ConcurrentSource, **kwargs: Any) -> None:
 35        """
 36        ConcurrentSourceAdapter is a Source that wraps a concurrent source and exposes it as a regular source.
 37
 38        The source's streams are still defined through the streams() method.
 39        Streams wrapped in a StreamFacade will be processed concurrently.
 40        Other streams will be processed sequentially as a later step.
 41        """
 42        self._concurrent_source = concurrent_source
 43        super().__init__(**kwargs)
 44
 45    def read(
 46        self,
 47        logger: logging.Logger,
 48        config: Mapping[str, Any],
 49        catalog: ConfiguredAirbyteCatalog,
 50        state: Optional[List[AirbyteStateMessage]] = None,
 51    ) -> Iterator[AirbyteMessage]:
 52        abstract_streams = self._select_abstract_streams(config, catalog)
 53        concurrent_stream_names = {stream.name for stream in abstract_streams}
 54        configured_catalog_for_regular_streams = ConfiguredAirbyteCatalog(
 55            streams=[
 56                stream
 57                for stream in catalog.streams
 58                if stream.stream.name not in concurrent_stream_names
 59            ]
 60        )
 61        if abstract_streams:
 62            yield from self._concurrent_source.read(abstract_streams)
 63        if configured_catalog_for_regular_streams.streams:
 64            yield from super().read(logger, config, configured_catalog_for_regular_streams, state)
 65
 66    def _select_abstract_streams(
 67        self, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog
 68    ) -> List[AbstractStream]:
 69        """
 70        Selects streams that can be processed concurrently and returns their abstract representations.
 71        """
 72        all_streams = self.streams(config)
 73        stream_name_to_instance: Mapping[str, Stream] = {s.name: s for s in all_streams}
 74        abstract_streams: List[AbstractStream] = []
 75        for configured_stream in configured_catalog.streams:
 76            stream_instance = stream_name_to_instance.get(configured_stream.stream.name)
 77            if not stream_instance:
 78                continue
 79
 80            if isinstance(stream_instance, AbstractStreamFacade):
 81                abstract_streams.append(stream_instance.get_underlying_stream())
 82        return abstract_streams
 83
 84    def convert_to_concurrent_stream(
 85        self,
 86        logger: logging.Logger,
 87        stream: Stream,
 88        state_manager: ConnectorStateManager,
 89        cursor: Optional[Cursor] = None,
 90    ) -> Stream:
 91        """
 92        Prepares a stream for concurrent processing by initializing or assigning a cursor,
 93        managing the stream's state, and returning an updated Stream instance.
 94        """
 95        state: MutableMapping[str, Any] = {}
 96
 97        if cursor:
 98            state = state_manager.get_stream_state(stream.name, stream.namespace)
 99
100            stream.cursor = cursor  # type: ignore[assignment]  # cursor is of type ConcurrentCursor, which inherits from Cursor
101            if hasattr(stream, "parent"):
102                stream.parent.cursor = cursor
103        else:
104            cursor = FinalStateCursor(
105                stream_name=stream.name,
106                stream_namespace=stream.namespace,
107                message_repository=self.message_repository,  # type: ignore[arg-type]  # _default_message_repository will be returned in the worst case
108            )
109        return StreamFacade.create_from_stream(stream, self, logger, state, cursor)
110
111    def initialize_cursor(
112        self,
113        stream: Stream,
114        state_manager: ConnectorStateManager,
115        converter: AbstractStreamStateConverter,
116        slice_boundary_fields: Optional[Tuple[str, str]],
117        start: Optional[CursorValueType],
118        end_provider: Callable[[], CursorValueType],
119        lookback_window: Optional[GapType] = None,
120        slice_range: Optional[GapType] = None,
121    ) -> Optional[ConcurrentCursor]:
122        lookback_window = lookback_window or timedelta(seconds=DEFAULT_LOOKBACK_SECONDS)
123
124        cursor_field_name = stream.cursor_field
125
126        if cursor_field_name:
127            if not isinstance(cursor_field_name, str):
128                raise ValueError(
129                    f"Cursor field type must be a string, but received {type(cursor_field_name).__name__}."
130                )
131
132            return ConcurrentCursor(
133                stream.name,
134                stream.namespace,
135                state_manager.get_stream_state(stream.name, stream.namespace),
136                self.message_repository,  # type: ignore[arg-type]  # _default_message_repository will be returned in the worst case
137                state_manager,
138                converter,
139                CursorField(cursor_field_name),
140                slice_boundary_fields,
141                start,
142                end_provider,
143                lookback_window,
144                slice_range,
145            )
146
147        return None