airbyte_cdk.sources.file_based.file_types

 47class AvroParser(FileTypeParser):
 48    ENCODING = None
 49
 50    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
 51        """
 52        AvroParser does not require config checks, implicit pydantic validation is enough.
 53        """
 54        return True, None
 55
 56    async def infer_schema(
 57        self,
 58        config: FileBasedStreamConfig,
 59        file: RemoteFile,
 60        stream_reader: AbstractFileBasedStreamReader,
 61        logger: logging.Logger,
 62    ) -> SchemaType:
 63        avro_format = config.format
 64        if not isinstance(avro_format, AvroFormat):
 65            raise ValueError(f"Expected ParquetFormat, got {avro_format}")
 66
 67        with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
 68            avro_reader = fastavro.reader(fp)  # type: ignore [arg-type]
 69            avro_schema = avro_reader.writer_schema
 70        if not avro_schema["type"] == "record":  # type: ignore [index, call-overload]
 71            unsupported_type = avro_schema["type"]  # type: ignore [index, call-overload]
 72            raise ValueError(
 73                f"Only record based avro files are supported. Found {unsupported_type}"
 74            )
 75        json_schema = {
 76            field["name"]: AvroParser._convert_avro_type_to_json(  # type: ignore [index]
 77                avro_format,
 78                field["name"],  # type: ignore [index]
 79                field["type"],  # type: ignore [index]
 80            )
 81            for field in avro_schema["fields"]  # type: ignore [index, call-overload]
 82        }
 83        return json_schema
 84
 85    @classmethod
 86    def _convert_avro_type_to_json(
 87        cls, avro_format: AvroFormat, field_name: str, avro_field: str
 88    ) -> Mapping[str, Any]:
 89        if isinstance(avro_field, str) and avro_field in AVRO_TYPE_TO_JSON_TYPE:
 90            # Legacy behavior to retain backwards compatibility. Long term we should always represent doubles as strings
 91            if avro_field == "double" and not avro_format.double_as_string:
 92                return {"type": "number"}
 93            return {"type": AVRO_TYPE_TO_JSON_TYPE[avro_field]}
 94        if isinstance(avro_field, Mapping):
 95            if avro_field["type"] == "record":
 96                return {
 97                    "type": "object",
 98                    "properties": {
 99                        object_field["name"]: AvroParser._convert_avro_type_to_json(
100                            avro_format, object_field["name"], object_field["type"]
101                        )
102                        for object_field in avro_field["fields"]
103                    },
104                }
105            elif avro_field["type"] == "array":
106                if "items" not in avro_field:
107                    raise ValueError(
108                        f"{field_name} array type does not have a required field items"
109                    )
110                return {
111                    "type": "array",
112                    "items": AvroParser._convert_avro_type_to_json(
113                        avro_format, "", avro_field["items"]
114                    ),
115                }
116            elif avro_field["type"] == "enum":
117                if "symbols" not in avro_field:
118                    raise ValueError(
119                        f"{field_name} enum type does not have a required field symbols"
120                    )
121                if "name" not in avro_field:
122                    raise ValueError(f"{field_name} enum type does not have a required field name")
123                return {"type": "string", "enum": avro_field["symbols"]}
124            elif avro_field["type"] == "map":
125                if "values" not in avro_field:
126                    raise ValueError(f"{field_name} map type does not have a required field values")
127                return {
128                    "type": "object",
129                    "additionalProperties": AvroParser._convert_avro_type_to_json(
130                        avro_format, "", avro_field["values"]
131                    ),
132                }
133            elif avro_field["type"] == "fixed" and avro_field.get("logicalType") != "duration":
134                if "size" not in avro_field:
135                    raise ValueError(f"{field_name} fixed type does not have a required field size")
136                if not isinstance(avro_field["size"], int):
137                    raise ValueError(f"{field_name} fixed type size value is not an integer")
138                return {
139                    "type": "string",
140                    "pattern": f"^[0-9A-Fa-f]{{{avro_field['size'] * 2}}}$",
141                }
142            elif avro_field.get("logicalType") == "decimal":
143                if "precision" not in avro_field:
144                    raise ValueError(
145                        f"{field_name} decimal type does not have a required field precision"
146                    )
147                if "scale" not in avro_field:
148                    raise ValueError(
149                        f"{field_name} decimal type does not have a required field scale"
150                    )
151                max_whole_number_range = avro_field["precision"] - avro_field["scale"]
152                decimal_range = avro_field["scale"]
153
154                # This regex looks like a mess, but it is validation for at least one whole number and optional fractional numbers
155                # For example: ^-?\d{1,5}(?:\.\d{1,3})?$ would accept 12345.123 and 123456.12345 would  be rejected
156                return {
157                    "type": "string",
158                    "pattern": f"^-?\\d{{{1,max_whole_number_range}}}(?:\\.\\d{1,decimal_range})?$",
159                }
160            elif "logicalType" in avro_field:
161                if avro_field["logicalType"] not in AVRO_LOGICAL_TYPE_TO_JSON:
162                    raise ValueError(
163                        f"{avro_field['logicalType']} is not a valid Avro logical type"
164                    )
165                return AVRO_LOGICAL_TYPE_TO_JSON[avro_field["logicalType"]]
166            else:
167                raise ValueError(f"Unsupported avro type: {avro_field}")
168        else:
169            raise ValueError(f"Unsupported avro type: {avro_field}")
170
171    def parse_records(
172        self,
173        config: FileBasedStreamConfig,
174        file: RemoteFile,
175        stream_reader: AbstractFileBasedStreamReader,
176        logger: logging.Logger,
177        discovered_schema: Optional[Mapping[str, SchemaType]],
178    ) -> Iterable[Dict[str, Any]]:
179        avro_format = config.format or AvroFormat(filetype="avro")
180        if not isinstance(avro_format, AvroFormat):
181            raise ValueError(f"Expected ParquetFormat, got {avro_format}")
182
183        line_no = 0
184        try:
185            with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
186                avro_reader = fastavro.reader(fp)  # type: ignore [arg-type]
187                schema = avro_reader.writer_schema
188                schema_field_name_to_type = {
189                    field["name"]: cast(dict[str, Any], field["type"])  # type: ignore [index]
190                    for field in schema["fields"]  # type: ignore [index, call-overload]  # If schema is not dict, it is not subscriptable by strings
191                }
192                for record in avro_reader:
193                    line_no += 1
194                    yield {
195                        record_field: self._to_output_value(
196                            avro_format,
197                            schema_field_name_to_type[record_field],  # type: ignore [index] # Any not subscriptable
198                            record[record_field],  # type: ignore [index] # Any not subscriptable
199                        )
200                        for record_field, record_value in schema_field_name_to_type.items()
201                    }
202        except Exception as exc:
203            raise RecordParseError(
204                FileBasedSourceError.ERROR_PARSING_RECORD, filename=file.uri, lineno=line_no
205            ) from exc
206
207    @property
208    def file_read_mode(self) -> FileReadMode:
209        return FileReadMode.READ_BINARY
210
211    @staticmethod
212    def _to_output_value(
213        avro_format: AvroFormat, record_type: Mapping[str, Any], record_value: Any
214    ) -> Any:
215        if isinstance(record_value, bytes):
216            return record_value.decode()
217        elif not isinstance(record_type, Mapping):
218            if record_type == "double" and avro_format.double_as_string:
219                return str(record_value)
220            return record_value
221        if record_type.get("logicalType") in ("decimal", "uuid"):
222            return str(record_value)
223        elif record_type.get("logicalType") == "date":
224            return record_value.isoformat()
225        elif record_type.get("logicalType") == "timestamp-millis":
226            return record_value.isoformat(sep="T", timespec="milliseconds")
227        elif record_type.get("logicalType") == "timestamp-micros":
228            return record_value.isoformat(sep="T", timespec="microseconds")
229        elif record_type.get("logicalType") == "local-timestamp-millis":
230            return record_value.isoformat(sep="T", timespec="milliseconds")
231        elif record_type.get("logicalType") == "local-timestamp-micros":
232            return record_value.isoformat(sep="T", timespec="microseconds")
233        else:
234            return record_value

155class CsvParser(FileTypeParser):
156    _MAX_BYTES_PER_FILE_FOR_SCHEMA_INFERENCE = 1_000_000
157
158    def __init__(self, csv_reader: Optional[_CsvReader] = None, csv_field_max_bytes: int = 2**31):
159        # Increase the maximum length of data that can be parsed in a single CSV field. The default is 128k, which is typically sufficient
160        # but given the use of Airbyte in loading a large variety of data it is best to allow for a larger maximum field size to avoid
161        # skipping data on load. https://stackoverflow.com/questions/15063936/csv-error-field-larger-than-field-limit-131072
162        csv.field_size_limit(csv_field_max_bytes)
163        self._csv_reader = csv_reader if csv_reader else _CsvReader()
164
165    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
166        """
167        CsvParser does not require config checks, implicit pydantic validation is enough.
168        """
169        return True, None
170
171    async def infer_schema(
172        self,
173        config: FileBasedStreamConfig,
174        file: RemoteFile,
175        stream_reader: AbstractFileBasedStreamReader,
176        logger: logging.Logger,
177    ) -> SchemaType:
178        input_schema = config.get_input_schema()
179        if input_schema:
180            return input_schema
181
182        # todo: the existing InMemoryFilesSource.open_file() test source doesn't currently require an encoding, but actual
183        #  sources will likely require one. Rather than modify the interface now we can wait until the real use case
184        config_format = _extract_format(config)
185        type_inferrer_by_field: Dict[str, _TypeInferrer] = defaultdict(
186            lambda: _JsonTypeInferrer(
187                config_format.true_values, config_format.false_values, config_format.null_values
188            )
189            if config_format.inference_type != InferenceType.NONE
190            else _DisabledTypeInferrer()
191        )
192        data_generator = self._csv_reader.read_data(
193            config, file, stream_reader, logger, self.file_read_mode
194        )
195        read_bytes = 0
196        for row in data_generator:
197            for header, value in row.items():
198                type_inferrer_by_field[header].add_value(value)
199                # This is not accurate as a representation of how many bytes were read because csv does some processing on the actual value
200                # before returning. Given we would like to be more accurate, we could wrap the IO file using a decorator
201                read_bytes += len(value)
202            read_bytes += len(row) - 1  # for separators
203            if read_bytes >= self._MAX_BYTES_PER_FILE_FOR_SCHEMA_INFERENCE:
204                break
205
206        if not type_inferrer_by_field:
207            raise AirbyteTracedException(
208                message=f"Could not infer schema as there are no rows in {file.uri}. If having an empty CSV file is expected, ignore this. "
209                f"Else, please contact Airbyte.",
210                failure_type=FailureType.config_error,
211            )
212        schema = {
213            header.strip(): {"type": type_inferred.infer()}
214            for header, type_inferred in type_inferrer_by_field.items()
215        }
216        data_generator.close()
217        return schema
218
219    def parse_records(
220        self,
221        config: FileBasedStreamConfig,
222        file: RemoteFile,
223        stream_reader: AbstractFileBasedStreamReader,
224        logger: logging.Logger,
225        discovered_schema: Optional[Mapping[str, SchemaType]],
226    ) -> Iterable[Dict[str, Any]]:
227        line_no = 0
228        try:
229            config_format = _extract_format(config)
230            if discovered_schema:
231                property_types = {
232                    col: prop["type"] for col, prop in discovered_schema["properties"].items()
233                }
234                deduped_property_types = CsvParser._pre_propcess_property_types(property_types)
235            else:
236                deduped_property_types = {}
237            cast_fn = CsvParser._get_cast_function(
238                deduped_property_types, config_format, logger, config.schemaless
239            )
240            data_generator = self._csv_reader.read_data(
241                config, file, stream_reader, logger, self.file_read_mode
242            )
243            for row in data_generator:
244                line_no += 1
245                yield CsvParser._to_nullable(
246                    cast_fn(row),
247                    deduped_property_types,
248                    config_format.null_values,
249                    config_format.strings_can_be_null,
250                )
251        except RecordParseError as parse_err:
252            raise RecordParseError(
253                FileBasedSourceError.ERROR_PARSING_RECORD, filename=file.uri, lineno=line_no
254            ) from parse_err
255        finally:
256            data_generator.close()
257
258    @property
259    def file_read_mode(self) -> FileReadMode:
260        return FileReadMode.READ
261
262    @staticmethod
263    def _get_cast_function(
264        deduped_property_types: Mapping[str, str],
265        config_format: CsvFormat,
266        logger: logging.Logger,
267        schemaless: bool,
268    ) -> Callable[[Mapping[str, str]], Mapping[str, str]]:
269        # Only cast values if the schema is provided
270        if deduped_property_types and not schemaless:
271            return partial(
272                CsvParser._cast_types,
273                deduped_property_types=deduped_property_types,
274                config_format=config_format,
275                logger=logger,
276            )
277        else:
278            # If no schema is provided, yield the rows as they are
279            return _no_cast
280
281    @staticmethod
282    def _to_nullable(
283        row: Mapping[str, str],
284        deduped_property_types: Mapping[str, str],
285        null_values: Set[str],
286        strings_can_be_null: bool,
287    ) -> Dict[str, Optional[str]]:
288        nullable = {
289            k: None
290            if CsvParser._value_is_none(
291                v, deduped_property_types.get(k), null_values, strings_can_be_null
292            )
293            else v
294            for k, v in row.items()
295        }
296        return nullable
297
298    @staticmethod
299    def _value_is_none(
300        value: Any,
301        deduped_property_type: Optional[str],
302        null_values: Set[str],
303        strings_can_be_null: bool,
304    ) -> bool:
305        return value in null_values and (strings_can_be_null or deduped_property_type != "string")
306
307    @staticmethod
308    def _pre_propcess_property_types(property_types: Dict[str, Any]) -> Mapping[str, str]:
309        """
310        Transform the property types to be non-nullable and remove duplicate types if any.
311        Sample input:
312        {
313        "col1": ["string", "null"],
314        "col2": ["string", "string", "null"],
315        "col3": "integer"
316        }
317
318        Sample output:
319        {
320        "col1": "string",
321        "col2": "string",
322        "col3": "integer",
323        }
324        """
325        output = {}
326        for prop, prop_type in property_types.items():
327            if isinstance(prop_type, list):
328                prop_type_distinct = set(prop_type)
329                prop_type_distinct.remove("null")
330                if len(prop_type_distinct) != 1:
331                    raise ValueError(f"Could not get non nullable type from {prop_type}")
332                output[prop] = next(iter(prop_type_distinct))
333            else:
334                output[prop] = prop_type
335        return output
336
337    @staticmethod
338    def _cast_types(
339        row: Dict[str, str],
340        deduped_property_types: Mapping[str, str],
341        config_format: CsvFormat,
342        logger: logging.Logger,
343    ) -> Dict[str, Any]:
344        """
345        Casts the values in the input 'row' dictionary according to the types defined in the JSON schema.
346
347        Array and object types are only handled if they can be deserialized as JSON.
348
349        If any errors are encountered, the value will be emitted as a string.
350        """
351        warnings = []
352        result = {}
353
354        for key, value in row.items():
355            prop_type = deduped_property_types.get(key)
356            cast_value: Any = value
357
358            if prop_type in TYPE_PYTHON_MAPPING and prop_type is not None:
359                _, python_type = TYPE_PYTHON_MAPPING[prop_type]
360
361                if python_type is None:
362                    if value == "":
363                        cast_value = None
364                    else:
365                        warnings.append(_format_warning(key, value, prop_type))
366
367                elif python_type is bool:
368                    try:
369                        cast_value = _value_to_bool(
370                            value, config_format.true_values, config_format.false_values
371                        )
372                    except ValueError:
373                        warnings.append(_format_warning(key, value, prop_type))
374
375                elif python_type is dict:
376                    try:
377                        # we don't re-use _value_to_object here because we type the column as object as long as there is only one object
378                        cast_value = orjson.loads(value)
379                    except orjson.JSONDecodeError:
380                        warnings.append(_format_warning(key, value, prop_type))
381
382                elif python_type is list:
383                    try:
384                        cast_value = _value_to_list(value)
385                    except (ValueError, json.JSONDecodeError):
386                        warnings.append(_format_warning(key, value, prop_type))
387
388                elif python_type:
389                    try:
390                        cast_value = _value_to_python_type(value, python_type)
391                    except ValueError:
392                        warnings.append(_format_warning(key, value, prop_type))
393
394                result[key] = cast_value
395
396        if warnings:
397            logger.warning(
398                f"{FileBasedSourceError.ERROR_CASTING_VALUE.value}: {','.join([w for w in warnings])}",
399            )
400        return result

 35class ExcelParser(FileTypeParser):
 36    ENCODING = None
 37
 38    def check_config(self, config: FileBasedStreamConfig) -> Tuple[bool, Optional[str]]:
 39        """
 40        ExcelParser does not require config checks, implicit pydantic validation is enough.
 41        """
 42        return True, None
 43
 44    async def infer_schema(
 45        self,
 46        config: FileBasedStreamConfig,
 47        file: RemoteFile,
 48        stream_reader: AbstractFileBasedStreamReader,
 49        logger: logging.Logger,
 50    ) -> SchemaType:
 51        """
 52        Infers the schema of the Excel file by examining its contents.
 53
 54        Args:
 55            config (FileBasedStreamConfig): Configuration for the file-based stream.
 56            file (RemoteFile): The remote file to be read.
 57            stream_reader (AbstractFileBasedStreamReader): Reader to read the file.
 58            logger (logging.Logger): Logger for logging information and errors.
 59
 60        Returns:
 61            SchemaType: Inferred schema of the Excel file.
 62        """
 63
 64        # Validate the format of the config
 65        self.validate_format(config.format, logger)
 66
 67        fields: Dict[str, str] = {}
 68
 69        with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
 70            df = self.open_and_parse_file(fp)
 71            for column, df_type in df.dtypes.items():
 72                # Choose the broadest data type if the column's data type differs in dataframes
 73                prev_frame_column_type = fields.get(column)  # type: ignore [call-overload]
 74                fields[column] = self.dtype_to_json_type(  # type: ignore [index]
 75                    prev_frame_column_type,
 76                    df_type,
 77                )
 78
 79        schema = {
 80            field: (
 81                {"type": "string", "format": "date-time"}
 82                if fields[field] == "date-time"
 83                else {"type": fields[field]}
 84            )
 85            for field in fields
 86        }
 87        return schema
 88
 89    def parse_records(
 90        self,
 91        config: FileBasedStreamConfig,
 92        file: RemoteFile,
 93        stream_reader: AbstractFileBasedStreamReader,
 94        logger: logging.Logger,
 95        discovered_schema: Optional[Mapping[str, SchemaType]] = None,
 96    ) -> Iterable[Dict[str, Any]]:
 97        """
 98        Parses records from an Excel file based on the provided configuration.
 99
100        Args:
101            config (FileBasedStreamConfig): Configuration for the file-based stream.
102            file (RemoteFile): The remote file to be read.
103            stream_reader (AbstractFileBasedStreamReader): Reader to read the file.
104            logger (logging.Logger): Logger for logging information and errors.
105            discovered_schema (Optional[Mapping[str, SchemaType]]): Discovered schema for validation.
106
107        Yields:
108            Iterable[Dict[str, Any]]: Parsed records from the Excel file.
109        """
110
111        # Validate the format of the config
112        self.validate_format(config.format, logger)
113
114        try:
115            # Open and parse the file using the stream reader
116            with stream_reader.open_file(file, self.file_read_mode, self.ENCODING, logger) as fp:
117                df = self.open_and_parse_file(fp)
118                # Yield records as dictionaries
119                # DataFrame.to_dict() method returns datetime values in pandas.Timestamp values, which are not serializable by orjson
120                # DataFrame.to_json() returns string with datetime values serialized to iso8601 with microseconds to align with pydantic behavior
121                # see PR description: https://github.com/airbytehq/airbyte/pull/44444/
122                yield from orjson.loads(
123                    df.to_json(orient="records", date_format="iso", date_unit="us")
124                )
125
126        except Exception as exc:
127            # Raise a RecordParseError if any exception occurs during parsing
128            raise RecordParseError(
129                FileBasedSourceError.ERROR_PARSING_RECORD, filename=file.uri
130            ) from exc
131
132    @property
133    def file_read_mode(self) -> FileReadMode:
134        """
135        Returns the file read mode for the Excel file.
136
137        Returns:
138            FileReadMode: The file read mode (binary).
139        """
140        return FileReadMode.READ_BINARY
141
142    @staticmethod
143    def dtype_to_json_type(
144        current_type: Optional[str],
145        dtype: dtype_,  # type: ignore [type-arg]
146    ) -> str:
147        """
148        Convert Pandas DataFrame types to Airbyte Types.
149
150        Args:
151            current_type (Optional[str]): One of the previous types based on earlier dataframes.
152            dtype: Pandas DataFrame type.
153
154        Returns:
155            str: Corresponding Airbyte Type.
156        """
157        number_types = ("int64", "float64")
158        if current_type == "string":
159            # Previous column values were of the string type, no need to look further.
160            return current_type
161        if dtype is object:
162            return "string"
163        if dtype in number_types and (not current_type or current_type == "number"):
164            return "number"
165        if dtype == "bool" and (not current_type or current_type == "boolean"):
166            return "boolean"
167        if issubdtype(dtype, datetime64):
168            return "date-time"
169        return "string"
170
171    @staticmethod
172    def validate_format(excel_format: BaseModel, logger: logging.Logger) -> None:
173        """
174        Validates if the given format is of type ExcelFormat.
175
176        Args:
177            excel_format (Any): The format to be validated.
178
179        Raises:
180            ConfigValidationError: If the format is not ExcelFormat.
181        """
182        if not isinstance(excel_format, ExcelFormat):
183            logger.info(f"Expected ExcelFormat, got {excel_format}")
184            raise ConfigValidationError(FileBasedSourceError.CONFIG_VALIDATION_ERROR)
185
186    @staticmethod
187    def open_and_parse_file(fp: Union[IOBase, str, Path]) -> pd.DataFrame:
188        """
189        Opens and parses the Excel file.
190
191        Args:
192            fp: File pointer to the Excel file.
193
194        Returns:
195            pd.DataFrame: Parsed data from the Excel file.
196        """
197        return pd.ExcelFile(fp, engine="calamine").parse()  # type: ignore [arg-type, call-overload, no-any-return]