airbyte_cdk.sources.streams.availability_strategy

 1#
 2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 3#
 4
 5import logging
 6import typing
 7from abc import ABC, abstractmethod
 8from typing import Any, Mapping, Optional, Tuple
 9
10from airbyte_cdk.models import SyncMode
11from airbyte_cdk.sources.streams.core import Stream, StreamData
12
13if typing.TYPE_CHECKING:
14    from airbyte_cdk.sources import Source
15
16
17class AvailabilityStrategy(ABC):
18    """
19    Abstract base class for checking stream availability.
20    """
21
22    @abstractmethod
23    def check_availability(
24        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
25    ) -> Tuple[bool, Optional[str]]:
26        """
27        Checks stream availability.
28
29        :param stream: stream
30        :param logger: source logger
31        :param source: (optional) source
32        :return: A tuple of (boolean, str). If boolean is true, then the stream
33          is available, and no str is required. Otherwise, the stream is unavailable
34          for some reason and the str should describe what went wrong and how to
35          resolve the unavailability, if possible.
36        """
37
38    @staticmethod
39    def get_first_stream_slice(stream: Stream) -> Optional[Mapping[str, Any]]:
40        """
41        Gets the first stream_slice from a given stream's stream_slices.
42        :param stream: stream
43        :raises StopIteration: if there is no first slice to return (the stream_slices generator is empty)
44        :return: first stream slice from 'stream_slices' generator (`None` is a valid stream slice)
45        """
46        # We wrap the return output of stream_slices() because some implementations return types that are iterable,
47        # but not iterators such as lists or tuples
48        slices = iter(
49            stream.stream_slices(
50                cursor_field=stream.cursor_field,  # type: ignore[arg-type]
51                sync_mode=SyncMode.full_refresh,
52            )
53        )
54        return next(slices)
55
56    @staticmethod
57    def get_first_record_for_slice(
58        stream: Stream, stream_slice: Optional[Mapping[str, Any]]
59    ) -> StreamData:
60        """
61        Gets the first record for a stream_slice of a stream.
62
63        :param stream: stream instance from which to read records
64        :param stream_slice: stream_slice parameters for slicing the stream
65        :raises StopIteration: if there is no first record to return (the read_records generator is empty)
66        :return: StreamData containing the first record in the slice
67        """
68        # Store the original value of exit_on_rate_limit
69        original_exit_on_rate_limit = stream.exit_on_rate_limit
70
71        try:
72            # Ensure exit_on_rate_limit is safely set to True if possible
73            stream.exit_on_rate_limit = True
74
75            # We wrap the return output of read_records() because some implementations return types that are iterable,
76            # but not iterators such as lists or tuples
77            records_for_slice = iter(
78                stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
79            )
80
81            return next(records_for_slice)
82        finally:
83            # Restore the original exit_on_rate_limit value
84            stream.exit_on_rate_limit = original_exit_on_rate_limit