fix: restrict chunk size to 10MB per quota (#18438)

Author: miketheman

tests/unit/search/test_tasks.py

@@ -243,7 +243,7 @@ def project_docs(db):
         class TestError(Exception):
             pass
 
-        def parallel_bulk(client, iterable, index=None):
+        def parallel_bulk(client, iterable, index=None, max_chunk_bytes=None):
             assert client is es_client
             assert iterable is docs
             assert index == "warehouse-cbcbcbcbcb"
@@ -303,15 +303,19 @@ def project_docs(db):
         )
         monkeypatch.setattr(warehouse.search.tasks, "SearchLock", NotLock)
 
-        parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])
+        parallel_bulk = pretend.call_recorder(
+            lambda client, iterable, index, max_chunk_bytes: [None]
+        )
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
 
         reindex(task, db_request)
 
         assert parallel_bulk.calls == [
-            pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")
+            pretend.call(
+                es_client, docs, index="warehouse-cbcbcbcbcb", max_chunk_bytes=10485760
+            )
         ]
         assert es_client.indices.create.calls == [
             pretend.call(
@@ -367,15 +371,19 @@ def project_docs(db):
         )
         monkeypatch.setattr(warehouse.search.tasks, "SearchLock", NotLock)
 
-        parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])
+        parallel_bulk = pretend.call_recorder(
+            lambda client, iterable, index, max_chunk_bytes: [None]
+        )
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
 
         reindex(task, db_request)
 
         assert parallel_bulk.calls == [
-            pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")
+            pretend.call(
+                es_client, docs, index="warehouse-cbcbcbcbcb", max_chunk_bytes=10485760
+            )
         ]
         assert es_client.indices.create.calls == [
             pretend.call(
@@ -432,7 +440,9 @@ def project_docs(db):
         )
         monkeypatch.setattr(warehouse.search.tasks, "SearchLock", NotLock)
 
-        parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])
+        parallel_bulk = pretend.call_recorder(
+            lambda client, iterable, index, max_chunk_bytes: [None]
+        )
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
@@ -455,7 +465,9 @@ def project_docs(db):
         ]
 
         assert parallel_bulk.calls == [
-            pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")
+            pretend.call(
+                es_client, docs, index="warehouse-cbcbcbcbcb", max_chunk_bytes=10485760
+            )
         ]
         assert es_client.indices.create.calls == [
             pretend.call(

warehouse/search/tasks.py

@@ -165,7 +165,10 @@ def reindex(self, request):
                 request.db.execute(text("SET statement_timeout = '600s'"))
 
                 for _ in parallel_bulk(
-                    client, _project_docs(request.db), index=new_index_name
+                    client,
+                    _project_docs(request.db),
+                    index=new_index_name,
+                    max_chunk_bytes=10 * 1024 * 1024,  # 10MB, per OpenSearch defaults
                 ):
                     pass
             except:  # noqa