diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 97dda3b98..89fae5352 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -4247,7 +4247,7 @@ definitions: - "$ref": "#/definitions/JsonlDecoder" CsvDecoder: title: CSV - description: "Select 'CSV' for response data that is formatted as CSV (comma-separated values). Can specify an encoding (default: 'utf-8') and a delimiter (default: ',')." + description: "Select 'CSV' for response data that is formatted as CSV (comma-separated values). Can specify an encoding (default: 'utf-8'), a delimiter (default: ','), and a maximum field size (default: 2147483648)." type: object required: - type @@ -4261,6 +4261,11 @@ definitions: delimiter: type: string default: "," + max_field_size: + title: Maximum Field Size + description: Maximum number of characters allowed in a CSV field. + type: integer + default: 2147483648 set_values_to_none: type: array items: diff --git a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py index 0561369a7..8b0749188 100644 --- a/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +++ b/airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py @@ -160,6 +160,7 @@ class CsvParser(Parser): encoding: Optional[str] = "utf-8" delimiter: Optional[str] = "," set_values_to_none: Optional[List[str]] = None + max_field_size: Optional[int] = 2**31 def _get_delimiter(self) -> Optional[str]: """ @@ -175,12 +176,28 @@ def parse(self, data: BufferedIOBase) -> PARSER_OUTPUT_TYPE: """ Parse CSV data from decompressed bytes. """ - text_data = TextIOWrapper(data, encoding=self.encoding) # type: ignore - reader = csv.DictReader(text_data, delimiter=self._get_delimiter() or ",") - for row in reader: - if self.set_values_to_none: - row = {k: (None if v in self.set_values_to_none else v) for k, v in row.items()} - yield row + max_field_size = self.max_field_size if self.max_field_size is not None else 2**31 + previous_field_size_limit = csv.field_size_limit(max_field_size) + try: + text_data = TextIOWrapper(data, encoding=self.encoding) # type: ignore + reader = csv.DictReader(text_data, delimiter=self._get_delimiter() or ",") + try: + for row in reader: + if self.set_values_to_none: + row = { + k: (None if v in self.set_values_to_none else v) for k, v in row.items() + } + yield row + except csv.Error as exc: + if not str(exc).startswith("field larger than field limit"): + raise + raise AirbyteTracedException( + message=f"CSV field exceeds the configured maximum size of {max_field_size} characters.", + internal_message=str(exc), + failure_type=FailureType.config_error, + ) from exc + finally: + csv.field_size_limit(previous_field_size_limit) class CompositeRawDecoder(Decoder): diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index a84412eaf..eda9707f5 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -1377,6 +1377,11 @@ class CsvDecoder(BaseModel): type: Literal["CsvDecoder"] encoding: Optional[str] = "utf-8" delimiter: Optional[str] = "," + max_field_size: Optional[int] = Field( + 2147483648, + description="Maximum number of characters allowed in a CSV field.", + title="Maximum Field Size", + ) set_values_to_none: Optional[List[str]] = None diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index b992b4201..82e6c5e9c 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2825,6 +2825,7 @@ def _get_parser(model: BaseModel, config: Config) -> Parser: encoding=model.encoding, delimiter=model.delimiter, set_values_to_none=model.set_values_to_none, + max_field_size=model.max_field_size, ) elif isinstance(model, GzipDecoderModel): return GzipParser( diff --git a/unit_tests/sources/declarative/decoders/test_composite_decoder.py b/unit_tests/sources/declarative/decoders/test_composite_decoder.py index 8af9a4b4c..c54226913 100644 --- a/unit_tests/sources/declarative/decoders/test_composite_decoder.py +++ b/unit_tests/sources/declarative/decoders/test_composite_decoder.py @@ -14,6 +14,7 @@ import pytest import requests +from airbyte_cdk.models import FailureType from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import ( CompositeRawDecoder, CsvParser, @@ -298,6 +299,36 @@ def test_composite_raw_decoder_parse_empty_strings( assert parsed_records == expected_data +def test_csv_parser_default_max_field_size_allows_large_fields(): + large_field = "a" * 200_000 + + parsed_records = list(CsvParser().parse(BytesIO(f"value\n{large_field}\n".encode()))) + + assert parsed_records == [{"value": large_field}] + + +def test_csv_parser_max_field_size_override_is_honored(): + previous_limit = csv.field_size_limit() + + with pytest.raises(AirbyteTracedException) as exception: + list(CsvParser(max_field_size=10).parse(BytesIO(b"value\nthis value is too long\n"))) + + assert ( + exception.value.message == "CSV field exceeds the configured maximum size of 10 characters." + ) + assert exception.value.failure_type == FailureType.config_error + assert exception.value.internal_message == "field larger than field limit (10)" + assert csv.field_size_limit() == previous_limit + + +def test_csv_parser_restores_global_field_size_limit_after_parsing(): + previous_limit = csv.field_size_limit() + + list(CsvParser(max_field_size=10).parse(BytesIO(b"value\nshort\n"))) + + assert csv.field_size_limit() == previous_limit + + class TestServer(BaseHTTPRequestHandler): __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name