airbyte_cdk.sources.streams.concurrent.default_stream 

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5from logging import Logger
  6from typing import Any, Callable, Iterable, List, Mapping, Optional, Union
  7
  8from airbyte_cdk.models import AirbyteStream, SyncMode
  9from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
 10    ConcurrentPerPartitionCursor,
 11)
 12from airbyte_cdk.sources.declarative.incremental.sync_token_cursor import SyncTokenCursor
 13from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 14from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
 15    StreamSlicerPartitionGenerator,
 16)
 17from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 18from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability
 19from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, CursorField
 20from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 21from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 22from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 23
 24
 25class DefaultStream(AbstractStream):
 26    def __init__(
 27        self,
 28        partition_generator: PartitionGenerator,
 29        name: str,
 30        json_schema: Union[Mapping[str, Any], Callable[[], Mapping[str, Any]]],
 31        primary_key: List[str],
 32        cursor_field: Optional[CursorField],
 33        logger: Logger,
 34        cursor: Cursor,
 35        namespace: Optional[str] = None,
 36        supports_file_transfer: bool = False,
 37        block_simultaneous_read: str = "",
 38        source_defined_incremental: bool = False,
 39    ) -> None:
 40        self._stream_partition_generator = partition_generator
 41        self._name = name
 42        self._json_schema = json_schema
 43        self._primary_key = primary_key
 44        self._cursor_field = cursor_field
 45        self._logger = logger
 46        self._cursor = cursor
 47        self._namespace = namespace
 48        self._supports_file_transfer = supports_file_transfer
 49        self._block_simultaneous_read = block_simultaneous_read
 50        self._source_defined_incremental = source_defined_incremental
 51
 52    def generate_partitions(self) -> Iterable[Partition]:
 53        yield from self._stream_partition_generator.generate()
 54
 55    @property
 56    def name(self) -> str:
 57        return self._name
 58
 59    @property
 60    def namespace(self) -> Optional[str]:
 61        return self._namespace
 62
 63    @property
 64    def cursor_field(self) -> Optional[str]:
 65        return self._cursor_field.cursor_field_key if self._cursor_field else None
 66
 67    def get_json_schema(self) -> Mapping[str, Any]:
 68        return self._json_schema() if callable(self._json_schema) else self._json_schema
 69
 70    def as_airbyte_stream(self) -> AirbyteStream:
 71        stream = AirbyteStream(
 72            name=self.name,
 73            json_schema=dict(self.get_json_schema()),
 74            supported_sync_modes=[SyncMode.full_refresh],
 75            is_resumable=False,
 76            is_file_based=self._supports_file_transfer,
 77        )
 78
 79        if self._namespace:
 80            stream.namespace = self._namespace
 81
 82        if self._cursor_field:
 83            stream.source_defined_cursor = (
 84                not self._cursor_field.supports_catalog_defined_cursor_field
 85            )
 86            stream.is_resumable = True
 87            stream.supported_sync_modes.append(SyncMode.incremental)
 88            stream.default_cursor_field = [self._cursor_field.cursor_field_key]
 89        elif self._source_defined_incremental:
 90            stream.source_defined_cursor = True
 91            stream.is_resumable = True
 92            stream.supported_sync_modes.append(SyncMode.incremental)
 93            stream.default_cursor_field = []
 94        keys = self._primary_key
 95        if keys and len(keys) > 0:
 96            stream.source_defined_primary_key = [[key] for key in keys]
 97
 98        return stream
 99
100    def log_stream_sync_configuration(self) -> None:
101        self._logger.debug(
102            f"Syncing stream instance: {self.name}",
103            extra={
104                "primary_key": self._primary_key,
105                "cursor_field": self.cursor_field,
106            },
107        )
108
109    @property
110    def cursor(self) -> Cursor:
111        return self._cursor
112
113    @property
114    def block_simultaneous_read(self) -> str:
115        """Returns the blocking group name for this stream, or empty string if no blocking"""
116        return self._block_simultaneous_read
117
118    @block_simultaneous_read.setter
119    def block_simultaneous_read(self, value: str) -> None:
120        self._block_simultaneous_read = value
121
122    def get_partition_router(self) -> PartitionRouter | None:
123        """Return the partition router for this stream, or None if not available."""
124        if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator):
125            return None
126        stream_slicer = self._stream_partition_generator._stream_slicer
127        if isinstance(stream_slicer, (SyncTokenCursor, ConcurrentPerPartitionCursor)):
128            return stream_slicer.partition_router
129        return None
130
131    def check_availability(self) -> StreamAvailability:
132        """
133        Check stream availability by attempting to read the first record of the stream.
134        """
135        try:
136            partition = next(iter(self.generate_partitions()))
137        except StopIteration:
138            # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
139            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
140            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
141            # without accounting for the case in which the parent stream is empty.
142            return StreamAvailability.unavailable(
143                f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
144            )
145        except AirbyteTracedException as error:
146            return StreamAvailability.unavailable(
147                error.message or error.internal_message or "<no error message>"
148            )
149
150        try:
151            next(iter(partition.read()))
152            return StreamAvailability.available()
153        except StopIteration:
154            self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
155            return StreamAvailability.available()
156        except AirbyteTracedException as error:
157            return StreamAvailability.unavailable(
158                error.message or error.internal_message or "<no error message>"
159            )
class DefaultStream(airbyte_cdk.sources.streams.concurrent.abstract_stream.AbstractStream): 
 26class DefaultStream(AbstractStream):
 27    def __init__(
 28        self,
 29        partition_generator: PartitionGenerator,
 30        name: str,
 31        json_schema: Union[Mapping[str, Any], Callable[[], Mapping[str, Any]]],
 32        primary_key: List[str],
 33        cursor_field: Optional[CursorField],
 34        logger: Logger,
 35        cursor: Cursor,
 36        namespace: Optional[str] = None,
 37        supports_file_transfer: bool = False,
 38        block_simultaneous_read: str = "",
 39        source_defined_incremental: bool = False,
 40    ) -> None:
 41        self._stream_partition_generator = partition_generator
 42        self._name = name
 43        self._json_schema = json_schema
 44        self._primary_key = primary_key
 45        self._cursor_field = cursor_field
 46        self._logger = logger
 47        self._cursor = cursor
 48        self._namespace = namespace
 49        self._supports_file_transfer = supports_file_transfer
 50        self._block_simultaneous_read = block_simultaneous_read
 51        self._source_defined_incremental = source_defined_incremental
 52
 53    def generate_partitions(self) -> Iterable[Partition]:
 54        yield from self._stream_partition_generator.generate()
 55
 56    @property
 57    def name(self) -> str:
 58        return self._name
 59
 60    @property
 61    def namespace(self) -> Optional[str]:
 62        return self._namespace
 63
 64    @property
 65    def cursor_field(self) -> Optional[str]:
 66        return self._cursor_field.cursor_field_key if self._cursor_field else None
 67
 68    def get_json_schema(self) -> Mapping[str, Any]:
 69        return self._json_schema() if callable(self._json_schema) else self._json_schema
 70
 71    def as_airbyte_stream(self) -> AirbyteStream:
 72        stream = AirbyteStream(
 73            name=self.name,
 74            json_schema=dict(self.get_json_schema()),
 75            supported_sync_modes=[SyncMode.full_refresh],
 76            is_resumable=False,
 77            is_file_based=self._supports_file_transfer,
 78        )
 79
 80        if self._namespace:
 81            stream.namespace = self._namespace
 82
 83        if self._cursor_field:
 84            stream.source_defined_cursor = (
 85                not self._cursor_field.supports_catalog_defined_cursor_field
 86            )
 87            stream.is_resumable = True
 88            stream.supported_sync_modes.append(SyncMode.incremental)
 89            stream.default_cursor_field = [self._cursor_field.cursor_field_key]
 90        elif self._source_defined_incremental:
 91            stream.source_defined_cursor = True
 92            stream.is_resumable = True
 93            stream.supported_sync_modes.append(SyncMode.incremental)
 94            stream.default_cursor_field = []
 95        keys = self._primary_key
 96        if keys and len(keys) > 0:
 97            stream.source_defined_primary_key = [[key] for key in keys]
 98
 99        return stream
100
101    def log_stream_sync_configuration(self) -> None:
102        self._logger.debug(
103            f"Syncing stream instance: {self.name}",
104            extra={
105                "primary_key": self._primary_key,
106                "cursor_field": self.cursor_field,
107            },
108        )
109
110    @property
111    def cursor(self) -> Cursor:
112        return self._cursor
113
114    @property
115    def block_simultaneous_read(self) -> str:
116        """Returns the blocking group name for this stream, or empty string if no blocking"""
117        return self._block_simultaneous_read
118
119    @block_simultaneous_read.setter
120    def block_simultaneous_read(self, value: str) -> None:
121        self._block_simultaneous_read = value
122
123    def get_partition_router(self) -> PartitionRouter | None:
124        """Return the partition router for this stream, or None if not available."""
125        if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator):
126            return None
127        stream_slicer = self._stream_partition_generator._stream_slicer
128        if isinstance(stream_slicer, (SyncTokenCursor, ConcurrentPerPartitionCursor)):
129            return stream_slicer.partition_router
130        return None
131
132    def check_availability(self) -> StreamAvailability:
133        """
134        Check stream availability by attempting to read the first record of the stream.
135        """
136        try:
137            partition = next(iter(self.generate_partitions()))
138        except StopIteration:
139            # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
140            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
141            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
142            # without accounting for the case in which the parent stream is empty.
143            return StreamAvailability.unavailable(
144                f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
145            )
146        except AirbyteTracedException as error:
147            return StreamAvailability.unavailable(
148                error.message or error.internal_message or "<no error message>"
149            )
150
151        try:
152            next(iter(partition.read()))
153            return StreamAvailability.available()
154        except StopIteration:
155            self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
156            return StreamAvailability.available()
157        except AirbyteTracedException as error:
158            return StreamAvailability.unavailable(
159                error.message or error.internal_message or "<no error message>"
160            )

AbstractStream is an experimental interface for streams developed as part of the Concurrent CDK. This interface is not yet stable and may change in the future. Use at your own risk.

Why create a new interface instead of adding concurrency capabilities the existing Stream? We learnt a lot since the initial design of the Stream interface, and we wanted to take the opportunity to improve.

High level, the changes we are targeting are:

  • Removing superfluous or leaky parameters from the methods' interfaces
  • Using composition instead of inheritance to add new capabilities

To allow us to iterate fast while ensuring backwards compatibility, we are creating a new interface with a facade object that will bridge the old and the new interfaces. Source connectors that wish to leverage concurrency need to implement this new interface. An example will be available shortly

Current restrictions on sources that implement this interface. Not all of these restrictions will be lifted in the future, but most will as we iterate on the design.

  • Only full refresh is supported. This will be addressed in the future.
  • The read method does not accept a cursor_field. Streams must be internally aware of the cursor field to use. User-defined cursor fields can be implemented by modifying the connector's main method to instantiate the streams with the configured cursor field.
  • Streams cannot return user-friendly messages by overriding Stream.get_error_display_message. This will be addressed in the future.
  • The Stream's behavior cannot depend on a namespace
  • TypeTransformer is not supported. This will be addressed in the future.
  • Nested cursor and primary keys are not supported
DefaultStream( partition_generator: airbyte_cdk.sources.streams.concurrent.partitions.partition_generator.PartitionGenerator, name: str, json_schema: Union[Mapping[str, Any], Callable[[], Mapping[str, Any]]], primary_key: List[str], cursor_field: Optional[airbyte_cdk.CursorField], logger: logging.Logger, cursor: airbyte_cdk.Cursor, namespace: Optional[str] = None, supports_file_transfer: bool = False, block_simultaneous_read: str = '', source_defined_incremental: bool = False) 
27    def __init__(
28        self,
29        partition_generator: PartitionGenerator,
30        name: str,
31        json_schema: Union[Mapping[str, Any], Callable[[], Mapping[str, Any]]],
32        primary_key: List[str],
33        cursor_field: Optional[CursorField],
34        logger: Logger,
35        cursor: Cursor,
36        namespace: Optional[str] = None,
37        supports_file_transfer: bool = False,
38        block_simultaneous_read: str = "",
39        source_defined_incremental: bool = False,
40    ) -> None:
41        self._stream_partition_generator = partition_generator
42        self._name = name
43        self._json_schema = json_schema
44        self._primary_key = primary_key
45        self._cursor_field = cursor_field
46        self._logger = logger
47        self._cursor = cursor
48        self._namespace = namespace
49        self._supports_file_transfer = supports_file_transfer
50        self._block_simultaneous_read = block_simultaneous_read
51        self._source_defined_incremental = source_defined_incremental
def generate_partitions( self) -> Iterable[airbyte_cdk.sources.streams.concurrent.partitions.partition.Partition]: 
53    def generate_partitions(self) -> Iterable[Partition]:
54        yield from self._stream_partition_generator.generate()

Generates the partitions that will be read by this stream.

Returns

An iterable of partitions.

name: str 
56    @property
57    def name(self) -> str:
58        return self._name
Returns

The stream name

namespace: Optional[str] 
60    @property
61    def namespace(self) -> Optional[str]:
62        return self._namespace
cursor_field: Optional[str] 
64    @property
65    def cursor_field(self) -> Optional[str]:
66        return self._cursor_field.cursor_field_key if self._cursor_field else None

Override to return the default cursor field used by this stream e.g: an API entity might always use created_at as the cursor field.

Returns

The name of the field used as a cursor. Nested cursor fields are not supported.

def get_json_schema(self) -> Mapping[str, Any]: 
68    def get_json_schema(self) -> Mapping[str, Any]:
69        return self._json_schema() if callable(self._json_schema) else self._json_schema
Returns

A dict of the JSON schema representing this stream.

def as_airbyte_stream( self) -> airbyte_protocol_dataclasses.models.airbyte_protocol.AirbyteStream: 
71    def as_airbyte_stream(self) -> AirbyteStream:
72        stream = AirbyteStream(
73            name=self.name,
74            json_schema=dict(self.get_json_schema()),
75            supported_sync_modes=[SyncMode.full_refresh],
76            is_resumable=False,
77            is_file_based=self._supports_file_transfer,
78        )
79
80        if self._namespace:
81            stream.namespace = self._namespace
82
83        if self._cursor_field:
84            stream.source_defined_cursor = (
85                not self._cursor_field.supports_catalog_defined_cursor_field
86            )
87            stream.is_resumable = True
88            stream.supported_sync_modes.append(SyncMode.incremental)
89            stream.default_cursor_field = [self._cursor_field.cursor_field_key]
90        elif self._source_defined_incremental:
91            stream.source_defined_cursor = True
92            stream.is_resumable = True
93            stream.supported_sync_modes.append(SyncMode.incremental)
94            stream.default_cursor_field = []
95        keys = self._primary_key
96        if keys and len(keys) > 0:
97            stream.source_defined_primary_key = [[key] for key in keys]
98
99        return stream
Returns

A dict of the JSON schema representing this stream.

def log_stream_sync_configuration(self) -> None: 
101    def log_stream_sync_configuration(self) -> None:
102        self._logger.debug(
103            f"Syncing stream instance: {self.name}",
104            extra={
105                "primary_key": self._primary_key,
106                "cursor_field": self.cursor_field,
107            },
108        )

Logs the stream's configuration for debugging purposes.

cursor: airbyte_cdk.Cursor 
110    @property
111    def cursor(self) -> Cursor:
112        return self._cursor
Returns

The cursor associated with this stream.

block_simultaneous_read: str 
114    @property
115    def block_simultaneous_read(self) -> str:
116        """Returns the blocking group name for this stream, or empty string if no blocking"""
117        return self._block_simultaneous_read

Returns the blocking group name for this stream, or empty string if no blocking

def get_partition_router( self) -> airbyte_cdk.sources.declarative.partition_routers.PartitionRouter | None: 
123    def get_partition_router(self) -> PartitionRouter | None:
124        """Return the partition router for this stream, or None if not available."""
125        if not isinstance(self._stream_partition_generator, StreamSlicerPartitionGenerator):
126            return None
127        stream_slicer = self._stream_partition_generator._stream_slicer
128        if isinstance(stream_slicer, (SyncTokenCursor, ConcurrentPerPartitionCursor)):
129            return stream_slicer.partition_router
130        return None

Return the partition router for this stream, or None if not available.

def check_availability( self) -> airbyte_cdk.sources.streams.concurrent.availability_strategy.StreamAvailability: 
132    def check_availability(self) -> StreamAvailability:
133        """
134        Check stream availability by attempting to read the first record of the stream.
135        """
136        try:
137            partition = next(iter(self.generate_partitions()))
138        except StopIteration:
139            # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
140            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
141            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
142            # without accounting for the case in which the parent stream is empty.
143            return StreamAvailability.unavailable(
144                f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
145            )
146        except AirbyteTracedException as error:
147            return StreamAvailability.unavailable(
148                error.message or error.internal_message or "<no error message>"
149            )
150
151        try:
152            next(iter(partition.read()))
153            return StreamAvailability.available()
154        except StopIteration:
155            self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
156            return StreamAvailability.available()
157        except AirbyteTracedException as error:
158            return StreamAvailability.unavailable(
159                error.message or error.internal_message or "<no error message>"
160            )

Check stream availability by attempting to read the first record of the stream.