diff --git a/complaint_search/defaults.py b/complaint_search/defaults.py index 366ab05..cace984 100644 --- a/complaint_search/defaults.py +++ b/complaint_search/defaults.py @@ -21,6 +21,7 @@ # The default result size matches the front-end default for users. # The trend_depth default limits display to 5 items in some Trends contexts. PAGINATION_BATCH = 100 +MAX_DOWNLOAD_SIZE = 100000 MAX_PAGINATION_DEPTH = 10000 RESULT_SIZE_DEFAULT = 25 RESULT_SIZE_OPTIONS = [10, 50, 100] diff --git a/complaint_search/export.py b/complaint_search/export.py index 4873a40..07e246f 100644 --- a/complaint_search/export.py +++ b/complaint_search/export.py @@ -4,6 +4,9 @@ from io import StringIO from django.http import StreamingHttpResponse +from django.http.response import Http404 + +from complaint_search.defaults import MAX_DOWNLOAD_SIZE class OpenSearchExporter(object): @@ -63,6 +66,9 @@ def stream(): # - total_count (int) # The total number of records to be output def export_json(self, scanResponse, total_count): + if not total_count or total_count > MAX_DOWNLOAD_SIZE: + return Http404 + def stream(): count = 0 # Write JSON diff --git a/complaint_search/tests/test_export.py b/complaint_search/tests/test_export.py index 1fb9ea9..8b2fb16 100644 --- a/complaint_search/tests/test_export.py +++ b/complaint_search/tests/test_export.py @@ -4,10 +4,12 @@ from collections import OrderedDict from django.http import StreamingHttpResponse +from django.http.response import Http404 from django.test import TestCase from parameterized import parameterized +from complaint_search.defaults import MAX_DOWNLOAD_SIZE from complaint_search.export import OpenSearchExporter @@ -76,6 +78,13 @@ def test_export_json_request_response(self, length): downloaded_file = io.BytesIO(b"".join(res.streaming_content)) self.assertFalse(downloaded_file is None) + def test_json_export_limit(self): + length = MAX_DOWNLOAD_SIZE + 1 + es_exporter = OpenSearchExporter() + gen = es_generator(length) + es_exporter.export_json(gen, length) + self.assertRaises(Http404) + class TestCSVExportWithUnicodeCharacters(TestCase): def test_export_contains_unicode_chacter(self):