airbyte.records

  1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
  2"""PyAirbyte Records module.
  3
  4## Understanding record handling in PyAirbyte
  5
  6PyAirbyte models record handling after Airbyte's "Destination V2" ("Dv2") record handling. This
  7includes the below implementation details.
  8
  9### Field Name Normalization
 10
 111. PyAirbyte normalizes top-level record keys to lowercase, replacing spaces and hyphens with
 12   underscores.
 132. PyAirbyte does not normalize nested keys on sub-properties.
 14
 15For example, the following record:
 16
 17```json
 18{
 19
 20    "My-Field": "value",
 21    "Nested": {
 22        "MySubField": "value"
 23    }
 24}
 25```
 26
 27Would be normalized to:
 28
 29```json
 30{
 31    "my_field": "value",
 32    "nested": {
 33        "MySubField": "value"
 34    }
 35}
 36```
 37
 38### Table Name Normalization
 39
 40Similar to column handling, PyAirbyte normalizes table names to the lowercase version of the stream
 41name and may remove or normalize special characters.
 42
 43### Airbyte-Managed Metadata Columns
 44
 45PyAirbyte adds the following columns to every record:
 46
 47- `ab_raw_id`: A unique identifier for the record.
 48- `ab_extracted_at`: The time the record was extracted.
 49- `ab_meta`: A dictionary of extra metadata about the record.
 50
 51The names of these columns are included in the `airbyte.constants` module for programmatic
 52reference.
 53
 54## Schema Evolution
 55
 56PyAirbyte supports a very basic form of schema evolution:
 57
 581. Columns are always auto-added to cache tables whenever newly arriving properties are detected
 59   as not present in the cache table.
 602. Column types will not be modified or expanded to fit changed types in the source catalog.
 61   - If column types change, we recommend user to manually alter the column types.
 623. At any time, users can run a full sync with a `WriteStrategy` of 'replace'. This will create a
 63   fresh table from scratch and then swap the old and new tables after table sync is complete.
 64
 65---
 66
 67"""
 68
 69from __future__ import annotations
 70
 71from datetime import datetime
 72from typing import TYPE_CHECKING, Any
 73
 74import pytz
 75from uuid_extensions import uuid7str
 76
 77from airbyte._util.name_normalizers import LowerCaseNormalizer, NameNormalizerBase
 78from airbyte.constants import (
 79    AB_EXTRACTED_AT_COLUMN,
 80    AB_INTERNAL_COLUMNS,
 81    AB_META_COLUMN,
 82    AB_RAW_ID_COLUMN,
 83)
 84
 85
 86if TYPE_CHECKING:
 87    from collections.abc import Iterator
 88
 89    from airbyte_protocol.models import (
 90        AirbyteRecordMessage,
 91    )
 92
 93
 94class StreamRecordHandler:
 95    """A class to handle processing of StreamRecords.
 96
 97    This is a long-lived object that can be used to process multiple StreamRecords, which
 98    saves on memory and processing time by reusing the same object for all records of the same type.
 99    """
100
101    def __init__(
102        self,
103        *,
104        json_schema: dict,
105        normalizer: type[NameNormalizerBase] = LowerCaseNormalizer,
106        normalize_keys: bool = False,
107        prune_extra_fields: bool,
108    ) -> None:
109        """Initialize the dictionary with the given data.
110
111        Args:
112            json_schema: The JSON Schema definition for this stream.
113            normalizer: The normalizer to use when normalizing keys. If not provided, the
114                LowerCaseNormalizer will be used.
115            normalize_keys: If `True`, the keys will be normalized using the given normalizer.
116            prune_extra_fields: If `True`, unexpected fields will be removed.
117        """
118        self._expected_keys: list[str] = list(json_schema.get("properties", {}).keys())
119        self._normalizer: type[NameNormalizerBase] = normalizer
120        self._normalize_keys: bool = normalize_keys
121        self.prune_extra_fields: bool = prune_extra_fields
122
123        self.index_keys: list[str] = [
124            self._normalizer.normalize(key) if self._normalize_keys else key
125            for key in self._expected_keys
126        ]
127        self.normalized_keys: list[str] = [
128            self._normalizer.normalize(key) for key in self._expected_keys
129        ]
130        self.quick_lookup: dict[str, str]
131
132        for internal_col in AB_INTERNAL_COLUMNS:
133            if internal_col not in self._expected_keys:
134                self._expected_keys.append(internal_col)
135
136        # Store a lookup from normalized keys to pretty cased (originally cased) keys.
137        self._pretty_case_lookup: dict[str, str] = {
138            self._normalizer.normalize(pretty_case.lower()): pretty_case
139            for pretty_case in self._expected_keys
140        }
141        # Store a map from all key versions (normalized and pretty-cased) to their normalized
142        # version.
143        self.quick_lookup = {
144            key: self._normalizer.normalize(key)
145            if self._normalize_keys
146            else self.to_display_case(key)
147            for key in set(self._expected_keys) | set(self._pretty_case_lookup.values())
148        }
149
150    def to_display_case(self, key: str) -> str:
151        """Return the original case of the key."""
152        return self._pretty_case_lookup[self._normalizer.normalize(key)]
153
154    def to_index_case(self, key: str) -> str:
155        """Return the internal case of the key.
156
157        If `normalize_keys` is True, returns the normalized key.
158        Otherwise, return the original case ("pretty case") of the key.
159        """
160        try:
161            return self.quick_lookup[key]
162        except KeyError:
163            result = (
164                self._normalizer.normalize(key)
165                if self._normalize_keys
166                else self.to_display_case(key)
167            )
168            self.quick_lookup[key] = result
169            return result
170
171
172class StreamRecord(dict[str, Any]):
173    """The StreamRecord class is a case-aware, case-insensitive dictionary implementation.
174
175    It has these behaviors:
176    - When a key is retrieved, deleted, or checked for existence, it is always checked in a
177      case-insensitive manner.
178    - The original case is stored in a separate dictionary, so that the original case can be
179      retrieved when needed.
180    - Because it is subclassed from `dict`, the `StreamRecord` class can be passed as a normal
181      Python dictionary.
182    - In addition to the properties of the stream's records, the dictionary also stores the Airbyte
183      metadata columns: `_airbyte_raw_id`, `_airbyte_extracted_at`, and `_airbyte_meta`.
184
185    This behavior mirrors how a case-aware, case-insensitive SQL database would handle column
186    references.
187
188    There are two ways this class can store keys internally:
189    - If normalize_keys is True, the keys are normalized using the given normalizer.
190    - If normalize_keys is False, the original case of the keys is stored.
191
192    In regards to missing values, the dictionary accepts an 'expected_keys' input. When set, the
193    dictionary will be initialized with the given keys. If a key is not found in the input data, it
194    will be initialized with a value of None. When provided, the 'expected_keys' input will also
195    determine the original case of the keys.
196    """
197
198    def __init__(
199        self,
200        from_dict: dict,
201        *,
202        stream_record_handler: StreamRecordHandler,
203        with_internal_columns: bool = True,
204        extracted_at: datetime | None = None,
205    ) -> None:
206        """Initialize the dictionary with the given data.
207
208        Args:
209            from_dict: The dictionary to initialize the StreamRecord with.
210            stream_record_handler: The StreamRecordHandler to use for processing the record.
211            with_internal_columns: If `True`, the internal columns will be added to the record.
212            extracted_at: The time the record was extracted. If not provided, the current time will
213                be used.
214        """
215        self._stream_handler: StreamRecordHandler = stream_record_handler
216
217        # Start by initializing all values to None
218        self.update(dict.fromkeys(stream_record_handler.index_keys))
219
220        # Update the dictionary with the given data
221        if self._stream_handler.prune_extra_fields:
222            self.update(
223                {
224                    self._stream_handler.to_index_case(k): v
225                    for k, v in from_dict.items()
226                    if self._stream_handler.to_index_case(k) in self._stream_handler.index_keys
227                }
228            )
229        else:
230            self.update({self._stream_handler.to_index_case(k): v for k, v in from_dict.items()})
231
232        if with_internal_columns:
233            self.update(
234                {
235                    AB_RAW_ID_COLUMN: uuid7str(),
236                    AB_EXTRACTED_AT_COLUMN: extracted_at or datetime.now(pytz.utc),
237                    AB_META_COLUMN: {},
238                }
239            )
240
241    @classmethod
242    def from_record_message(
243        cls,
244        record_message: AirbyteRecordMessage,
245        *,
246        stream_record_handler: StreamRecordHandler,
247    ) -> StreamRecord:
248        """Return a StreamRecord from a RecordMessage."""
249        data_dict: dict[str, Any] = record_message.data.copy()
250        return cls(
251            from_dict=data_dict,
252            stream_record_handler=stream_record_handler,
253            with_internal_columns=True,
254            extracted_at=datetime.fromtimestamp(record_message.emitted_at / 1000, tz=pytz.utc),
255        )
256
257    def __getitem__(self, key: str) -> Any:  # noqa: ANN401
258        """Return the item with the given key."""
259        try:
260            return super().__getitem__(key)
261        except KeyError:
262            return super().__getitem__(self._stream_handler.to_index_case(key))
263
264    def __setitem__(self, key: str, value: Any) -> None:  # noqa: ANN401
265        """Set the item with the given key to the given value."""
266        index_case_key = self._stream_handler.to_index_case(key)
267        if (
268            self._stream_handler.prune_extra_fields
269            and index_case_key not in self._stream_handler.index_keys
270        ):
271            return
272
273        super().__setitem__(index_case_key, value)
274
275    def __delitem__(self, key: str) -> None:
276        """Delete the item with the given key."""
277        try:
278            super().__delitem__(key)
279        except KeyError:
280            index_case_key = self._stream_handler.to_index_case(key)
281            if super().__contains__(index_case_key):
282                super().__delitem__(index_case_key)
283                return
284        else:
285            # No failure. Key was deleted.
286            return
287
288        raise KeyError(key)
289
290    def __contains__(self, key: object) -> bool:
291        """Return whether the dictionary contains the given key."""
292        assert isinstance(key, str), "Key must be a string."
293        return super().__contains__(key) or super().__contains__(
294            self._stream_handler.to_index_case(key)
295        )
296
297    def __iter__(self) -> Iterator[str]:
298        """Return an iterator over the keys of the dictionary."""
299        return iter(super().__iter__())
300
301    def __len__(self) -> int:
302        """Return the number of items in the dictionary."""
303        return super().__len__()
304
305    def __eq__(self, other: object) -> bool:
306        """Return whether the StreamRecord is equal to the given dict or StreamRecord object."""
307        if isinstance(other, StreamRecord):
308            return dict(self) == dict(other)
309
310        if isinstance(other, dict):
311            return {k.lower(): v for k, v in self.items()} == {
312                k.lower(): v for k, v in other.items()
313            }
314        return False
315
316    def __hash__(self) -> int:  # type: ignore [override]  # Doesn't match superclass (dict)
317        """Return the hash of the dictionary with keys sorted."""
318        items = [(k, v) for k, v in self.items() if not isinstance(v, dict)]
319        return hash(tuple(sorted(items)))