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

18class HttpAvailabilityStrategy(AvailabilityStrategy):
19    def check_availability(
20        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
21    ) -> Tuple[bool, Optional[str]]:
22        """
23        Check stream availability by attempting to read the first record of the
24        stream.
25
26        :param stream: stream
27        :param logger: source logger
28        :param source: (optional) source
29        :return: A tuple of (boolean, str). If boolean is true, then the stream
30          is available, and no str is required. Otherwise, the stream is unavailable
31          for some reason and the str should describe what went wrong and how to
32          resolve the unavailability, if possible.
33        """
34        reason: Optional[str]
35        try:
36            # Some streams need a stream slice to read records (e.g. if they have a SubstreamPartitionRouter)
37            # Streams that don't need a stream slice will return `None` as their first stream slice.
38            stream_slice = self.get_first_stream_slice(stream)
39        except StopIteration:
40            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
41            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
42            # without accounting for the case in which the parent stream is empty.
43            reason = f"Cannot attempt to connect to stream {stream.name} - no stream slices were found, likely because the parent stream is empty."
44            return False, reason
45        except AirbyteTracedException as error:
46            return False, error.message
47
48        try:
49            self.get_first_record_for_slice(stream, stream_slice)
50            return True, None
51        except StopIteration:
52            logger.info(f"Successfully connected to stream {stream.name}, but got 0 records.")
53            return True, None
54        except AirbyteTracedException as error:
55            return False, error.message

11@dataclass
12class ResumableFullRefreshCursor(Cursor):
13    """
14    Cursor that allows for the checkpointing of sync progress according to a synthetic cursor based on the pagination state
15    of the stream. Resumable full refresh syncs are only intended to retain state in between sync attempts of the same job
16    with the platform responsible for removing said state.
17    """
18
19    def __init__(self) -> None:
20        self._cursor: StreamState = {}
21
22    def get_stream_state(self) -> StreamState:
23        return self._cursor
24
25    def set_initial_state(self, stream_state: StreamState) -> None:
26        self._cursor = stream_state
27
28    def observe(self, stream_slice: StreamSlice, record: Record) -> None:
29        """
30        Resumable full refresh manages state using a page number so it does not need to update state by observing incoming records.
31        """
32        pass
33
34    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
35        self._cursor = stream_slice.cursor_slice
36
37    def should_be_synced(self, record: Record) -> bool:
38        """
39        Unlike date-based cursors which filter out records outside slice boundaries, resumable full refresh records exist within pages
40        that don't have filterable bounds. We should always return them.
41        """
42        return True
43
44    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
45        """
46        RFR record don't have ordering to be compared between one another.
47        """
48        return False
49
50    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
51        # A top-level RFR cursor only manages the state of a single partition
52        return self._cursor

128class ConcurrentCursor(Cursor):
129    _START_BOUNDARY = 0
130    _END_BOUNDARY = 1
131
132    def __init__(
133        self,
134        stream_name: str,
135        stream_namespace: Optional[str],
136        stream_state: Any,
137        message_repository: MessageRepository,
138        connector_state_manager: ConnectorStateManager,
139        connector_state_converter: AbstractStreamStateConverter,
140        cursor_field: CursorField,
141        slice_boundary_fields: Optional[Tuple[str, str]],
142        start: Optional[CursorValueType],
143        end_provider: Callable[[], CursorValueType],
144        lookback_window: Optional[GapType] = None,
145        slice_range: Optional[GapType] = None,
146        cursor_granularity: Optional[GapType] = None,
147        clamping_strategy: ClampingStrategy = NoClamping(),
148    ) -> None:
149        self._stream_name = stream_name
150        self._stream_namespace = stream_namespace
151        self._message_repository = message_repository
152        self._connector_state_converter = connector_state_converter
153        self._connector_state_manager = connector_state_manager
154        self._cursor_field = cursor_field
155        # To see some example where the slice boundaries might not be defined, check https://github.com/airbytehq/airbyte/blob/1ce84d6396e446e1ac2377362446e3fb94509461/airbyte-integrations/connectors/source-stripe/source_stripe/streams.py#L363-L379
156        self._slice_boundary_fields = slice_boundary_fields
157        self._start = start
158        self._end_provider = end_provider
159        self.start, self._concurrent_state = self._get_concurrent_state(stream_state)
160        self._lookback_window = lookback_window
161        self._slice_range = slice_range
162        self._most_recent_cursor_value_per_partition: MutableMapping[
163            Union[StreamSlice, Mapping[str, Any], None], Any
164        ] = {}
165        self._has_closed_at_least_one_slice = False
166        self._cursor_granularity = cursor_granularity
167        # Flag to track if the logger has been triggered (per stream)
168        self._should_be_synced_logger_triggered = False
169        self._clamping_strategy = clamping_strategy
170
171    @property
172    def state(self) -> MutableMapping[str, Any]:
173        return self._connector_state_converter.convert_to_state_message(
174            self.cursor_field, self._concurrent_state
175        )
176
177    @property
178    def cursor_field(self) -> CursorField:
179        return self._cursor_field
180
181    @property
182    def _slice_boundary_fields_wrapper(self) -> Tuple[str, str]:
183        return (
184            self._slice_boundary_fields
185            if self._slice_boundary_fields
186            else (
187                self._connector_state_converter.START_KEY,
188                self._connector_state_converter.END_KEY,
189            )
190        )
191
192    def _get_concurrent_state(
193        self, state: MutableMapping[str, Any]
194    ) -> Tuple[CursorValueType, MutableMapping[str, Any]]:
195        if self._connector_state_converter.is_state_message_compatible(state):
196            return (
197                self._start or self._connector_state_converter.zero_value,
198                self._connector_state_converter.deserialize(state),
199            )
200        return self._connector_state_converter.convert_from_sequential_state(
201            self._cursor_field, state, self._start
202        )
203
204    def observe(self, record: Record) -> None:
205        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
206            record.associated_slice
207        )
208        try:
209            cursor_value = self._extract_cursor_value(record)
210
211            if most_recent_cursor_value is None or most_recent_cursor_value < cursor_value:
212                self._most_recent_cursor_value_per_partition[record.associated_slice] = cursor_value
213        except ValueError:
214            self._log_for_record_without_cursor_value()
215
216    def _extract_cursor_value(self, record: Record) -> Any:
217        return self._connector_state_converter.parse_value(self._cursor_field.extract_value(record))
218
219    def close_partition(self, partition: Partition) -> None:
220        slice_count_before = len(self._concurrent_state.get("slices", []))
221        self._add_slice_to_state(partition)
222        if slice_count_before < len(
223            self._concurrent_state["slices"]
224        ):  # only emit if at least one slice has been processed
225            self._merge_partitions()
226            self._emit_state_message()
227        self._has_closed_at_least_one_slice = True
228
229    def _add_slice_to_state(self, partition: Partition) -> None:
230        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
231            partition.to_slice()
232        )
233
234        if self._slice_boundary_fields:
235            if "slices" not in self._concurrent_state:
236                raise RuntimeError(
237                    f"The state for stream {self._stream_name} should have at least one slice to delineate the sync start time, but no slices are present. This is unexpected. Please contact Support."
238                )
239            self._concurrent_state["slices"].append(
240                {
241                    self._connector_state_converter.START_KEY: self._extract_from_slice(
242                        partition, self._slice_boundary_fields[self._START_BOUNDARY]
243                    ),
244                    self._connector_state_converter.END_KEY: self._extract_from_slice(
245                        partition, self._slice_boundary_fields[self._END_BOUNDARY]
246                    ),
247                    self._connector_state_converter.MOST_RECENT_RECORD_KEY: most_recent_cursor_value,
248                }
249            )
250        elif most_recent_cursor_value:
251            if self._has_closed_at_least_one_slice:
252                # If we track state value using records cursor field, we can only do that if there is one partition. This is because we save
253                # the state every time we close a partition. We assume that if there are multiple slices, they need to be providing
254                # boundaries. There are cases where partitions could not have boundaries:
255                # * The cursor should be per-partition
256                # * The stream state is actually the parent stream state
257                # There might be other cases not listed above. Those are not supported today hence the stream should not use this cursor for
258                # state management. For the specific user that was affected with this issue, we need to:
259                # * Fix state tracking (which is currently broken)
260                # * Make the new version available
261                # * (Probably) ask the user to reset the stream to avoid data loss
262                raise ValueError(
263                    "Given that slice_boundary_fields is not defined and that per-partition state is not supported, only one slice is "
264                    "expected. Please contact the Airbyte team."
265                )
266
267            self._concurrent_state["slices"].append(
268                {
269                    self._connector_state_converter.START_KEY: self.start,
270                    self._connector_state_converter.END_KEY: most_recent_cursor_value,
271                    self._connector_state_converter.MOST_RECENT_RECORD_KEY: most_recent_cursor_value,
272                }
273            )
274
275    def _emit_state_message(self) -> None:
276        self._connector_state_manager.update_state_for_stream(
277            self._stream_name,
278            self._stream_namespace,
279            self.state,
280        )
281        state_message = self._connector_state_manager.create_state_message(
282            self._stream_name, self._stream_namespace
283        )
284        self._message_repository.emit_message(state_message)
285
286    def _merge_partitions(self) -> None:
287        self._concurrent_state["slices"] = self._connector_state_converter.merge_intervals(
288            self._concurrent_state["slices"]
289        )
290
291    def _extract_from_slice(self, partition: Partition, key: str) -> CursorValueType:
292        try:
293            _slice = partition.to_slice()
294            if not _slice:
295                raise KeyError(f"Could not find key `{key}` in empty slice")
296            return self._connector_state_converter.parse_value(_slice[key])  # type: ignore  # we expect the devs to specify a key that would return a CursorValueType
297        except KeyError as exception:
298            raise KeyError(
299                f"Partition is expected to have key `{key}` but could not be found"
300            ) from exception
301
302    def ensure_at_least_one_state_emitted(self) -> None:
303        """
304        The platform expect to have at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
305        called.
306        """
307        self._emit_state_message()
308
309    def stream_slices(self) -> Iterable[StreamSlice]:
310        """
311        Generating slices based on a few parameters:
312        * lookback_window: Buffer to remove from END_KEY of the highest slice
313        * slice_range: Max difference between two slices. If the difference between two slices is greater, multiple slices will be created
314        * start: `_split_per_slice_range` will clip any value to `self._start which means that:
315          * if upper is less than self._start, no slices will be generated
316          * if lower is less than self._start, self._start will be used as the lower boundary (lookback_window will not be considered in that case)
317
318        Note that the slices will overlap at their boundaries. We therefore expect to have at least the lower or the upper boundary to be
319        inclusive in the API that is queried.
320        """
321        self._merge_partitions()
322
323        if self._start is not None and self._is_start_before_first_slice():
324            yield from self._split_per_slice_range(
325                self._start,
326                self._concurrent_state["slices"][0][self._connector_state_converter.START_KEY],
327                False,
328            )
329
330        if len(self._concurrent_state["slices"]) == 1:
331            yield from self._split_per_slice_range(
332                self._calculate_lower_boundary_of_last_slice(
333                    self._concurrent_state["slices"][0][self._connector_state_converter.END_KEY]
334                ),
335                self._end_provider(),
336                True,
337            )
338        elif len(self._concurrent_state["slices"]) > 1:
339            for i in range(len(self._concurrent_state["slices"]) - 1):
340                if self._cursor_granularity:
341                    yield from self._split_per_slice_range(
342                        self._concurrent_state["slices"][i][self._connector_state_converter.END_KEY]
343                        + self._cursor_granularity,
344                        self._concurrent_state["slices"][i + 1][
345                            self._connector_state_converter.START_KEY
346                        ],
347                        False,
348                    )
349                else:
350                    yield from self._split_per_slice_range(
351                        self._concurrent_state["slices"][i][
352                            self._connector_state_converter.END_KEY
353                        ],
354                        self._concurrent_state["slices"][i + 1][
355                            self._connector_state_converter.START_KEY
356                        ],
357                        False,
358                    )
359            yield from self._split_per_slice_range(
360                self._calculate_lower_boundary_of_last_slice(
361                    self._concurrent_state["slices"][-1][self._connector_state_converter.END_KEY]
362                ),
363                self._end_provider(),
364                True,
365            )
366        else:
367            raise ValueError("Expected at least one slice")
368
369    def _is_start_before_first_slice(self) -> bool:
370        return (
371            self._start is not None
372            and self._start
373            < self._concurrent_state["slices"][0][self._connector_state_converter.START_KEY]
374        )
375
376    def _calculate_lower_boundary_of_last_slice(
377        self, lower_boundary: CursorValueType
378    ) -> CursorValueType:
379        if self._lookback_window:
380            return lower_boundary - self._lookback_window
381        return lower_boundary
382
383    def _split_per_slice_range(
384        self, lower: CursorValueType, upper: CursorValueType, upper_is_end: bool
385    ) -> Iterable[StreamSlice]:
386        if lower >= upper:
387            return
388
389        if self._start and upper < self._start:
390            return
391
392        lower = max(lower, self._start) if self._start else lower
393        if not self._slice_range or self._evaluate_upper_safely(lower, self._slice_range) >= upper:
394            clamped_lower = self._clamping_strategy.clamp(lower)
395            clamped_upper = self._clamping_strategy.clamp(upper)
396            start_value, end_value = (
397                (clamped_lower, clamped_upper - self._cursor_granularity)
398                if self._cursor_granularity and not upper_is_end
399                else (clamped_lower, clamped_upper)
400            )
401            yield StreamSlice(
402                partition={},
403                cursor_slice={
404                    self._slice_boundary_fields_wrapper[
405                        self._START_BOUNDARY
406                    ]: self._connector_state_converter.output_format(start_value),
407                    self._slice_boundary_fields_wrapper[
408                        self._END_BOUNDARY
409                    ]: self._connector_state_converter.output_format(end_value),
410                },
411            )
412        else:
413            stop_processing = False
414            current_lower_boundary = lower
415            while not stop_processing:
416                current_upper_boundary = min(
417                    self._evaluate_upper_safely(current_lower_boundary, self._slice_range), upper
418                )
419                has_reached_upper_boundary = current_upper_boundary >= upper
420
421                clamped_upper = (
422                    self._clamping_strategy.clamp(current_upper_boundary)
423                    if current_upper_boundary != upper
424                    else current_upper_boundary
425                )
426                clamped_lower = self._clamping_strategy.clamp(current_lower_boundary)
427                if clamped_lower >= clamped_upper:
428                    # clamping collapsed both values which means that it is time to stop processing
429                    # FIXME should this be replace by proper end_provider
430                    break
431                start_value, end_value = (
432                    (clamped_lower, clamped_upper - self._cursor_granularity)
433                    if self._cursor_granularity
434                    and (not upper_is_end or not has_reached_upper_boundary)
435                    else (clamped_lower, clamped_upper)
436                )
437                yield StreamSlice(
438                    partition={},
439                    cursor_slice={
440                        self._slice_boundary_fields_wrapper[
441                            self._START_BOUNDARY
442                        ]: self._connector_state_converter.output_format(start_value),
443                        self._slice_boundary_fields_wrapper[
444                            self._END_BOUNDARY
445                        ]: self._connector_state_converter.output_format(end_value),
446                    },
447                )
448                current_lower_boundary = clamped_upper
449                if current_upper_boundary >= upper:
450                    stop_processing = True
451
452    def _evaluate_upper_safely(self, lower: CursorValueType, step: GapType) -> CursorValueType:
453        """
454        Given that we set the default step at datetime.timedelta.max, we will generate an OverflowError when evaluating the next start_date
455        This method assumes that users would never enter a step that would generate an overflow. Given that would be the case, the code
456        would have broken anyway.
457        """
458        try:
459            return lower + step
460        except OverflowError:
461            return self._end_provider()
462
463    def should_be_synced(self, record: Record) -> bool:
464        """
465        Determines if a record should be synced based on its cursor value.
466        :param record: The record to evaluate
467
468        :return: True if the record's cursor value falls within the sync boundaries
469        """
470        try:
471            record_cursor_value: CursorValueType = self._extract_cursor_value(record)
472        except ValueError:
473            self._log_for_record_without_cursor_value()
474            return True
475        return self.start <= record_cursor_value <= self._end_provider()
476
477    def _log_for_record_without_cursor_value(self) -> None:
478        if not self._should_be_synced_logger_triggered:
479            LOGGER.warning(
480                f"Could not find cursor field `{self.cursor_field.cursor_field_key}` in record for stream {self._stream_name}. The incremental sync will assume it needs to be synced"
481            )
482            self._should_be_synced_logger_triggered = True