airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5from abc import ABC, abstractmethod
  6from enum import Enum
  7from typing import TYPE_CHECKING, Any, Callable, List, MutableMapping, Optional, Tuple
  8
  9if TYPE_CHECKING:
 10    from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 11
 12
 13class ConcurrencyCompatibleStateType(Enum):
 14    date_range = "date-range"
 15    integer = "integer"
 16
 17
 18class AbstractStreamStateConverter(ABC):
 19    START_KEY = "start"
 20    END_KEY = "end"
 21    MOST_RECENT_RECORD_KEY = "most_recent_cursor_value"
 22
 23    @abstractmethod
 24    def _from_state_message(self, value: Any) -> Any:
 25        pass
 26
 27    @abstractmethod
 28    def _to_state_message(self, value: Any) -> Any:
 29        pass
 30
 31    def __init__(self, is_sequential_state: bool = True):
 32        self._is_sequential_state = is_sequential_state
 33
 34    def convert_to_state_message(
 35        self, cursor_field: "CursorField", stream_state: MutableMapping[str, Any]
 36    ) -> MutableMapping[str, Any]:
 37        """
 38        Convert the state message from the concurrency-compatible format to the stream's original format.
 39
 40        e.g.
 41        { "created": "2021-01-18T21:18:20.000Z" }
 42        """
 43        if self.is_state_message_compatible(stream_state) and self._is_sequential_state:
 44            legacy_state = stream_state.get("legacy", {})
 45            latest_complete_time = self._get_latest_complete_time(stream_state.get("slices", []))
 46            if latest_complete_time is not None:
 47                legacy_state.update(
 48                    {cursor_field.cursor_field_key: self._to_state_message(latest_complete_time)}
 49                )
 50            return legacy_state or {}
 51        else:
 52            return self.serialize(stream_state, ConcurrencyCompatibleStateType.date_range)
 53
 54    def _get_latest_complete_time(self, slices: List[MutableMapping[str, Any]]) -> Any:
 55        """
 56        Get the latest time before which all records have been processed.
 57        """
 58        if not slices:
 59            raise RuntimeError(
 60                "Expected at least one slice but there were none. This is unexpected; please contact Support."
 61            )
 62        merged_intervals = self.merge_intervals(slices)
 63        first_interval = merged_intervals[0]
 64
 65        return first_interval.get("most_recent_cursor_value") or first_interval[self.START_KEY]
 66
 67    def deserialize(self, state: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
 68        """
 69        Perform any transformations needed for compatibility with the converter.
 70        """
 71        for stream_slice in state.get("slices", []):
 72            stream_slice[self.START_KEY] = self._from_state_message(stream_slice[self.START_KEY])
 73            stream_slice[self.END_KEY] = self._from_state_message(stream_slice[self.END_KEY])
 74            if self.MOST_RECENT_RECORD_KEY in stream_slice:
 75                stream_slice[self.MOST_RECENT_RECORD_KEY] = self._from_state_message(
 76                    stream_slice[self.MOST_RECENT_RECORD_KEY]
 77                )
 78        return state
 79
 80    def serialize(
 81        self, state: MutableMapping[str, Any], state_type: ConcurrencyCompatibleStateType
 82    ) -> MutableMapping[str, Any]:
 83        """
 84        Perform any transformations needed for compatibility with the converter.
 85        """
 86        serialized_slices = []
 87        for stream_slice in state.get("slices", []):
 88            serialized_slice = {
 89                self.START_KEY: self._to_state_message(stream_slice[self.START_KEY]),
 90                self.END_KEY: self._to_state_message(stream_slice[self.END_KEY]),
 91            }
 92            if stream_slice.get(self.MOST_RECENT_RECORD_KEY):
 93                serialized_slice[self.MOST_RECENT_RECORD_KEY] = self._to_state_message(
 94                    stream_slice[self.MOST_RECENT_RECORD_KEY]
 95                )
 96            serialized_slices.append(serialized_slice)
 97        return {"slices": serialized_slices, "state_type": state_type.value}
 98
 99    @staticmethod
100    def is_state_message_compatible(state: MutableMapping[str, Any]) -> bool:
101        return bool(state) and state.get("state_type") in [
102            t.value for t in ConcurrencyCompatibleStateType
103        ]
104
105    @abstractmethod
106    def convert_from_sequential_state(
107        self,
108        cursor_field: "CursorField",  # to deprecate as it is only needed for sequential state
109        stream_state: MutableMapping[str, Any],
110        start: Optional[Any],
111    ) -> Tuple[Any, MutableMapping[str, Any]]:
112        """
113        Convert the state message to the format required by the ConcurrentCursor.
114
115        e.g.
116        {
117            "state_type": ConcurrencyCompatibleStateType.date_range.value,
118            "metadata": { … },
119            "slices": [
120                {starts: 0, end: 1617030403, finished_processing: true}]
121        }
122        """
123        ...
124
125    @abstractmethod
126    def increment(self, value: Any) -> Any:
127        """
128        Increment a timestamp by a single unit.
129        """
130        ...
131
132    @abstractmethod
133    def output_format(self, value: Any) -> Any:
134        """
135        Convert the cursor value type to a JSON valid type.
136        """
137        ...
138
139    def merge_intervals(
140        self, intervals: List[MutableMapping[str, Any]]
141    ) -> List[MutableMapping[str, Any]]:
142        """
143        Compute and return a list of merged intervals.
144
145        Intervals may be merged if the start time of the second interval is 1 unit or less (as defined by the
146        `increment` method) than the end time of the first interval.
147        """
148        if not intervals:
149            return []
150
151        sorted_intervals = sorted(
152            intervals, key=lambda interval: (interval[self.START_KEY], interval[self.END_KEY])
153        )
154        merged_intervals = [sorted_intervals[0]]
155
156        for current_interval in sorted_intervals[1:]:
157            last_interval = merged_intervals[-1]
158            last_interval_end = last_interval[self.END_KEY]
159            current_interval_start = current_interval[self.START_KEY]
160
161            if self.increment(last_interval_end) >= current_interval_start:
162                last_interval[self.END_KEY] = max(last_interval_end, current_interval[self.END_KEY])
163                last_interval_cursor_value = last_interval.get("most_recent_cursor_value")
164                current_interval_cursor_value = current_interval.get("most_recent_cursor_value")
165
166                last_interval["most_recent_cursor_value"] = (
167                    max(current_interval_cursor_value, last_interval_cursor_value)
168                    if current_interval_cursor_value and last_interval_cursor_value
169                    else current_interval_cursor_value or last_interval_cursor_value
170                )
171            else:
172                # Add a new interval if no overlap
173                merged_intervals.append(current_interval)
174
175        return merged_intervals
176
177    @abstractmethod
178    def parse_value(self, value: Any) -> Any:
179        """
180        Parse the value of the cursor field into a comparable value.
181        """
182        ...
183
184    @property
185    @abstractmethod
186    def zero_value(self) -> Any: ...