diff --git a/airbyte_cdk/sources/file_based/file_types/csv_parser.py b/airbyte_cdk/sources/file_based/file_types/csv_parser.py index 3b3dc4a0f7..b8c720642a 100644 --- a/airbyte_cdk/sources/file_based/file_types/csv_parser.py +++ b/airbyte_cdk/sources/file_based/file_types/csv_parser.py @@ -70,7 +70,9 @@ def read_data( file, file_read_mode, config_format.encoding, logger ) as fp: try: - headers = self._get_headers(fp, config_format, dialect_name) + headers, raw_headers = self._read_and_validate_headers( + fp, config_format, dialect_name + ) except UnicodeError: raise AirbyteTracedException( message=f"{FileBasedSourceError.ENCODING_ERROR.value} Expected encoding: {config_format.encoding}", @@ -84,9 +86,21 @@ def read_data( self._skip_rows(fp, rows_to_skip) lineno += rows_to_skip - reader = csv.DictReader(fp, dialect=dialect_name, fieldnames=headers) # type: ignore - for row in reader: + reader = csv.reader(fp, dialect=dialect_name) # type: ignore + trailing_empty_header_count = self._get_trailing_empty_header_count(raw_headers) + for values in reader: lineno += 1 + if not values: + continue + + values = self._strip_trailing_empty_header_values( + values, + headers, + trailing_empty_header_count, + file.uri, + lineno, + ) + row = self._row_to_dict(headers, values) # The row was not properly parsed if any of the values are None. This will most likely occur if there are more columns # than headers or more headers dans columns @@ -118,32 +132,114 @@ def read_data( def _get_headers(self, fp: IOBase, config_format: CsvFormat, dialect_name: str) -> List[str]: """Assumes the fp is pointing to the beginning of the files and will reset it as such.""" - # Note that this method assumes the dialect has already been registered if we're parsing the headers + headers, _ = self._read_and_validate_headers(fp, config_format, dialect_name) + return headers + + def _read_and_validate_headers( + self, fp: IOBase, config_format: CsvFormat, dialect_name: str + ) -> Tuple[List[str], List[str]]: + """Reads headers from file, validates, strips trailing empties, and resets fp to 0. + + Returns (clean_headers, raw_headers). For user-provided or autogenerated headers, + raw_headers == clean_headers (no trailing empties to strip). + """ if isinstance(config_format.header_definition, CsvHeaderUserProvided): - return config_format.header_definition.column_names + fp.seek(0) + columns = config_format.header_definition.column_names + return columns, columns if isinstance(config_format.header_definition, CsvHeaderAutogenerated): self._skip_rows( fp, config_format.skip_rows_before_header + config_format.skip_rows_after_header ) headers = self._auto_generate_headers(fp, dialect_name) - else: - # Then read the header - self._skip_rows(fp, config_format.skip_rows_before_header) - reader = csv.reader(fp, dialect=dialect_name) # type: ignore - headers = list(next(reader)) - - empty_count = sum(1 for h in headers if not h or h.isspace()) - if empty_count: - raise AirbyteTracedException( - message="CSV header row contains empty column name(s). Remove trailing delimiters or empty columns from the header row.", - internal_message=f"Found {empty_count} empty/whitespace-only column name(s) in header: {headers}", - failure_type=FailureType.config_error, - ) + fp.seek(0) + return headers, headers + + self._skip_rows(fp, config_format.skip_rows_before_header) + reader = csv.reader(fp, dialect=dialect_name) # type: ignore + raw_headers = list(next(reader)) + + headers = self._validate_trailing_headers(raw_headers) fp.seek(0) + return headers, raw_headers + + def _validate_trailing_headers(self, raw_headers: List[str]) -> List[str]: + """Validates header names and strips trailing empty headers.""" + trailing_empty_header_count = self._get_trailing_empty_header_count(raw_headers) + headers = ( + raw_headers[: len(raw_headers) - trailing_empty_header_count] + if trailing_empty_header_count + else raw_headers + ) + + if trailing_empty_header_count == len(raw_headers): + raise AirbyteTracedException( + message="CSV header row contains only empty column name(s).", + internal_message=f"Found {trailing_empty_header_count} empty/whitespace-only column name(s) in header: {raw_headers}", + failure_type=FailureType.config_error, + ) + + interior_empty_count = sum(1 for h in headers if self._is_empty_csv_value(h)) + if interior_empty_count: + raise AirbyteTracedException( + message="CSV header row contains interior empty column name(s).", + internal_message=f"Found {interior_empty_count} empty/whitespace-only column name(s) before trailing blank columns in header: {raw_headers}", + failure_type=FailureType.config_error, + ) + return headers + def _strip_trailing_empty_header_values( + self, + values: List[str], + headers: List[str], + trailing_empty_header_count: int, + filename: str, + lineno: int, + ) -> List[str]: + if not trailing_empty_header_count: + return values + + strip_start = len(headers) + strip_end = strip_start + trailing_empty_header_count + trailing_values = values[strip_start:strip_end] + non_empty_trailing_values = [ + value for value in trailing_values if not self._is_empty_csv_value(value) + ] + if non_empty_trailing_values: + raise AirbyteTracedException( + message="CSV data row contains non-empty value(s) under trailing empty header column(s).", + internal_message=f"Found {len(non_empty_trailing_values)} non-empty value(s) in stripped trailing column(s) on line {lineno} of file {filename}.", + failure_type=FailureType.config_error, + ) + + return values[:strip_start] + values[strip_end:] + + @staticmethod + def _row_to_dict(headers: List[str], values: List[str]) -> Dict[Any, Any]: + row: Dict[Any, Any] = dict(zip(headers, values)) + if len(values) > len(headers): + row[None] = values[len(headers) :] + elif len(values) < len(headers): + for header in headers[len(values) :]: + row[header] = None + return row + + @staticmethod + def _get_trailing_empty_header_count(headers: List[str]) -> int: + trailing_empty_header_count = 0 + for header in reversed(headers): + if not _CsvReader._is_empty_csv_value(header): + break + trailing_empty_header_count += 1 + return trailing_empty_header_count + + @staticmethod + def _is_empty_csv_value(value: Optional[str]) -> bool: + return not value or value.isspace() + def _auto_generate_headers(self, fp: IOBase, dialect_name: str) -> List[str]: """ Generates field names as [f0, f1, ...] in the same way as pyarrow's csv reader with autogenerate_column_names=True. diff --git a/unit_tests/sources/file_based/file_types/test_csv_parser.py b/unit_tests/sources/file_based/file_types/test_csv_parser.py index 3238a96bb3..f12f6613e8 100644 --- a/unit_tests/sources/file_based/file_types/test_csv_parser.py +++ b/unit_tests/sources/file_based/file_types/test_csv_parser.py @@ -648,7 +648,7 @@ def test_read_data_with_encoding_error(self) -> None: self._stream_reader.open_file.return_value = ( CsvFileBuilder().with_data(["something"]).build() ) - self._csv_reader._get_headers = Mock( + self._csv_reader._read_and_validate_headers = Mock( side_effect=UnicodeDecodeError("encoding", b"", 0, 1, "reason") ) @@ -657,7 +657,7 @@ def test_read_data_with_encoding_error(self) -> None: assert len(list(data_generator)) == 0 assert "encoding" in ate.value.message - assert self._csv_reader._get_headers.called + assert self._csv_reader._read_and_validate_headers.called def _read_data(self) -> Generator[Dict[str, str], None, None]: data_generator = self._csv_reader.read_data( @@ -791,10 +791,10 @@ def test_encoding_is_passed_to_stream_reader() -> None: @pytest.mark.parametrize( "header_row, expected_empty_count", [ - pytest.param("col1,col2,col3,,,", 3, id="trailing_empty_columns"), pytest.param("col1,,col3", 1, id="middle_empty_column"), pytest.param(",col2,col3", 1, id="leading_empty_column"), - pytest.param("col1,col2, ", 1, id="whitespace_only_column"), + pytest.param("col1, ,col3", 1, id="middle_whitespace_only_column"), + pytest.param(" , , ", 3, id="only_empty_columns"), ], ) def test_get_headers_raises_on_empty_column_names( @@ -825,6 +825,41 @@ def test_get_headers_raises_on_empty_column_names( csv.unregister_dialect(dialect_name) +@pytest.mark.parametrize( + "header_row, expected_headers", + [ + pytest.param("col1,col2,col3,,,", ["col1", "col2", "col3"], id="empty_headers"), + pytest.param( + "col1,col2,col3, , ", + ["col1", "col2", "col3"], + id="whitespace_only_headers", + ), + ], +) +def test_get_headers_strips_trailing_empty_column_names( + header_row: str, expected_headers: List[str] +) -> None: + csv_reader = _CsvReader() + config_format = CsvFormat() + fp = io.StringIO(header_row) + + dialect_name = f"test_{uuid4()}" + csv.register_dialect( + dialect_name, + delimiter=config_format.delimiter, + quotechar=config_format.quote_char, + escapechar=config_format.escape_char, + doublequote=config_format.double_quote, + quoting=csv.QUOTE_MINIMAL, + ) + + try: + headers = csv_reader._get_headers(fp, config_format, dialect_name) + assert headers == expected_headers + finally: + csv.unregister_dialect(dialect_name) + + def test_get_headers_accepts_valid_headers() -> None: csv_reader = _CsvReader() config_format = CsvFormat() @@ -847,7 +882,94 @@ def test_get_headers_accepts_valid_headers() -> None: csv.unregister_dialect(dialect_name) -def test_read_data_raises_on_empty_column_names() -> None: +@pytest.mark.parametrize( + "data, expected_records", + [ + pytest.param( + ["col1,col2,col3,,", "v1,v2,v3,,", "v4,v5,v6"], + [ + {"col1": "v1", "col2": "v2", "col3": "v3"}, + {"col1": "v4", "col2": "v5", "col3": "v6"}, + ], + id="trailing_empty_headers_and_cells", + ), + pytest.param( + ["col1,col2,col3, , ", "v1,v2,v3, ,", "v4,v5,v6"], + [ + {"col1": "v1", "col2": "v2", "col3": "v3"}, + {"col1": "v4", "col2": "v5", "col3": "v6"}, + ], + id="trailing_whitespace_headers_and_cells", + ), + pytest.param( + ["col1,col2,col3", "v1,v2,v3"], + [{"col1": "v1", "col2": "v2", "col3": "v3"}], + id="valid_headers", + ), + pytest.param( + ["col1,col2", "v1,v2", "", "v3,v4"], + [ + {"col1": "v1", "col2": "v2"}, + {"col1": "v3", "col2": "v4"}, + ], + id="blank_line_mid_file", + ), + pytest.param( + ["col1,col2", "v1,v2", ""], + [{"col1": "v1", "col2": "v2"}], + id="trailing_blank_line", + ), + ], +) +def test_read_data_strips_trailing_empty_headers_and_values( + data: List[str], expected_records: List[Dict[str, str]] +) -> None: + config_format = CsvFormat() + config = Mock() + config.name = "config_name" + config.format = config_format + + file = RemoteFile(uri="test.csv", last_modified=datetime.now()) + stream_reader = Mock(spec=AbstractFileBasedStreamReader) + logger = Mock(spec=logging.Logger) + csv_reader = _CsvReader() + + stream_reader.open_file.return_value = CsvFileBuilder().with_data(data).build() + + assert ( + list(csv_reader.read_data(config, file, stream_reader, logger, FileReadMode.READ)) + == expected_records + ) + + +@pytest.mark.parametrize( + "data", + [ + pytest.param(["col1,,col3", "v1,,v3"], id="interior_empty_header"), + pytest.param(["col1, ,col3", "v1, ,v3"], id="interior_whitespace_header"), + ], +) +def test_read_data_raises_on_interior_empty_column_names(data: List[str]) -> None: + config_format = CsvFormat() + config = Mock() + config.name = "config_name" + config.format = config_format + + file = RemoteFile(uri="test.csv", last_modified=datetime.now()) + stream_reader = Mock(spec=AbstractFileBasedStreamReader) + logger = Mock(spec=logging.Logger) + csv_reader = _CsvReader() + + stream_reader.open_file.return_value = CsvFileBuilder().with_data(data).build() + + with pytest.raises(AirbyteTracedException) as exc_info: + list(csv_reader.read_data(config, file, stream_reader, logger, FileReadMode.READ)) + + assert exc_info.value.failure_type == FailureType.config_error + assert "interior empty column name" in exc_info.value.message + + +def test_read_data_raises_on_non_empty_values_under_trailing_empty_headers() -> None: config_format = CsvFormat() config = Mock() config.name = "config_name" @@ -859,7 +981,7 @@ def test_read_data_raises_on_empty_column_names() -> None: csv_reader = _CsvReader() stream_reader.open_file.return_value = ( - CsvFileBuilder().with_data(["col1,col2,col3,,,", "v1,v2,v3,v4,v5,v6"]).build() + CsvFileBuilder().with_data(["col1,col2,col3,,", "v1,v2,v3,v4,"]).build() ) with pytest.raises(AirbyteTracedException) as exc_info: @@ -874,7 +996,7 @@ def test_read_data_raises_on_empty_column_names() -> None: ) assert exc_info.value.failure_type == FailureType.config_error - assert "empty column name" in exc_info.value.message + assert "non-empty value" in exc_info.value.message def test_parse_records_preserves_mismatch_error_detail() -> None: