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._finished_generating_slices = False
196        self._logged_regarding_datetime_format_error = False
197
198    @property
199    def cursor_field(self) -> CursorField:
200        return self._cursor_field
201
202    @property
203    def partition_router(self) -> PartitionRouter:
204        return self._partition_router
205
206    @property
207    def state(self) -> MutableMapping[str, Any]:
208        state: dict[str, Any] = {"use_global_cursor": self._use_global_cursor}
209        if not self._use_global_cursor:
210            states = []
211            for partition_tuple, cursor in self._cursor_per_partition.items():
212                if cursor.state:
213                    states.append(
214                        {
215                            "partition": self._to_dict(partition_tuple),
216                            "cursor": copy.deepcopy(cursor.state),
217                        }
218                    )
219            state[self._PERPARTITION_STATE_KEY] = states
220
221        if self._global_cursor:
222            state[self._GLOBAL_STATE_KEY] = self._global_cursor
223        if self._lookback_window is not None:
224            state["lookback_window"] = self._lookback_window
225        if self._parent_state is not None:
226            state["parent_state"] = self._parent_state
227        return state
228
229    def close_partition(self, partition: Partition) -> None:
230        # Attempt to retrieve the stream slice
231        stream_slice: Optional[StreamSlice] = partition.to_slice()  # type: ignore[assignment]
232
233        # Ensure stream_slice is not None
234        if stream_slice is None:
235            raise ValueError("stream_slice cannot be None")
236
237        partition_key = self._to_partition_key(stream_slice.partition)
238        with self._lock:
239            self._semaphore_per_partition[partition_key].acquire()
240            if not self._use_global_cursor:
241                cursor = self._cursor_per_partition[partition_key]
242                cursor.close_partition(partition=partition)
243                if (
244                    partition_key in self._partitions_done_generating_stream_slices
245                    and self._semaphore_per_partition[partition_key]._value == 0
246                ):
247                    self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
248
249            # Clean up the partition if it is fully processed
250            self._cleanup_if_done(partition_key)
251
252            self._check_and_update_parent_state()
253
254            self._emit_state_message()
255
256    def _check_and_update_parent_state(self) -> None:
257        last_closed_state = None
258
259        while self._partition_parent_state_map:
260            earliest_key, (candidate_state, candidate_seq) = next(
261                iter(self._partition_parent_state_map.items())
262            )
263
264            # if any partition that started <= candidate_seq is still open, we must wait
265            if (
266                self._processing_partitions_indexes
267                and self._processing_partitions_indexes[0] <= candidate_seq
268            ):
269                break
270
271            # safe to pop
272            self._partition_parent_state_map.popitem(last=False)
273            last_closed_state = candidate_state
274
275        if last_closed_state is not None:
276            self._parent_state = last_closed_state
277
278    def ensure_at_least_one_state_emitted(self) -> None:
279        """
280        The platform expects at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
281        called.
282
283        If the stream did not complete (a partition was never closed, or slice generation stopped early), the global cursor is not
284        advanced, because partitions that were never generated have not been read up to it. If the state then holds no
285        checkpoint, nothing is emitted. Such a state records no progress, but it is not empty, so the next attempt would treat
286        the stream as already synced once: a StateDelegatingStream would switch to its incremental stream and never read the
287        slices this attempt missed. This covers global-cursor streams; in per-partition mode the cursor of every generated
288        partition reports its start value, so the state is still emitted.
289
290        If slice generation stopped early, the parent state is not refreshed from the partition router: the router may already have closed a parent slice whose last child slice was never generated, so only the parent state recorded by `_check_and_update_parent_state` for closed child partitions is kept.
291        """
292        all_partitions_closed = not any(
293            semaphore_item[1]._value for semaphore_item in self._semaphore_per_partition.items()
294        )
295        stream_completed = all_partitions_closed and self._finished_generating_slices
296        if all_partitions_closed:
297            if self._synced_some_data and stream_completed:
298                # we only update those if we actually synced some data
299                self._global_cursor = self._new_global_cursor
300                self._lookback_window = self._timer.finish()
301            if stream_completed:
302                self._parent_state = self._partition_router.get_stream_state()
303        if not stream_completed and not self._has_checkpoint():
304            return
305        self._emit_state_message(throttle=False)
306
307    def _has_checkpoint(self) -> bool:
308        state = self.state
309        return bool(
310            state.get(self._GLOBAL_STATE_KEY)
311            or state.get(self._PERPARTITION_STATE_KEY)
312            or state.get("parent_state")
313        )
314
315    def _throttle_state_message(self) -> Optional[float]:
316        """
317        Throttles the state message emission to once every 600 seconds.
318        """
319        current_time = time.time()
320        if current_time - self._last_emission_time <= 600:
321            return None
322        return current_time
323
324    def _emit_state_message(self, throttle: bool = True) -> None:
325        if throttle:
326            current_time = self._throttle_state_message()
327            if current_time is None:
328                return
329            self._last_emission_time = current_time
330            # Skip state emit for global cursor if parent state is empty
331            if self._use_global_cursor and not self._parent_state:
332                return
333
334        self._connector_state_manager.update_state_for_stream(
335            self._stream_name,
336            self._stream_namespace,
337            self.state,
338        )
339        state_message = self._connector_state_manager.create_state_message(
340            self._stream_name, self._stream_namespace
341        )
342        self._message_repository.emit_message(state_message)
343
344    def stream_slices(self) -> Iterable[StreamSlice]:
345        if self._timer.is_running():
346            raise RuntimeError("stream_slices has been executed more than once.")
347
348        slices = self._partition_router.stream_slices()
349        self._timer.start()
350        for partition, last, parent_state in iterate_with_last_flag_and_state(
351            slices, self._partition_router.get_stream_state
352        ):
353            yield from self._generate_slices_from_partition(partition, parent_state)
354        self._finished_generating_slices = True
355
356    def _generate_slices_from_partition(
357        self, partition: StreamSlice, parent_state: Mapping[str, Any]
358    ) -> Iterable[StreamSlice]:
359        # Ensure the maximum number of partitions is not exceeded
360        self._ensure_partition_limit()
361
362        partition_key = self._to_partition_key(partition.partition)
363
364        cursor = self._cursor_per_partition.get(self._to_partition_key(partition.partition))
365        if not cursor:
366            cursor = self._create_cursor(
367                self._global_cursor,
368                self._lookback_window if self._global_cursor else 0,
369            )
370            with self._lock:
371                self._number_of_partitions += 1
372                self._cursor_per_partition[partition_key] = cursor
373
374        if partition_key in self._semaphore_per_partition:
375            if not self._IS_PARTITION_DUPLICATION_LOGGED:
376                logger.warning(f"Partition duplication detected for stream {self._stream_name}")
377                self._IS_PARTITION_DUPLICATION_LOGGED = True
378            return
379        else:
380            self._semaphore_per_partition[partition_key] = threading.Semaphore(0)
381
382        with self._lock:
383            seq = self._generated_partitions_count
384            self._generated_partitions_count += 1
385            self._processing_partitions_indexes.append(seq)
386            self._partition_key_to_index[partition_key] = seq
387
388            if (
389                len(self._partition_parent_state_map) == 0
390                or self._partition_parent_state_map[
391                    next(reversed(self._partition_parent_state_map))
392                ][self._PARENT_STATE]
393                != parent_state
394            ):
395                self._partition_parent_state_map[partition_key] = (deepcopy(parent_state), seq)
396
397        for cursor_slice, is_last_slice, _ in iterate_with_last_flag_and_state(
398            cursor.stream_slices(),
399            lambda: None,
400        ):
401            self._semaphore_per_partition[partition_key].release()
402            if is_last_slice:
403                self._partitions_done_generating_stream_slices.add(partition_key)
404            yield StreamSlice(
405                partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
406            )
407
408    def _ensure_partition_limit(self) -> None:
409        """
410        Ensure the maximum number of partitions does not exceed the predefined limit.
411
412        Steps:
413        1. Attempt to remove partitions that are marked as finished in `_finished_partitions`.
414           These partitions are considered processed and safe to delete.
415        2. If the limit is still exceeded and no finished partitions are available for removal,
416           remove the oldest partition unconditionally. We expect failed partitions to be removed.
417
418        Logging:
419        - Logs a warning each time a partition is removed, indicating whether it was finished
420          or removed due to being the oldest.
421        """
422        if not self._use_global_cursor and self.limit_reached():
423            logger.info(
424                f"Exceeded the 'SWITCH_TO_GLOBAL_LIMIT' of {self.SWITCH_TO_GLOBAL_LIMIT}. "
425                f"Switching to global cursor for {self._stream_name}."
426            )
427            self._use_global_cursor = True
428
429        with self._lock:
430            while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
431                # Try removing finished partitions first
432                for partition_key in list(self._cursor_per_partition.keys()):
433                    if partition_key not in self._partition_key_to_index:
434                        oldest_partition = self._cursor_per_partition.pop(
435                            partition_key
436                        )  # Remove the oldest partition
437                        logger.debug(
438                            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}."
439                        )
440                        break
441                else:
442                    # If no finished partitions can be removed, fall back to removing the oldest partition
443                    oldest_partition = self._cursor_per_partition.popitem(last=False)[
444                        1
445                    ]  # Remove the oldest partition
446                    logger.warning(
447                        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}."
448                    )
449
450    def _set_initial_state(self, stream_state: StreamState) -> None:
451        """
452        Initialize the cursor's state using the provided `stream_state`.
453
454        This method supports global and per-partition state initialization.
455
456        - **Global State**: If `states` is missing, the `state` is treated as global and applied to all partitions.
457          The `global state` holds a single cursor position representing the latest processed record across all partitions.
458
459        - **Lookback Window**: Configured via `lookback_window`, it defines the period (in seconds) for reprocessing records.
460          This ensures robustness in case of upstream data delays or reordering. If not specified, it defaults to 0.
461
462        - **Per-Partition State**: If `states` is present, each partition's cursor state is initialized separately.
463
464        - **Parent State**: (if available) Used to initialize partition routers based on parent streams.
465
466        Args:
467            stream_state (StreamState): The state of the streams to be set. The format of the stream state should be:
468                {
469                    "states": [
470                        {
471                            "partition": {
472                                "partition_key": "value"
473                            },
474                            "cursor": {
475                                "last_updated": "2023-05-27T00:00:00Z"
476                            }
477                        }
478                    ],
479                    "state": {
480                        "last_updated": "2023-05-27T00:00:00Z"
481                    },
482                    lookback_window: 10,
483                    "parent_state": {
484                        "parent_stream_name": {
485                            "last_updated": "2023-05-27T00:00:00Z"
486                        }
487                    }
488                }
489        """
490        if not stream_state:
491            return
492
493        if (
494            self._PERPARTITION_STATE_KEY not in stream_state
495            and self._GLOBAL_STATE_KEY not in stream_state
496        ):
497            # We assume that `stream_state` is in a global format that can be applied to all partitions.
498            # Example: {"global_state_format_key": "global_state_format_value"}
499            self._set_global_state(stream_state)
500
501        else:
502            self._use_global_cursor = stream_state.get("use_global_cursor", False)
503
504            self._lookback_window = int(stream_state.get("lookback_window", 0))
505
506            for state in stream_state.get(self._PERPARTITION_STATE_KEY, []):
507                self._number_of_partitions += 1
508                self._cursor_per_partition[self._to_partition_key(state["partition"])] = (
509                    self._create_cursor(state["cursor"])
510                )
511
512            # set default state for missing partitions if it is per partition with fallback to global
513            if self._GLOBAL_STATE_KEY in stream_state:
514                self._set_global_state(stream_state[self._GLOBAL_STATE_KEY])
515
516        # Set initial parent state
517        if stream_state.get("parent_state"):
518            self._parent_state = stream_state["parent_state"]
519
520    def _set_global_state(self, stream_state: Mapping[str, Any]) -> None:
521        """
522        Initializes the global cursor state from the provided stream state.
523
524        If the cursor field key is present in the stream state, its value is parsed,
525        formatted, and stored as the global cursor. This ensures consistency in state
526        representation across partitions.
527        """
528        if self.cursor_field.cursor_field_key in stream_state:
529            global_state_value = stream_state[self.cursor_field.cursor_field_key]
530            final_format_global_state_value = self._connector_state_converter.output_format(
531                self._connector_state_converter.parse_value(global_state_value)
532            )
533
534            fixed_global_state = {
535                self.cursor_field.cursor_field_key: final_format_global_state_value
536            }
537
538            self._global_cursor = deepcopy(fixed_global_state)
539            self._new_global_cursor = deepcopy(fixed_global_state)
540
541    def observe(self, record: Record) -> None:
542        if not record.associated_slice:
543            raise ValueError(
544                "Invalid state as stream slices that are emitted should refer to an existing cursor"
545            )
546
547        # if the current record has no cursor value, we cannot meaningfully update the state based on it, so there is nothing more to do
548        try:
549            record_cursor_value = self._cursor_field.extract_value(record)
550        except ValueError:
551            return
552
553        try:
554            record_cursor = self._connector_state_converter.output_format(
555                self._connector_state_converter.parse_value(record_cursor_value)
556            )
557        except ValueError as exception:
558            if not self._logged_regarding_datetime_format_error:
559                logger.warning(
560                    "Skipping cursor update for stream '%s': failed to parse cursor field '%s' value %r: %s",
561                    self._stream_name,
562                    self._cursor_field.cursor_field_key,
563                    record_cursor_value,
564                    exception,
565                )
566                self._logged_regarding_datetime_format_error = True
567            return
568
569        self._synced_some_data = True
570        self._update_global_cursor(record_cursor)
571        if not self._use_global_cursor:
572            self._cursor_per_partition[
573                self._to_partition_key(record.associated_slice.partition)
574            ].observe(record)
575
576    def _update_global_cursor(self, value: Any) -> None:
577        if (
578            self._new_global_cursor is None
579            or self._new_global_cursor[self.cursor_field.cursor_field_key] < value
580        ):
581            self._new_global_cursor = {self.cursor_field.cursor_field_key: copy.deepcopy(value)}
582
583    def _cleanup_if_done(self, partition_key: str) -> None:
584        """
585        Free every in-memory structure that belonged to a completed partition:
586        cursor, semaphore, flag inside `_finished_partitions`
587        """
588        if not (
589            partition_key in self._partitions_done_generating_stream_slices
590            and self._semaphore_per_partition[partition_key]._value == 0
591        ):
592            return
593
594        self._semaphore_per_partition.pop(partition_key, None)
595        self._partitions_done_generating_stream_slices.discard(partition_key)
596
597        seq = self._partition_key_to_index.pop(partition_key)
598        self._processing_partitions_indexes.remove(seq)
599
600        logger.debug(f"Partition {partition_key} fully processed and cleaned up.")
601
602    def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
603        return self._partition_serializer.to_partition_key(partition)
604
605    def _to_dict(self, partition_key: str) -> Mapping[str, Any]:
606        return self._partition_serializer.to_partition(partition_key)
607
608    def _create_cursor(
609        self, cursor_state: Any, runtime_lookback_window: int = 0
610    ) -> ConcurrentCursor:
611        cursor = self._cursor_factory.create(
612            stream_state=deepcopy(cursor_state),
613            runtime_lookback_window=timedelta(seconds=runtime_lookback_window),
614        )
615        return cursor
616
617    def should_be_synced(self, record: Record) -> bool:
618        return self._get_cursor(record).should_be_synced(record)
619
620    def _get_cursor(self, record: Record) -> ConcurrentCursor:
621        if not record.associated_slice:
622            raise ValueError(
623                "Invalid state as stream slices that are emitted should refer to an existing cursor"
624            )
625
626        if self._use_global_cursor:
627            return self._create_cursor(
628                self._global_cursor,
629                self._lookback_window if self._global_cursor else 0,
630            )
631
632        partition_key = self._to_partition_key(record.associated_slice.partition)
633        if (
634            partition_key not in self._cursor_per_partition
635            and not self._attempt_to_create_cursor_if_not_provided
636        ):
637            raise ValueError(
638                "Invalid state as stream slices that are emitted should refer to an existing cursor"
639            )
640        elif partition_key not in self._cursor_per_partition:
641            return self._create_cursor(
642                self._global_cursor,
643                self._lookback_window if self._global_cursor else 0,
644            )
645        else:
646            return self._cursor_per_partition[partition_key]
647
648    def limit_reached(self) -> bool:
649        return self._number_of_partitions > self.SWITCH_TO_GLOBAL_LIMIT
650
651    @staticmethod
652    def get_parent_state(
653        stream_state: Optional[StreamState], parent_stream_name: str
654    ) -> Optional[AirbyteStateMessage]:
655        if not stream_state:
656            return None
657
658        if "parent_state" not in stream_state:
659            logger.warning(
660                f"Trying to get_parent_state for stream `{parent_stream_name}` when there are not parent state in the state"
661            )
662            return None
663        elif parent_stream_name not in stream_state["parent_state"]:
664            logger.info(
665                f"Could not find parent state for stream `{parent_stream_name}`. On parents available are {list(stream_state['parent_state'].keys())}"
666            )
667            return None
668
669        return AirbyteStateMessage(
670            type=AirbyteStateType.STREAM,
671            stream=AirbyteStreamState(
672                stream_descriptor=StreamDescriptor(parent_stream_name, None),
673                stream_state=AirbyteStateBlob(stream_state["parent_state"][parent_stream_name]),
674            ),
675        )
676
677    @staticmethod
678    def get_global_state(
679        stream_state: Optional[StreamState], parent_stream_name: str
680    ) -> Optional[AirbyteStateMessage]:
681        return (
682            AirbyteStateMessage(
683                type=AirbyteStateType.STREAM,
684                stream=AirbyteStreamState(
685                    stream_descriptor=StreamDescriptor(parent_stream_name, None),
686                    stream_state=AirbyteStateBlob(stream_state["state"]),
687                ),
688            )
689            if stream_state and "state" in stream_state
690            else None
691        )
692
693    def get_cursor_datetime_from_state(
694        self, stream_state: Mapping[str, Any]
695    ) -> datetime.datetime | None:
696        """Extract and parse the cursor datetime from the global cursor in per-partition state.
697
698        For per-partition cursors, the global cursor is stored under the "state" key.
699        This method delegates to the underlying cursor factory to parse the datetime.
700
701        Returns None if the global cursor is not present or cannot be parsed.
702        """
703        global_state = stream_state.get(self._GLOBAL_STATE_KEY)
704        if not global_state or not isinstance(global_state, dict):
705            return None
706
707        # Create a cursor to delegate the parsing
708        cursor = self._cursor_factory.create(stream_state={}, runtime_lookback_window=None)
709        return cursor.get_cursor_datetime_from_state(global_state)
710
711    def split_request_window(
712        self, stream_slice: StreamSlice, min_split_window: Optional[timedelta] = None
713    ) -> Optional[List[StreamSlice]]:
714        """
715        Split `stream_slice` in half the way any of this stream's per-partition cursors would.
716
717        `ConcurrentCursor.split_request_window` only reads the slice's boundaries and the cursor's granularity
718        and format, never the cursor's state, and every per-partition cursor is built by the same
719        `_cursor_factory`. So one state-free cursor can split the slices of every partition, including after
720        the switch to a global cursor, when the partition's own cursor may already have been evicted. The
721        children keep `stream_slice`'s `partition` and `extra_fields`, and `SimpleRetriever` re-associates the
722        records read from them with the original slice, so `observe` and `close_partition` are unaffected.
723
724        As with an unpartitioned stream, the children are not clamped again to the cursor's start or end.
725        """
726        return self._get_request_window_splitting_cursor().split_request_window(
727            stream_slice, min_split_window
728        )
729
730    def _get_request_window_splitting_cursor(self) -> ConcurrentCursor:
731        # Built on first use so that streams that never split a window do not pay for an extra cursor.
732        with self._lock:
733            if self._request_window_splitting_cursor is None:
734                self._request_window_splitting_cursor = self._cursor_factory.create(
735                    stream_state={}, runtime_lookback_window=None
736                )
737            return self._request_window_splitting_cursor

 44class SyncTokenCursor(Cursor):
 45    """Track opaque API-issued sync tokens and commit them after a complete partition read.
 46
 47    Unpartitioned state is ``{"state_type": "sync_token", "token": "..."}``; partitioned
 48    state stores a ``states`` list pairing each partition with its token. Tokens become committed
 49    state only in ``close_partition``. Partitioned state is capped at 25,000 entries.
 50    """
 51
 52    STATE_TYPE = "sync_token"
 53
 54    def __init__(
 55        self,
 56        *,
 57        stream_name: str,
 58        stream_namespace: Optional[str],
 59        stream_state: Optional[StreamState],
 60        message_repository: MessageRepository,
 61        connector_state_manager: ConnectorStateManager,
 62        partition_router: PartitionRouter,
 63        token_path: list[str],
 64        decoder: Decoder,
 65        config: Config,
 66        parameters: Mapping[str, Any],
 67        completion_condition: Optional[str] = None,
 68        cursor_field: Optional[str] = None,
 69        max_partitions: int = ConcurrentPerPartitionCursor.DEFAULT_MAX_PARTITIONS_NUMBER,
 70    ) -> None:
 71        self._stream_name = stream_name
 72        self._stream_namespace = stream_namespace
 73        self._message_repository = message_repository
 74        self._connector_state_manager = connector_state_manager
 75        self._partition_router = partition_router
 76        self._partitioned = not isinstance(partition_router, SinglePartitionRouter)
 77        self._token_path = [
 78            InterpolatedString.create(path, parameters=parameters) for path in token_path
 79        ]
 80        self._decoder = decoder
 81        self._config = config
 82        self._completion_condition = (
 83            InterpolatedBoolean(completion_condition, parameters=parameters)
 84            if completion_condition is not None
 85            else None
 86        )
 87        self.cursor_field = CursorField(cursor_field) if cursor_field else None
 88        self._max_partitions = max_partitions
 89        self._serializer = PerPartitionKeySerializer()
 90        self._committed: OrderedDict[str, dict[str, Any]] = OrderedDict()
 91        self._pending: dict[str, dict[str, Any]] = {}
 92        self._page_counts: dict[str, int] = {}
 93        self._lock = threading.RLock()
 94        self._last_emission_time = 0.0
 95        self._logged_bootstrap_reasons: set[str] = set()
 96        self._state_is_foreign = False
 97        self._state_emitted = False
 98        self._warned_no_token = False
 99        self._warned_cap = False
100        self._partitions_over_cap: set[str] = set()
101        self._invalid_state_reason: Optional[str] = None
102        self._load_state(stream_state or {})
103
104    @property
105    def decoder(self) -> Decoder:
106        return self._decoder
107
108    @property
109    def partitioned(self) -> bool:
110        return self._partitioned
111
112    @property
113    def partition_router(self) -> PartitionRouter:
114        return self._partition_router
115
116    @property
117    def state(self) -> MutableMapping[str, Any]:
118        with self._lock:
119            if not self._partitioned:
120                state: dict[str, Any] = {"state_type": self.STATE_TYPE}
121                if self._committed:
122                    state.update(copy.deepcopy(next(iter(self._committed.values()))))
123                return state
124            return {
125                "state_type": self.STATE_TYPE,
126                "states": [
127                    {
128                        "partition": self._serializer.to_partition(partition_key),
129                        "cursor": copy.deepcopy(cursor),
130                    }
131                    for partition_key, cursor in self._committed.items()
132                ],
133            }
134
135    def stream_slices(self) -> Iterable[StreamSlice]:
136        if self._invalid_state_reason is not None:
137            message = (
138                f"Stream {self._stream_name}: the saved sync token state is malformed "
139                f"({self._invalid_state_reason}). Refresh the stream to start a new full sync."
140            )
141            raise AirbyteTracedException(
142                message=message, internal_message=message, failure_type=FailureType.system_error
143            )
144        seen_partitions: set[str] = set()
145        for stream_slice in self._partition_router.stream_slices():
146            partition = dict(stream_slice.partition or {})
147            if "token" in partition:
148                message = f"Stream {self._stream_name}: partition router uses reserved key 'token'"
149                raise AirbyteTracedException(
150                    message=message,
151                    internal_message=message,
152                    failure_type=FailureType.config_error,
153                )
154            partition_key = self._serializer.to_partition_key(partition)
155            if partition_key in seen_partitions:
156                continue
157            seen_partitions.add(partition_key)
158            with self._lock:
159                committed = self._committed.get(partition_key)
160                token = committed["token"] if committed else None
161                at_partition_cap = len(self._committed) >= self._max_partitions
162                if token:
163                    self._log_replay(partition, committed or {})
164                elif at_partition_cap:
165                    self._partitions_over_cap.add(partition_key)
166                    self._log_bootstrap("partition cap")
167                elif not self._state_is_foreign:
168                    self._log_bootstrap("none")
169            yield StreamSlice(
170                partition=partition,
171                cursor_slice={"token": token} if token else {},
172                extra_fields=stream_slice.extra_fields,
173            )
174
175        if self._partitioned:
176            with self._lock:
177                vanished = set(self._committed) - seen_partitions
178                for partition_key in vanished:
179                    del self._committed[partition_key]
180                if vanished:
181                    logger.info(
182                        "Stream %s: dropped the stored tokens of %s partitions the router no longer returns.",
183                        self._stream_name,
184                        len(vanished),
185                    )
186
187    def on_page_completed(self, page: PageCompleted) -> None:
188        partition_key = self._serializer.to_partition_key(page.stream_slice.partition or {})
189        with self._lock:
190            self._page_counts[partition_key] = self._page_counts.get(partition_key, 0) + 1
191
192        if not page.pagination_exhausted or page.response is None:
193            return
194
195        body = next(
196            (
197                decoded
198                for decoded in self._decoder.decode(page.response)
199                if isinstance(decoded, Mapping)
200            ),
201            {},
202        )
203        path = [str(part.eval(self._config)) for part in self._token_path]
204        token_value = dpath.get(body, path, default=None) if isinstance(body, Mapping) else None
205        if self._completion_condition is not None:
206            complete = self._completion_condition.eval(
207                self._config,
208                response=body,
209                headers=page.response.headers,
210                stream_slice=page.stream_slice,
211            )
212        else:
213            complete = True
214        if not complete:
215            self._warn_missing_token()
216            return
217
218        token = self._coerce_token(token_value, path)
219        if token is None:
220            self._warn_missing_token()
221            return
222
223        with self._lock:
224            self._pending[partition_key] = {
225                "token": token,
226                "committed_at": datetime.datetime.now(datetime.timezone.utc)
227                .isoformat(timespec="milliseconds")
228                .replace("+00:00", "Z"),
229            }
230
231    def _coerce_token(self, value: Any, path: list[str]) -> Optional[str]:
232        if value is None or value == "":
233            return None
234        if isinstance(value, bool):
235            value_type = "boolean"
236        elif isinstance(value, str):
237            return value
238        elif isinstance(value, int):
239            return str(value)
240        elif isinstance(value, float):
241            value_type = "float"
242        elif isinstance(value, Mapping):
243            value_type = "dict"
244        elif isinstance(value, list):
245            value_type = "list"
246        else:
247            value_type = type(value).__name__
248        message = (
249            f"Stream {self._stream_name}: token_path {json.dumps(path)} holds a {value_type}, "
250            "not a string or integer. Check token_path."
251        )
252        raise AirbyteTracedException(
253            internal_message=message,
254            message=message,
255            failure_type=FailureType.system_error,
256        )
257
258    def _log_bootstrap(self, reason: str) -> None:
259        if reason not in self._logged_bootstrap_reasons:
260            logger.info(
261                "Stream %s: no stored sync token (reason: %s); reading all records.",
262                self._stream_name,
263                reason,
264            )
265            self._logged_bootstrap_reasons.add(reason)
266
267    def _log_foreign_state(self) -> None:
268        self._state_is_foreign = True
269        logger.warning(
270            "Stream %s: the saved state belongs to another cursor type; reading all records once.",
271            self._stream_name,
272        )
273
274    def _log_replay(self, partition: Mapping[str, Any], committed: Mapping[str, Any]) -> None:
275        age = self._human_age(committed.get("committed_at"))
276        if self._partitioned:
277            logger.debug(
278                "Stream %s: replaying the sync token committed %s ago for partition %s.",
279                self._stream_name,
280                age,
281                dict(partition),
282            )
283        else:
284            logger.info(
285                "Stream %s: replaying the sync token committed %s ago.",
286                self._stream_name,
287                age,
288            )
289
290    @staticmethod
291    def _human_age(committed_at: Any) -> str:
292        if not isinstance(committed_at, str) or not committed_at:
293            return "an unknown time"
294        try:
295            committed_time = datetime.datetime.fromisoformat(committed_at.replace("Z", "+00:00"))
296            if committed_time.tzinfo is None:
297                committed_time = committed_time.replace(tzinfo=datetime.timezone.utc)
298            age_seconds = max(
299                0,
300                int(
301                    (datetime.datetime.now(datetime.timezone.utc) - committed_time).total_seconds()
302                ),
303            )
304        except (OverflowError, ValueError):
305            return "an unknown time"
306
307        units = (
308            ("day", 24 * 60 * 60),
309            ("hour", 60 * 60),
310            ("minute", 60),
311            ("second", 1),
312        )
313        parts = []
314        for unit, unit_seconds in units:
315            count, age_seconds = divmod(age_seconds, unit_seconds)
316            if count:
317                parts.append(f"{count} {unit}{'' if count == 1 else 's'}")
318            if len(parts) == 2:
319                break
320        return " ".join(parts) if parts else "0 seconds"
321
322    def _warn_missing_token(self) -> None:
323        if not self._warned_no_token:
324            logger.warning(
325                "Stream %s: pagination ended without a new sync token; kept the stored token. "
326                "The next sync re-reads from it.",
327                self._stream_name,
328            )
329            self._warned_no_token = True
330
331    def close_partition(self, partition: Partition) -> None:
332        stream_slice = partition.to_slice()
333        if stream_slice is None:
334            raise ValueError("stream_slice cannot be None")
335        partition_values = getattr(stream_slice, "partition", None)
336        if partition_values is None:
337            partition_values = stream_slice.get("partition", {})
338        partition_key = self._serializer.to_partition_key(partition_values)
339        with self._lock:
340            pending = self._pending.pop(partition_key, None)
341            page_count = self._page_counts.pop(partition_key, 0)
342            if pending is not None:
343                if partition_key in self._committed or len(self._committed) < self._max_partitions:
344                    self._committed[partition_key] = pending
345                    self._partitions_over_cap.discard(partition_key)
346                    logger.debug(
347                        "Stream %s: committed sync token %s for partition %s after %s pages.",
348                        self._stream_name,
349                        pending["token"],
350                        self._serializer.to_partition(partition_key),
351                        page_count,
352                    )
353                else:
354                    self._partitions_over_cap.add(partition_key)
355            if not self._partitioned:
356                self._emit_state_message()
357            elif time.time() - self._last_emission_time > 600:
358                self._emit_state_message()
359
360    def _emit_state_message(self) -> None:
361        self._last_emission_time = time.time()
362        self._connector_state_manager.update_state_for_stream(
363            self._stream_name, self._stream_namespace, self.state
364        )
365        self._message_repository.emit_message(
366            self._connector_state_manager.create_state_message(
367                self._stream_name, self._stream_namespace
368            )
369        )
370        self._state_emitted = True
371
372    def ensure_at_least_one_state_emitted(self) -> None:
373        if self._invalid_state_reason is not None:
374            return  # leave the malformed state in place for diagnosis
375        if self._partitioned or not self._state_emitted:
376            self._emit_state_message()
377        if self._partitions_over_cap and not self._warned_cap:
378            logger.warning(
379                "Stream %s: %s partitions already store tokens; %s more read in full on every sync.",
380                self._stream_name,
381                len(self._committed),
382                len(self._partitions_over_cap),
383            )
384            self._warned_cap = True
385
386    def observe(self, record: Record) -> None:
387        return
388
389    def should_be_synced(self, record: Record) -> bool:
390        return True
391
392    def _load_state(self, stream_state: StreamState) -> None:
393        if not stream_state:
394            self._log_bootstrap("none")
395            return
396        if stream_state.get("state_type") != self.STATE_TYPE:
397            self._log_foreign_state()
398            return
399        if self._partitioned:
400            states = stream_state.get("states")
401            if not isinstance(states, list) or "token" in stream_state:
402                if "token" not in stream_state and "states" not in stream_state:
403                    self._log_bootstrap("none")
404                else:
405                    self._invalid_state_reason = "unpartitioned shape on a partitioned stream"
406                return
407            committed: dict[str, dict[str, Any]] = {}
408            for state in states:
409                if not isinstance(state, Mapping):
410                    self._invalid_state_reason = "a partition entry is not an object"
411                    return
412                partition = state.get("partition")
413                cursor = state.get("cursor")
414                if (
415                    isinstance(partition, Mapping)
416                    and isinstance(cursor, Mapping)
417                    and isinstance(cursor.get("token"), str)
418                    and cursor.get("token")
419                ):
420                    committed[self._serializer.to_partition_key(partition)] = {
421                        "token": cursor["token"],
422                        "committed_at": cursor.get("committed_at", ""),
423                    }
424                else:
425                    self._invalid_state_reason = "a partition entry has no valid token"
426                    return
427            self._committed.update(committed)
428        else:
429            if "states" in stream_state:
430                self._invalid_state_reason = "partitioned shape on an unpartitioned stream"
431                return
432            token = stream_state.get("token")
433            if isinstance(token, str) and token:
434                self._committed["{}"] = {
435                    "token": token,
436                    "committed_at": stream_state.get("committed_at", ""),
437                }
438            elif "token" in stream_state:
439                self._invalid_state_reason = "the token is not a non-empty string"