airbyte_cdk.sources.declarative.extractors

11@dataclass
12class TypeTransformer(ABC):
13    """
14    Abstract base class for implementing type transformation logic.
15
16    This class provides a blueprint for defining custom transformations
17    on data records based on a provided schema. Implementing classes
18    must override the `transform` method to specify the transformation
19    logic.
20
21    Attributes:
22        None explicitly defined, as this is a dataclass intended to be
23        subclassed.
24
25    Methods:
26        transform(record: Dict[str, Any], schema: Mapping[str, Any]) -> None:
27            Abstract method that must be implemented by subclasses.
28            It performs a transformation on a given data record based
29            on the provided schema.
30
31    Usage:
32        To use this class, create a subclass that implements the
33        `transform` method with the desired transformation logic.
34    """
35
36    @abstractmethod
37    def transform(
38        self,
39        record: Dict[str, Any],
40        schema: Mapping[str, Any],
41    ) -> None:
42        """
43        Perform a transformation on a data record based on a given schema.
44
45        Args:
46            record (Dict[str, Any]): The data record to be transformed.
47            schema (Mapping[str, Any]): The schema that dictates how
48                the record should be transformed.
49
50        Returns:
51            None
52
53        Raises:
54            NotImplementedError: If the method is not implemented
55                by a subclass.
56        """

14class HttpSelector:
15    """
16    Responsible for translating an HTTP response into a list of records by extracting records from the response and optionally filtering
17    records based on a heuristic.
18    """
19
20    @abstractmethod
21    def select_records(
22        self,
23        response: requests.Response,
24        stream_state: StreamState,
25        records_schema: Mapping[str, Any],
26        stream_slice: Optional[StreamSlice] = None,
27        next_page_token: Optional[Mapping[str, Any]] = None,
28    ) -> Iterable[Record]:
29        """
30        Selects records from the response
31        :param response: The response to select the records from
32        :param stream_state: The stream state
33        :param records_schema: json schema of records to return
34        :param stream_slice: The stream slice
35        :param next_page_token: The paginator token
36        :return: List of Records selected from the response
37        """
38        pass

 19@dataclass
 20class DpathExtractor(RecordExtractor):
 21    """
 22    Record extractor that searches a decoded response over a path defined as an array of fields.
 23
 24    If the field path points to an array, that array is returned.
 25    If the field path points to an object, that object is returned wrapped as an array.
 26    If the field path points to an empty object, an empty array is returned.
 27    If the field path points to a non-existing path, an empty array is returned.
 28
 29    Optionally, records can be expanded by providing a RecordExpander component.
 30    When record_expander is configured, each extracted record is passed through the
 31    expander which extracts items from nested array fields and emits each item as a
 32    separate record.
 33
 34    Examples of instantiating this transform:
 35    ```
 36      extractor:
 37        type: DpathExtractor
 38        field_path:
 39          - "root"
 40          - "data"
 41    ```
 42
 43    ```
 44      extractor:
 45        type: DpathExtractor
 46        field_path:
 47          - "root"
 48          - "{{ parameters['field'] }}"
 49    ```
 50
 51    ```
 52      extractor:
 53        type: DpathExtractor
 54        field_path: []
 55    ```
 56
 57    ```
 58      extractor:
 59        type: DpathExtractor
 60        field_path:
 61          - "data"
 62          - "object"
 63        record_expander:
 64          type: RecordExpander
 65          expand_records_from_field:
 66            - "lines"
 67            - "data"
 68          remain_original_record: true
 69    ```
 70
 71    Attributes:
 72        field_path (Union[InterpolatedString, str]): Path to the field that should be extracted
 73        config (Config): The user-provided configuration as specified by the source's spec
 74        decoder (Decoder): The decoder responsible to transfom the response in a Mapping
 75        record_expander (Optional[RecordExpander]): Optional component to expand records by extracting items from nested array fields
 76    """
 77
 78    field_path: List[Union[InterpolatedString, str]]
 79    config: Config
 80    parameters: InitVar[Mapping[str, Any]]
 81    decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={}))
 82    record_expander: Optional[RecordExpander] = None
 83
 84    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 85        self._field_path = [
 86            InterpolatedString.create(path, parameters=parameters) for path in self.field_path
 87        ]
 88        for path_index in range(len(self.field_path)):
 89            if isinstance(self.field_path[path_index], str):
 90                self._field_path[path_index] = InterpolatedString.create(
 91                    self.field_path[path_index], parameters=parameters
 92                )
 93
 94    def extract_records(self, response: requests.Response) -> Iterable[MutableMapping[Any, Any]]:
 95        for body in self.decoder.decode(response):
 96            if len(self._field_path) == 0:
 97                extracted = body
 98            else:
 99                path = [path.eval(self.config) for path in self._field_path]
100                if "*" in path:
101                    extracted = dpath.values(body, path)
102                else:
103                    extracted = dpath.get(body, path, default=[])  # type: ignore # extracted will be a MutableMapping, given input data structure
104            if isinstance(extracted, list):
105                if not self.record_expander:
106                    yield from extracted
107                else:
108                    for record in extracted:
109                        yield from self.record_expander.expand_record(record)
110            elif extracted:
111                if self.record_expander:
112                    yield from self.record_expander.expand_record(extracted)
113                else:
114                    yield extracted
115            else:
116                yield from []

13@dataclass
14class RecordFilter:
15    """
16    Filter applied on a list of Records
17
18    config (Config): The user-provided configuration as specified by the source's spec
19    condition (str): The string representing the predicate to filter a record. Records will be removed if evaluated to False
20    """
21
22    parameters: InitVar[Mapping[str, Any]]
23    config: Config
24    condition: str = ""
25
26    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
27        self._filter_interpolator = InterpolatedBoolean(
28            condition=self.condition, parameters=parameters
29        )
30
31    def filter_records(
32        self,
33        records: Iterable[Mapping[str, Any]],
34        stream_state: StreamState,
35        stream_slice: Optional[StreamSlice] = None,
36        next_page_token: Optional[Mapping[str, Any]] = None,
37    ) -> Iterable[Mapping[str, Any]]:
38        kwargs = {
39            "stream_state": stream_state,
40            "stream_slice": stream_slice,
41            "next_page_token": next_page_token,
42            "stream_slice.extra_fields": stream_slice.extra_fields if stream_slice else {},
43        }
44        for record in records:
45            if self._filter_interpolator.eval(self.config, record=record, **kwargs):
46                yield record

 25@dataclass
 26class RecordSelector(HttpSelector):
 27    """
 28    Responsible for translating an HTTP response into a list of records by extracting records from the response and optionally filtering
 29    records based on a heuristic.
 30
 31    Attributes:
 32        extractor (RecordExtractor): The record extractor responsible for extracting records from a response
 33        schema_normalization (TypeTransformer): The record normalizer responsible for casting record values to stream schema types
 34        record_filter (RecordFilter): The record filter responsible for filtering extracted records
 35        transformations (List[RecordTransformation]): The transformations to be done on the records
 36    """
 37
 38    extractor: RecordExtractor
 39    config: Config
 40    parameters: InitVar[Mapping[str, Any]]
 41    schema_normalization: Union[TypeTransformer, DeclarativeTypeTransformer]
 42    name: str
 43    _name: Union[InterpolatedString, str] = field(init=False, repr=False, default="")
 44    record_filter: Optional[RecordFilter] = None
 45    transformations: List[RecordTransformation] = field(default_factory=lambda: [])
 46    transform_before_filtering: bool = False
 47    file_uploader: Optional[DefaultFileUploader] = None
 48
 49    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 50        self._parameters = parameters
 51        self._name = (
 52            InterpolatedString(self._name, parameters=parameters)
 53            if isinstance(self._name, str)
 54            else self._name
 55        )
 56
 57    @property  # type: ignore
 58    def name(self) -> str:
 59        """
 60        :return: Stream name
 61        """
 62        return (
 63            str(self._name.eval(self.config))
 64            if isinstance(self._name, InterpolatedString)
 65            else self._name
 66        )
 67
 68    @name.setter
 69    def name(self, value: str) -> None:
 70        if not isinstance(value, property):
 71            self._name = value
 72
 73    def select_records(
 74        self,
 75        response: requests.Response,
 76        stream_state: StreamState,
 77        records_schema: Mapping[str, Any],
 78        stream_slice: Optional[StreamSlice] = None,
 79        next_page_token: Optional[Mapping[str, Any]] = None,
 80    ) -> Iterable[Record]:
 81        """
 82        Selects records from the response
 83        :param response: The response to select the records from
 84        :param stream_state: The stream state
 85        :param records_schema: json schema of records to return
 86        :param stream_slice: The stream slice
 87        :param next_page_token: The paginator token
 88        :return: List of Records selected from the response
 89        """
 90        all_data: Iterable[Mapping[str, Any]] = self.extractor.extract_records(response)
 91        yield from self.filter_and_transform(
 92            all_data, stream_state, records_schema, stream_slice, next_page_token
 93        )
 94
 95    def filter_and_transform(
 96        self,
 97        all_data: Iterable[Mapping[str, Any]],
 98        stream_state: StreamState,
 99        records_schema: Mapping[str, Any],
100        stream_slice: Optional[StreamSlice] = None,
101        next_page_token: Optional[Mapping[str, Any]] = None,
102    ) -> Iterable[Record]:
103        """
104        There is an issue with the selector as of 2024-08-30: it does technology-agnostic processing like filtering, transformation and
105        normalization with an API that is technology-specific (as requests.Response is only for HTTP communication using the requests
106        library).
107
108        Until we decide to move this logic away from the selector, we made this method public so that users like AsyncJobRetriever could
109        share the logic of doing transformations on a set of records.
110        """
111        if self.transform_before_filtering:
112            transformed_data = self._transform(all_data, stream_state, stream_slice)
113            transformed_filtered_data = self._filter(
114                transformed_data, stream_state, stream_slice, next_page_token
115            )
116        else:
117            filtered_data = self._filter(all_data, stream_state, stream_slice, next_page_token)
118            transformed_filtered_data = self._transform(filtered_data, stream_state, stream_slice)
119        normalized_data = self._normalize_by_schema(
120            transformed_filtered_data, schema=records_schema
121        )
122        for data in normalized_data:
123            record = Record(data=data, stream_name=self.name, associated_slice=stream_slice)
124            if self.file_uploader:
125                self.file_uploader.upload(record)
126            yield record
127
128    def _normalize_by_schema(
129        self, records: Iterable[Mapping[str, Any]], schema: Optional[Mapping[str, Any]]
130    ) -> Iterable[Mapping[str, Any]]:
131        if schema:
132            # record has type Mapping[str, Any], but dict[str, Any] expected
133            for record in records:
134                normalized_record = dict(record)
135                self.schema_normalization.transform(normalized_record, schema)
136                yield normalized_record
137        else:
138            yield from records
139
140    def _filter(
141        self,
142        records: Iterable[Mapping[str, Any]],
143        stream_state: StreamState,
144        stream_slice: Optional[StreamSlice],
145        next_page_token: Optional[Mapping[str, Any]],
146    ) -> Iterable[Mapping[str, Any]]:
147        if self.record_filter:
148            yield from self.record_filter.filter_records(
149                records,
150                stream_state=stream_state,
151                stream_slice=stream_slice,
152                next_page_token=next_page_token,
153            )
154        else:
155            yield from records
156
157    def _transform(
158        self,
159        records: Iterable[Mapping[str, Any]],
160        stream_state: StreamState,
161        stream_slice: Optional[StreamSlice] = None,
162    ) -> Iterable[Mapping[str, Any]]:
163        for record in records:
164            for transformation in self.transformations:
165                transformation.transform(
166                    record,  # type: ignore  # record has type Mapping[str, Any], but Dict[str, Any] expected
167                    config=self.config,
168                    stream_state=stream_state,
169                    stream_slice=stream_slice,
170                )
171            yield record

 23@dataclass
 24class ResponseToFileExtractor(RecordExtractor):
 25    """
 26    This class is used when having very big HTTP responses (usually streamed) which would require too much memory so we use disk space as
 27    a tradeoff.
 28
 29    Eventually, we want to support multiple file type by re-using the file based CDK parsers if possible. However, the lift is too high for
 30    a first iteration so we will only support CSV parsing using pandas as salesforce and sendgrid were doing.
 31
 32    By default, pandas interprets strings such as "NA", "N/A", "null" and "NaN" as missing values and converts them to null. When
 33    ``preserve_na_values`` is True, these strings are kept as-is (only empty cells are treated as null). This is opt-in to avoid changing
 34    behavior for existing connectors.
 35    """
 36
 37    parameters: InitVar[Mapping[str, Any]]
 38    preserve_na_values: bool = False
 39
 40    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 41        self.logger = logging.getLogger("airbyte")
 42
 43    def _get_response_encoding(self, headers: Dict[str, Any]) -> str:
 44        """
 45        Get the encoding of the response based on the provided headers. This method is heavily inspired by the requests library
 46        implementation.
 47
 48        Args:
 49            headers (Dict[str, Any]): The headers of the response.
 50        Returns:
 51            str: The encoding of the response.
 52        """
 53
 54        content_type = headers.get("content-type")
 55
 56        if not content_type:
 57            return DEFAULT_ENCODING
 58
 59        content_type, params = requests.utils.parse_header_links(content_type)
 60
 61        if "charset" in params:
 62            return params["charset"].strip("'\"")  # type: ignore  # we assume headers are returned as str
 63
 64        return DEFAULT_ENCODING
 65
 66    def _filter_null_bytes(self, b: bytes) -> bytes:
 67        """
 68        Filter out null bytes from a bytes object.
 69
 70        Args:
 71            b (bytes): The input bytes object.
 72        Returns:
 73            bytes: The filtered bytes object with null bytes removed.
 74
 75        Referenced Issue:
 76            https://github.com/airbytehq/airbyte/issues/8300
 77        """
 78
 79        res = b.replace(b"\x00", b"")
 80        if len(res) < len(b):
 81            self.logger.warning(
 82                "Filter 'null' bytes from string, size reduced %d -> %d chars", len(b), len(res)
 83            )
 84        return res
 85
 86    def _save_to_file(self, response: requests.Response) -> Tuple[str, str]:
 87        """
 88        Saves the binary data from the given response to a temporary file and returns the filepath and response encoding.
 89
 90        Args:
 91            response (Optional[requests.Response]): The response object containing the binary data. Defaults to None.
 92
 93        Returns:
 94            Tuple[str, str]: A tuple containing the filepath of the temporary file and the response encoding.
 95
 96        Raises:
 97            ValueError: If the temporary file does not exist after saving the binary data.
 98        """
 99        # set filepath for binary data from response
100        decompressor = zlib.decompressobj(zlib.MAX_WBITS | 32)
101        needs_decompression = True  # we will assume at first that the response is compressed and change the flag if not
102
103        tmp_file = str(uuid.uuid4())
104        with closing(response) as response, open(tmp_file, "wb") as data_file:
105            response_encoding = self._get_response_encoding(dict(response.headers or {}))
106            for chunk in response.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
107                try:
108                    if needs_decompression:
109                        data_file.write(decompressor.decompress(chunk))
110                        needs_decompression = True
111                    else:
112                        data_file.write(self._filter_null_bytes(chunk))
113                except zlib.error:
114                    data_file.write(self._filter_null_bytes(chunk))
115                    needs_decompression = False
116
117        # check the file exists
118        if os.path.isfile(tmp_file):
119            return tmp_file, response_encoding
120        else:
121            raise ValueError(
122                f"The IO/Error occured while verifying binary data. Tmp file {tmp_file} doesn't exist."
123            )
124
125    def _read_with_chunks(
126        self, path: str, file_encoding: str, chunk_size: int = 100
127    ) -> Iterable[Mapping[str, Any]]:
128        """
129        Reads data from a file in chunks and yields each row as a dictionary.
130
131        Args:
132            path (str): The path to the file to be read.
133            file_encoding (str): The encoding of the file.
134            chunk_size (int, optional): The size of each chunk to be read. Defaults to 100.
135
136        Yields:
137            Mapping[str, Any]: A dictionary representing each row of data.
138
139        Raises:
140            ValueError: If an IO/Error occurs while reading the temporary data.
141        """
142
143        try:
144            with open(path, "r", encoding=file_encoding) as data:
145                # When preserving NA values, disable pandas' default NA parsing so strings like "NA"/"N/A" are kept as-is,
146                # and explicitly map empty cells to None. Otherwise, keep the historical behavior where pandas converts its
147                # default NA tokens (including empty cells) to NaN, which we then map to None.
148                na_replacements = {nan: None, "": None} if self.preserve_na_values else {nan: None}
149                chunks = pd.read_csv(
150                    data,
151                    chunksize=chunk_size,
152                    iterator=True,
153                    dialect="unix",
154                    dtype=object,
155                    keep_default_na=not self.preserve_na_values,
156                )
157                for chunk in chunks:
158                    chunk = chunk.replace(na_replacements).to_dict(orient="records")
159                    for row in chunk:
160                        yield row
161        except pd.errors.EmptyDataError as e:
162            self.logger.info(f"Empty data received. {e}")
163            yield from []
164        except IOError as ioe:
165            raise ValueError(f"The IO/Error occured while reading tmp data. Called: {path}", ioe)
166        finally:
167            # remove binary tmp file, after data is read
168            os.remove(path)
169
170    def extract_records(
171        self, response: Optional[requests.Response] = None
172    ) -> Iterable[Mapping[str, Any]]:
173        """
174        Extracts records from the given response by:
175            1) Saving the result to a tmp file
176            2) Reading from saved file by chunks to avoid OOM
177
178        Args:
179            response (Optional[requests.Response]): The response object containing the data. Defaults to None.
180
181        Yields:
182            Iterable[Mapping[str, Any]]: An iterable of mappings representing the extracted records.
183
184        Returns:
185            None
186        """
187        if response:
188            file_path, encoding = self._save_to_file(response)
189            yield from self._read_with_chunks(file_path, encoding)
190        else:
191            yield from []