airbyte_cdk.sources.declarative.incremental

 96class ConcurrentCursorFactory:
 97    def __init__(self, create_function: Callable[..., ConcurrentCursor]):
 98        self._create_function = create_function
 99
100    def create(
101        self, stream_state: Mapping[str, Any], runtime_lookback_window: Optional[timedelta]
102    ) -> ConcurrentCursor:
103        return self._create_function(
104            stream_state=stream_state, runtime_lookback_window=runtime_lookback_window
105        )

108class ConcurrentPerPartitionCursor(Cursor):
109    """
110    Manages state per partition when a stream has many partitions, preventing data loss or duplication.
111
112    Attributes:
113        DEFAULT_MAX_PARTITIONS_NUMBER (int): Maximum number of partitions to retain in memory (default is 10,000). This limit needs to be higher than the number of threads we might enqueue (which is represented by ThreadPoolManager.DEFAULT_MAX_QUEUE_SIZE). If not, we could have partitions that have been generated and submitted to the ThreadPool but got deleted from the ConcurrentPerPartitionCursor and when closing them, it will generate KeyError.
114
115    - **Partition Limitation Logic**
116      Ensures the number of tracked partitions does not exceed the specified limit to prevent memory overuse. Oldest partitions are removed when the limit is reached.
117
118    - **Global Cursor Fallback**
119      New partitions use global state as the initial state to progress the state for deleted or new partitions. The history data added after the initial sync will be missing.
120
121    CurrentPerPartitionCursor expects the state of the ConcurrentCursor to follow the format {cursor_field: cursor_value}.
122    """
123
124    DEFAULT_MAX_PARTITIONS_NUMBER = 25_000
125    SWITCH_TO_GLOBAL_LIMIT = 10_000
126    _NO_STATE: Mapping[str, Any] = {}
127    _NO_CURSOR_STATE: Mapping[str, Any] = {}
128    _GLOBAL_STATE_KEY = "state"
129    _PERPARTITION_STATE_KEY = "states"
130    _IS_PARTITION_DUPLICATION_LOGGED = False
131    _PARENT_STATE = 0
132    _GENERATION_SEQUENCE = 1
133
134    def __init__(
135        self,
136        cursor_factory: ConcurrentCursorFactory,
137        partition_router: PartitionRouter,
138        stream_name: str,
139        stream_namespace: Optional[str],
140        stream_state: Any,
141        message_repository: MessageRepository,
142        connector_state_manager: ConnectorStateManager,
143        connector_state_converter: AbstractStreamStateConverter,
144        cursor_field: CursorField,
145        use_global_cursor: bool = False,
146        attempt_to_create_cursor_if_not_provided: bool = False,
147    ) -> None:
148        self._global_cursor: Optional[StreamState] = {}
149        self._stream_name = stream_name
150        self._stream_namespace = stream_namespace
151        self._message_repository = message_repository
152        self._connector_state_manager = connector_state_manager
153        self._connector_state_converter = connector_state_converter
154        self._cursor_field = cursor_field
155
156        self._cursor_factory = cursor_factory  # self._cursor_factory is flagged as private but is used in model_to_component_factory to ease pagination reset instantiation
157        self._partition_router = partition_router
158
159        # The dict is ordered to ensure that once the maximum number of partitions is reached,
160        # the oldest partitions can be efficiently removed, maintaining the most recent partitions.
161        self._cursor_per_partition: OrderedDict[str, ConcurrentCursor] = OrderedDict()
162        self._semaphore_per_partition: OrderedDict[str, threading.Semaphore] = OrderedDict()
163
164        # Parent-state tracking: store each partition’s parent state in creation order
165        self._partition_parent_state_map: OrderedDict[str, tuple[Mapping[str, Any], int]] = (
166            OrderedDict()
167        )
168        self._parent_state: Optional[StreamState] = None
169
170        # Tracks when the last slice for partition is emitted
171        self._partitions_done_generating_stream_slices: set[str] = set()
172        # Used to track the index of partitions that are not closed yet
173        self._processing_partitions_indexes: List[int] = list()
174        self._generated_partitions_count: int = 0
175        # Dictionary to map partition keys to their index
176        self._partition_key_to_index: dict[str, int] = {}
177
178        self._lock = threading.Lock()
179        self._request_window_splitting_cursor: Optional[ConcurrentCursor] = None
180        self._lookback_window: int = 0
181        self._new_global_cursor: Optional[StreamState] = None
182        self._number_of_partitions: int = 0
183        self._use_global_cursor: bool = use_global_cursor
184        self._partition_serializer = PerPartitionKeySerializer()
185
186        # Track the last time a state message was emitted
187        self._last_emission_time: float = 0.0
188        self._timer = Timer()
189
190        self._set_initial_state(stream_state)
191
192        # FIXME this is a temporary field the time of the migration from declarative cursors to concurrent ones
193        self._attempt_to_create_cursor_if_not_provided = attempt_to_create_cursor_if_not_provided
194        self._synced_some_data = False
195        self._logged_regarding_datetime_format_error = False
196
197    @property
198    def cursor_field(self) -> CursorField:
199        return self._cursor_field
200
201    @property
202    def state(self) -> MutableMapping[str, Any]:
203        state: dict[str, Any] = {"use_global_cursor": self._use_global_cursor}
204        if not self._use_global_cursor:
205            states = []
206            for partition_tuple, cursor in self._cursor_per_partition.items():
207                if cursor.state:
208                    states.append(
209                        {
210                            "partition": self._to_dict(partition_tuple),
211                            "cursor": copy.deepcopy(cursor.state),
212                        }
213                    )
214            state[self._PERPARTITION_STATE_KEY] = states
215
216        if self._global_cursor:
217            state[self._GLOBAL_STATE_KEY] = self._global_cursor
218        if self._lookback_window is not None:
219            state["lookback_window"] = self._lookback_window
220        if self._parent_state is not None:
221            state["parent_state"] = self._parent_state
222        return state
223
224    def close_partition(self, partition: Partition) -> None:
225        # Attempt to retrieve the stream slice
226        stream_slice: Optional[StreamSlice] = partition.to_slice()  # type: ignore[assignment]
227
228        # Ensure stream_slice is not None
229        if stream_slice is None:
230            raise ValueError("stream_slice cannot be None")
231
232        partition_key = self._to_partition_key(stream_slice.partition)
233        with self._lock:
234            self._semaphore_per_partition[partition_key].acquire()
235            if not self._use_global_cursor:
236                cursor = self._cursor_per_partition[partition_key]
237                cursor.close_partition(partition=partition)
238                if (
239                    partition_key in self._partitions_done_generating_stream_slices
240                    and self._semaphore_per_partition[partition_key]._value == 0
241                ):
242                    self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
243
244            # Clean up the partition if it is fully processed
245            self._cleanup_if_done(partition_key)
246
247            self._check_and_update_parent_state()
248
249            self._emit_state_message()
250
251    def _check_and_update_parent_state(self) -> None:
252        last_closed_state = None
253
254        while self._partition_parent_state_map:
255            earliest_key, (candidate_state, candidate_seq) = next(
256                iter(self._partition_parent_state_map.items())
257            )
258
259            # if any partition that started <= candidate_seq is still open, we must wait
260            if (
261                self._processing_partitions_indexes
262                and self._processing_partitions_indexes[0] <= candidate_seq
263            ):
264                break
265
266            # safe to pop
267            self._partition_parent_state_map.popitem(last=False)
268            last_closed_state = candidate_state
269
270        if last_closed_state is not None:
271            self._parent_state = last_closed_state
272
273    def ensure_at_least_one_state_emitted(self) -> None:
274        """
275        The platform expects at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
276        called.
277        """
278        if not any(
279            semaphore_item[1]._value for semaphore_item in self._semaphore_per_partition.items()
280        ):
281            if self._synced_some_data:
282                # we only update those if we actually synced some data
283                self._global_cursor = self._new_global_cursor
284                self._lookback_window = self._timer.finish()
285            self._parent_state = self._partition_router.get_stream_state()
286        self._emit_state_message(throttle=False)
287
288    def _throttle_state_message(self) -> Optional[float]:
289        """
290        Throttles the state message emission to once every 600 seconds.
291        """
292        current_time = time.time()
293        if current_time - self._last_emission_time <= 600:
294            return None
295        return current_time
296
297    def _emit_state_message(self, throttle: bool = True) -> None:
298        if throttle:
299            current_time = self._throttle_state_message()
300            if current_time is None:
301                return
302            self._last_emission_time = current_time
303            # Skip state emit for global cursor if parent state is empty
304            if self._use_global_cursor and not self._parent_state:
305                return
306
307        self._connector_state_manager.update_state_for_stream(
308            self._stream_name,
309            self._stream_namespace,
310            self.state,
311        )
312        state_message = self._connector_state_manager.create_state_message(
313            self._stream_name, self._stream_namespace
314        )
315        self._message_repository.emit_message(state_message)
316
317    def stream_slices(self) -> Iterable[StreamSlice]:
318        if self._timer.is_running():
319            raise RuntimeError("stream_slices has been executed more than once.")
320
321        slices = self._partition_router.stream_slices()
322        self._timer.start()
323        for partition, last, parent_state in iterate_with_last_flag_and_state(
324            slices, self._partition_router.get_stream_state
325        ):
326            yield from self._generate_slices_from_partition(partition, parent_state)
327
328    def _generate_slices_from_partition(
329        self, partition: StreamSlice, parent_state: Mapping[str, Any]
330    ) -> Iterable[StreamSlice]:
331        # Ensure the maximum number of partitions is not exceeded
332        self._ensure_partition_limit()
333
334        partition_key = self._to_partition_key(partition.partition)
335
336        cursor = self._cursor_per_partition.get(self._to_partition_key(partition.partition))
337        if not cursor:
338            cursor = self._create_cursor(
339                self._global_cursor,
340                self._lookback_window if self._global_cursor else 0,
341            )
342            with self._lock:
343                self._number_of_partitions += 1
344                self._cursor_per_partition[partition_key] = cursor
345
346        if partition_key in self._semaphore_per_partition:
347            if not self._IS_PARTITION_DUPLICATION_LOGGED:
348                logger.warning(f"Partition duplication detected for stream {self._stream_name}")
349                self._IS_PARTITION_DUPLICATION_LOGGED = True
350            return
351        else:
352            self._semaphore_per_partition[partition_key] = threading.Semaphore(0)
353
354        with self._lock:
355            seq = self._generated_partitions_count
356            self._generated_partitions_count += 1
357            self._processing_partitions_indexes.append(seq)
358            self._partition_key_to_index[partition_key] = seq
359
360            if (
361                len(self._partition_parent_state_map) == 0
362                or self._partition_parent_state_map[
363                    next(reversed(self._partition_parent_state_map))
364                ][self._PARENT_STATE]
365                != parent_state
366            ):
367                self._partition_parent_state_map[partition_key] = (deepcopy(parent_state), seq)
368
369        for cursor_slice, is_last_slice, _ in iterate_with_last_flag_and_state(
370            cursor.stream_slices(),
371            lambda: None,
372        ):
373            self._semaphore_per_partition[partition_key].release()
374            if is_last_slice:
375                self._partitions_done_generating_stream_slices.add(partition_key)
376            yield StreamSlice(
377                partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
378            )
379
380    def _ensure_partition_limit(self) -> None:
381        """
382        Ensure the maximum number of partitions does not exceed the predefined limit.
383
384        Steps:
385        1. Attempt to remove partitions that are marked as finished in `_finished_partitions`.
386           These partitions are considered processed and safe to delete.
387        2. If the limit is still exceeded and no finished partitions are available for removal,
388           remove the oldest partition unconditionally. We expect failed partitions to be removed.
389
390        Logging:
391        - Logs a warning each time a partition is removed, indicating whether it was finished
392          or removed due to being the oldest.
393        """
394        if not self._use_global_cursor and self.limit_reached():
395            logger.info(
396                f"Exceeded the 'SWITCH_TO_GLOBAL_LIMIT' of {self.SWITCH_TO_GLOBAL_LIMIT}. "
397                f"Switching to global cursor for {self._stream_name}."
398            )
399            self._use_global_cursor = True
400
401        with self._lock:
402            while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
403                # Try removing finished partitions first
404                for partition_key in list(self._cursor_per_partition.keys()):
405                    if partition_key not in self._partition_key_to_index:
406                        oldest_partition = self._cursor_per_partition.pop(
407                            partition_key
408                        )  # Remove the oldest partition
409                        logger.debug(
410                            f"The maximum number of partitions has been reached. Dropping the oldest finished partition: {oldest_partition}. Over limit: {self._number_of_partitions - self.DEFAULT_MAX_PARTITIONS_NUMBER}."
411                        )
412                        break
413                else:
414                    # If no finished partitions can be removed, fall back to removing the oldest partition
415                    oldest_partition = self._cursor_per_partition.popitem(last=False)[
416                        1
417                    ]  # Remove the oldest partition
418                    logger.warning(
419                        f"The maximum number of partitions has been reached. Dropping the oldest partition: {oldest_partition}. Over limit: {self._number_of_partitions - self.DEFAULT_MAX_PARTITIONS_NUMBER}."
420                    )
421
422    def _set_initial_state(self, stream_state: StreamState) -> None:
423        """
424        Initialize the cursor's state using the provided `stream_state`.
425
426        This method supports global and per-partition state initialization.
427
428        - **Global State**: If `states` is missing, the `state` is treated as global and applied to all partitions.
429          The `global state` holds a single cursor position representing the latest processed record across all partitions.
430
431        - **Lookback Window**: Configured via `lookback_window`, it defines the period (in seconds) for reprocessing records.
432          This ensures robustness in case of upstream data delays or reordering. If not specified, it defaults to 0.
433
434        - **Per-Partition State**: If `states` is present, each partition's cursor state is initialized separately.
435
436        - **Parent State**: (if available) Used to initialize partition routers based on parent streams.
437
438        Args:
439            stream_state (StreamState): The state of the streams to be set. The format of the stream state should be:
440                {
441                    "states": [
442                        {
443                            "partition": {
444                                "partition_key": "value"
445                            },
446                            "cursor": {
447                                "last_updated": "2023-05-27T00:00:00Z"
448                            }
449                        }
450                    ],
451                    "state": {
452                        "last_updated": "2023-05-27T00:00:00Z"
453                    },
454                    lookback_window: 10,
455                    "parent_state": {
456                        "parent_stream_name": {
457                            "last_updated": "2023-05-27T00:00:00Z"
458                        }
459                    }
460                }
461        """
462        if not stream_state:
463            return
464
465        if (
466            self._PERPARTITION_STATE_KEY not in stream_state
467            and self._GLOBAL_STATE_KEY not in stream_state
468        ):
469            # We assume that `stream_state` is in a global format that can be applied to all partitions.
470            # Example: {"global_state_format_key": "global_state_format_value"}
471            self._set_global_state(stream_state)
472
473        else:
474            self._use_global_cursor = stream_state.get("use_global_cursor", False)
475
476            self._lookback_window = int(stream_state.get("lookback_window", 0))
477
478            for state in stream_state.get(self._PERPARTITION_STATE_KEY, []):
479                self._number_of_partitions += 1
480                self._cursor_per_partition[self._to_partition_key(state["partition"])] = (
481                    self._create_cursor(state["cursor"])
482                )
483
484            # set default state for missing partitions if it is per partition with fallback to global
485            if self._GLOBAL_STATE_KEY in stream_state:
486                self._set_global_state(stream_state[self._GLOBAL_STATE_KEY])
487
488        # Set initial parent state
489        if stream_state.get("parent_state"):
490            self._parent_state = stream_state["parent_state"]
491
492    def _set_global_state(self, stream_state: Mapping[str, Any]) -> None:
493        """
494        Initializes the global cursor state from the provided stream state.
495
496        If the cursor field key is present in the stream state, its value is parsed,
497        formatted, and stored as the global cursor. This ensures consistency in state
498        representation across partitions.
499        """
500        if self.cursor_field.cursor_field_key in stream_state:
501            global_state_value = stream_state[self.cursor_field.cursor_field_key]
502            final_format_global_state_value = self._connector_state_converter.output_format(
503                self._connector_state_converter.parse_value(global_state_value)
504            )
505
506            fixed_global_state = {
507                self.cursor_field.cursor_field_key: final_format_global_state_value
508            }
509
510            self._global_cursor = deepcopy(fixed_global_state)
511            self._new_global_cursor = deepcopy(fixed_global_state)
512
513    def observe(self, record: Record) -> None:
514        if not record.associated_slice:
515            raise ValueError(
516                "Invalid state as stream slices that are emitted should refer to an existing cursor"
517            )
518
519        # if the current record has no cursor value, we cannot meaningfully update the state based on it, so there is nothing more to do
520        try:
521            record_cursor_value = self._cursor_field.extract_value(record)
522        except ValueError:
523            return
524
525        try:
526            record_cursor = self._connector_state_converter.output_format(
527                self._connector_state_converter.parse_value(record_cursor_value)
528            )
529        except ValueError as exception:
530            if not self._logged_regarding_datetime_format_error:
531                logger.warning(
532                    "Skipping cursor update for stream '%s': failed to parse cursor field '%s' value %r: %s",
533                    self._stream_name,
534                    self._cursor_field.cursor_field_key,
535                    record_cursor_value,
536                    exception,
537                )
538                self._logged_regarding_datetime_format_error = True
539            return
540
541        self._synced_some_data = True
542        self._update_global_cursor(record_cursor)
543        if not self._use_global_cursor:
544            self._cursor_per_partition[
545                self._to_partition_key(record.associated_slice.partition)
546            ].observe(record)
547
548    def _update_global_cursor(self, value: Any) -> None:
549        if (
550            self._new_global_cursor is None
551            or self._new_global_cursor[self.cursor_field.cursor_field_key] < value
552        ):
553            self._new_global_cursor = {self.cursor_field.cursor_field_key: copy.deepcopy(value)}
554
555    def _cleanup_if_done(self, partition_key: str) -> None:
556        """
557        Free every in-memory structure that belonged to a completed partition:
558        cursor, semaphore, flag inside `_finished_partitions`
559        """
560        if not (
561            partition_key in self._partitions_done_generating_stream_slices
562            and self._semaphore_per_partition[partition_key]._value == 0
563        ):
564            return
565
566        self._semaphore_per_partition.pop(partition_key, None)
567        self._partitions_done_generating_stream_slices.discard(partition_key)
568
569        seq = self._partition_key_to_index.pop(partition_key)
570        self._processing_partitions_indexes.remove(seq)
571
572        logger.debug(f"Partition {partition_key} fully processed and cleaned up.")
573
574    def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
575        return self._partition_serializer.to_partition_key(partition)
576
577    def _to_dict(self, partition_key: str) -> Mapping[str, Any]:
578        return self._partition_serializer.to_partition(partition_key)
579
580    def _create_cursor(
581        self, cursor_state: Any, runtime_lookback_window: int = 0
582    ) -> ConcurrentCursor:
583        cursor = self._cursor_factory.create(
584            stream_state=deepcopy(cursor_state),
585            runtime_lookback_window=timedelta(seconds=runtime_lookback_window),
586        )
587        return cursor
588
589    def should_be_synced(self, record: Record) -> bool:
590        return self._get_cursor(record).should_be_synced(record)
591
592    def _get_cursor(self, record: Record) -> ConcurrentCursor:
593        if not record.associated_slice:
594            raise ValueError(
595                "Invalid state as stream slices that are emitted should refer to an existing cursor"
596            )
597
598        if self._use_global_cursor:
599            return self._create_cursor(
600                self._global_cursor,
601                self._lookback_window if self._global_cursor else 0,
602            )
603
604        partition_key = self._to_partition_key(record.associated_slice.partition)
605        if (
606            partition_key not in self._cursor_per_partition
607            and not self._attempt_to_create_cursor_if_not_provided
608        ):
609            raise ValueError(
610                "Invalid state as stream slices that are emitted should refer to an existing cursor"
611            )
612        elif partition_key not in self._cursor_per_partition:
613            return self._create_cursor(
614                self._global_cursor,
615                self._lookback_window if self._global_cursor else 0,
616            )
617        else:
618            return self._cursor_per_partition[partition_key]
619
620    def limit_reached(self) -> bool:
621        return self._number_of_partitions > self.SWITCH_TO_GLOBAL_LIMIT
622
623    @staticmethod
624    def get_parent_state(
625        stream_state: Optional[StreamState], parent_stream_name: str
626    ) -> Optional[AirbyteStateMessage]:
627        if not stream_state:
628            return None
629
630        if "parent_state" not in stream_state:
631            logger.warning(
632                f"Trying to get_parent_state for stream `{parent_stream_name}` when there are not parent state in the state"
633            )
634            return None
635        elif parent_stream_name not in stream_state["parent_state"]:
636            logger.info(
637                f"Could not find parent state for stream `{parent_stream_name}`. On parents available are {list(stream_state['parent_state'].keys())}"
638            )
639            return None
640
641        return AirbyteStateMessage(
642            type=AirbyteStateType.STREAM,
643            stream=AirbyteStreamState(
644                stream_descriptor=StreamDescriptor(parent_stream_name, None),
645                stream_state=AirbyteStateBlob(stream_state["parent_state"][parent_stream_name]),
646            ),
647        )
648
649    @staticmethod
650    def get_global_state(
651        stream_state: Optional[StreamState], parent_stream_name: str
652    ) -> Optional[AirbyteStateMessage]:
653        return (
654            AirbyteStateMessage(
655                type=AirbyteStateType.STREAM,
656                stream=AirbyteStreamState(
657                    stream_descriptor=StreamDescriptor(parent_stream_name, None),
658                    stream_state=AirbyteStateBlob(stream_state["state"]),
659                ),
660            )
661            if stream_state and "state" in stream_state
662            else None
663        )
664
665    def get_cursor_datetime_from_state(
666        self, stream_state: Mapping[str, Any]
667    ) -> datetime.datetime | None:
668        """Extract and parse the cursor datetime from the global cursor in per-partition state.
669
670        For per-partition cursors, the global cursor is stored under the "state" key.
671        This method delegates to the underlying cursor factory to parse the datetime.
672
673        Returns None if the global cursor is not present or cannot be parsed.
674        """
675        global_state = stream_state.get(self._GLOBAL_STATE_KEY)
676        if not global_state or not isinstance(global_state, dict):
677            return None
678
679        # Create a cursor to delegate the parsing
680        cursor = self._cursor_factory.create(stream_state={}, runtime_lookback_window=None)
681        return cursor.get_cursor_datetime_from_state(global_state)
682
683    def split_request_window(
684        self, stream_slice: StreamSlice, min_split_window: Optional[timedelta] = None
685    ) -> Optional[List[StreamSlice]]:
686        """
687        Split `stream_slice` in half the way any of this stream's per-partition cursors would.
688
689        `ConcurrentCursor.split_request_window` only reads the slice's boundaries and the cursor's granularity
690        and format, never the cursor's state, and every per-partition cursor is built by the same
691        `_cursor_factory`. So one state-free cursor can split the slices of every partition, including after
692        the switch to a global cursor, when the partition's own cursor may already have been evicted. The
693        children keep `stream_slice`'s `partition` and `extra_fields`, and `SimpleRetriever` re-associates the
694        records read from them with the original slice, so `observe` and `close_partition` are unaffected.
695
696        As with an unpartitioned stream, the children are not clamped again to the cursor's start or end.
697        """
698        return self._get_request_window_splitting_cursor().split_request_window(
699            stream_slice, min_split_window
700        )
701
702    def _get_request_window_splitting_cursor(self) -> ConcurrentCursor:
703        # Built on first use so that streams that never split a window do not pay for an extra cursor.
704        with self._lock:
705            if self._request_window_splitting_cursor is None:
706                self._request_window_splitting_cursor = self._cursor_factory.create(
707                    stream_state={}, runtime_lookback_window=None
708                )
709            return self._request_window_splitting_cursor