diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 7d99a01881..8a46a9e9be 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -1982,6 +1982,11 @@ definitions: type: type: string enum: [ResponseToFileExtractor] + preserve_na_values: + title: Preserve NA Values + description: When enabled, string values such as "NA", "N/A", "NULL", "None" and "NaN" are kept as-is instead of being interpreted as missing and converted to null. Empty cells are still treated as null. Defaults to false to preserve historical behavior. + type: boolean + default: false $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py b/airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py index eea4b80b2a..f9a655db0c 100644 --- a/airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py +++ b/airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py @@ -27,9 +27,14 @@ class ResponseToFileExtractor(RecordExtractor): 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 a first iteration so we will only support CSV parsing using pandas as salesforce and sendgrid were doing. + + By default, pandas interprets strings such as "NA", "N/A", "null" and "NaN" as missing values and converts them to null. When + ``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 + behavior for existing connectors. """ parameters: InitVar[Mapping[str, Any]] + preserve_na_values: bool = False def __post_init__(self, parameters: Mapping[str, Any]) -> None: self.logger = logging.getLogger("airbyte") @@ -136,11 +141,20 @@ def _read_with_chunks( try: with open(path, "r", encoding=file_encoding) as data: + # When preserving NA values, disable pandas' default NA parsing so strings like "NA"/"N/A" are kept as-is, + # and explicitly map empty cells to None. Otherwise, keep the historical behavior where pandas converts its + # default NA tokens (including empty cells) to NaN, which we then map to None. + na_replacements = {nan: None, "": None} if self.preserve_na_values else {nan: None} chunks = pd.read_csv( - data, chunksize=chunk_size, iterator=True, dialect="unix", dtype=object + data, + chunksize=chunk_size, + iterator=True, + dialect="unix", + dtype=object, + keep_default_na=not self.preserve_na_values, ) for chunk in chunks: - chunk = chunk.replace({nan: None}).to_dict(orient="records") + chunk = chunk.replace(na_replacements).to_dict(orient="records") for row in chunk: yield row except pd.errors.EmptyDataError as e: diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 3bee30ca68..d7fead317d 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -503,6 +503,11 @@ class Config: class ResponseToFileExtractor(BaseModel): type: Literal["ResponseToFileExtractor"] + preserve_na_values: Optional[bool] = Field( + False, + description='When enabled, string values such as "NA", "N/A", "NULL", "None" and "NaN" are kept as-is instead of being interpreted as missing and converted to null. Empty cells are still treated as null. Defaults to false to preserve historical behavior.', + title="Preserve NA Values", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 1e5801bb06..5a75912b36 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2438,7 +2438,10 @@ def create_response_to_file_extractor( model: ResponseToFileExtractorModel, **kwargs: Any, ) -> ResponseToFileExtractor: - return ResponseToFileExtractor(parameters=model.parameters or {}) + return ResponseToFileExtractor( + parameters=model.parameters or {}, + preserve_na_values=model.preserve_na_values or False, + ) @staticmethod def create_exponential_backoff_strategy( diff --git a/unit_tests/sources/declarative/extractors/test_response_to_file_extractor.py b/unit_tests/sources/declarative/extractors/test_response_to_file_extractor.py index eeefe08174..6294db49d8 100644 --- a/unit_tests/sources/declarative/extractors/test_response_to_file_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_response_to_file_extractor.py @@ -15,6 +15,7 @@ class ResponseToFileExtractorTest(TestCase): def setUp(self) -> None: self._extractor = ResponseToFileExtractor({}) + self._preserving_extractor = ResponseToFileExtractor({}, preserve_na_values=True) self._http_mocker = requests_mock.Mocker() self._http_mocker.__enter__() @@ -39,6 +40,45 @@ def test_text_response_with_null_bytes(self) -> None: assert extracted_records == [{"FIRST_NAME": "a first name", "LAST_NAME": "a last name"}] + def test_na_string_values_are_preserved_when_enabled(self) -> None: + pandas_na_tokens = ["NA", "N/A", "NULL", "None", "nan", "null", "#N/A", ""] + csv_rows = "\n".join(f'"{token}","{token} value"' for token in pandas_na_tokens) + csv_with_na = f'"region","name"\n{csv_rows}\n' + response = self._mock_streamed_response(BytesIO(csv_with_na.encode("utf-8"))) + + extracted_records = list(self._preserving_extractor.extract_records(response)) + + assert len(extracted_records) == len(pandas_na_tokens) + for record, token in zip(extracted_records, pandas_na_tokens): + assert record["region"] == token, ( + f"Expected '{token}' to be preserved but got {record['region']}" + ) + assert record["name"] == f"{token} value" + + def test_na_string_values_are_converted_to_none_by_default(self) -> None: + pandas_na_tokens = ["NA", "N/A", "NULL", "None", "nan", "null", "#N/A", ""] + csv_rows = "\n".join(f'"{token}","preserved"' for token in pandas_na_tokens) + csv_with_na = f'"region","name"\n{csv_rows}\n' + response = self._mock_streamed_response(BytesIO(csv_with_na.encode("utf-8"))) + + extracted_records = list(self._extractor.extract_records(response)) + + assert len(extracted_records) == len(pandas_na_tokens) + for record in extracted_records: + assert record["region"] is None + assert record["name"] == "preserved" + + def test_empty_fields_are_converted_to_none_when_preserving(self) -> None: + csv_with_empty = '"region","name"\n"NA","North America"\n,""\n' + response = self._mock_streamed_response(BytesIO(csv_with_empty.encode("utf-8"))) + + extracted_records = list(self._preserving_extractor.extract_records(response)) + + assert extracted_records == [ + {"region": "NA", "name": "North America"}, + {"region": None, "name": None}, + ] + def _test_folder_path(self) -> Path: return Path(__file__).parent.resolve() diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index 5b39b069e5..5149f46d2b 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -5561,3 +5561,23 @@ def get_retriever(stream: Union[DeclarativeStream, DefaultStream]): if isinstance(stream, DeclarativeStream) else stream._stream_partition_generator._partition_factory._retriever ) + + +def test_create_response_to_file_extractor_preserve_na_values(): + from airbyte_cdk.sources.declarative.extractors import ResponseToFileExtractor + from airbyte_cdk.sources.declarative.models import ( + ResponseToFileExtractor as ResponseToFileExtractorModel, + ) + + factory = ModelToComponentFactory() + + default_component = factory.create_response_to_file_extractor( + ResponseToFileExtractorModel(type="ResponseToFileExtractor") + ) + assert isinstance(default_component, ResponseToFileExtractor) + assert default_component.preserve_na_values is False + + enabled_component = factory.create_response_to_file_extractor( + ResponseToFileExtractorModel(type="ResponseToFileExtractor", preserve_na_values=True) + ) + assert enabled_component.preserve_na_values is True