From 87b9fcae8cffb25bbd0eda0cabdb617b72ad6d82 Mon Sep 17 00:00:00 2001
From: Andy Chosak <andy.chosak@cfpb.gov>
Date: Fri, 24 Jul 2026 17:11:41 -0400
Subject: [PATCH] Apply MAX_DOWNLOAD_SIZE to CSV exports

Copies changes from b1429f0520d034f8f20ec847bd89ca3af928c5ae (and the
related fix in 0938dff2cce060652f58b36913a67848b576757c) to CSV exports,
not just JSON.
---
 complaint_search/es_interface.py      | 14 +++++++------
 complaint_search/export.py            | 29 ++++++++++++++++++---------
 complaint_search/tests/test_export.py | 12 +++++++++--
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/complaint_search/es_interface.py b/complaint_search/es_interface.py
index 7e609e1..23bb74c 100644
--- a/complaint_search/es_interface.py
+++ b/complaint_search/es_interface.py
@@ -347,14 +347,16 @@ def search(agg_exclude=None, **kwargs):
 
         exporter = OpenSearchExporter()
 
+        # Determine the total number of hits to enforce MAX_DOWNLOAD_SIZE.
+        if "highlight" in body:
+            del body["highlight"]
+        body.update({"size": 0, "track_total_hits": True})
+        count_res = _get_es().search(index=_COMPLAINT_ES_INDEX, body=body)
+        hit_total = count_res["hits"]["total"]["value"]
+
         if params.get("format") == "csv":
-            res = exporter.export_csv(scan_response, CSV_ORDERED_HEADERS)
+            res = exporter.export_csv(scan_response, CSV_ORDERED_HEADERS, hit_total)
         elif params.get("format") == "json":
-            if "highlight" in body:
-                del body["highlight"]
-            body.update({"size": 0, "track_total_hits": True})
-            count_res = _get_es().search(index=_COMPLAINT_ES_INDEX, body=body)
-            hit_total = count_res["hits"]["total"]["value"]
             res = exporter.export_json(scan_response, hit_total)
 
     return res
diff --git a/complaint_search/export.py b/complaint_search/export.py
index c676a34..1a5f94d 100644
--- a/complaint_search/export.py
+++ b/complaint_search/export.py
@@ -11,6 +11,20 @@
 
 
 class OpenSearchExporter(object):
+    def _check_download_size(self, total_count):
+        """Reject downloads if size exceeds MAX_DOWNLOAD_SIZE.
+
+        Raises ValidationError to be consistent with other DRF errors.
+        """
+        if total_count and total_count > MAX_DOWNLOAD_SIZE:
+            raise ValidationError(
+                {
+                    "size": [
+                        f"Result set of {total_count} exceeds the export limit of {MAX_DOWNLOAD_SIZE}"
+                    ]
+                }
+            )
+
     # export_csv - Stream an OpenSearch response as a CSV file
     #
     # Parameters:
@@ -19,7 +33,11 @@ class OpenSearchExporter(object):
     # - header_dict (OrderedDict)
     #   The ordered dictionary where the key is the OpenSearch field name
     #   and the value is the CSV column header for that field
-    def export_csv(self, scanResponse, header_dict):
+    # - total_count (int)
+    #   The total number of records to be output
+    def export_csv(self, scanResponse, header_dict, total_count):
+        self._check_download_size(total_count)
+
         def read_and_flush(writer, buffer_, row):
             writer.writerow(row)
             buffer_.seek(0)
@@ -66,14 +84,7 @@ def stream():
     # - total_count (int)
     #   The total number of records to be output
     def export_json(self, scanResponse, total_count):
-        if total_count and total_count > MAX_DOWNLOAD_SIZE:
-            raise ValidationError(
-                {
-                    "size": [
-                        f"Result set of {total_count} exceeds the export limit of {MAX_DOWNLOAD_SIZE}"
-                    ]
-                }
-            )
+        self._check_download_size(total_count)
 
         def stream():
             count = 0
diff --git a/complaint_search/tests/test_export.py b/complaint_search/tests/test_export.py
index 3a4fc94..e7e2f25 100644
--- a/complaint_search/tests/test_export.py
+++ b/complaint_search/tests/test_export.py
@@ -45,7 +45,7 @@ def test_export_csv_request_response(self, length):
         gen = es_generator(length)
 
         # act
-        res = es_exporter.export_csv(gen, TEST_HEADERS)
+        res = es_exporter.export_csv(gen, TEST_HEADERS, length)
 
         # assert
         self.assertTrue(isinstance(res, StreamingHttpResponse))
@@ -58,6 +58,14 @@ def test_export_csv_request_response(self, length):
         downloaded_file = io.BytesIO(b"".join(res.streaming_content))
         self.assertFalse(downloaded_file is None)
 
+    def test_csv_export_limit(self):
+        length = MAX_DOWNLOAD_SIZE + 1
+        es_exporter = OpenSearchExporter()
+        gen = es_generator(length)
+        with self.assertRaises(ValidationError) as context:
+            es_exporter.export_csv(gen, TEST_HEADERS, length)
+        self.assertEqual(context.exception.get_codes(), {"size": ["invalid"]})
+
     @parameterized.expand([[10], [5010], [100000]])
     def test_export_json_request_response(self, length):
         # arrange
@@ -110,6 +118,6 @@ def unicode_results():
             }
 
         exporter = OpenSearchExporter()
-        response = exporter.export_csv(unicode_results(), headers)
+        response = exporter.export_csv(unicode_results(), headers, 1)
         content = io.BytesIO(b"".join(response.streaming_content)).read()
         self.assertEqual(content, b"Key\r\n\xe2\x80\x99\r\n")