fix(file-based): Allow safe trailing empty CSV headers

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.

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: