Apply MAX_DOWNLOAD_SIZE to CSV exports

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

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

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")