Merge pull request #33767 from dimagi/ml/add-chunksize-option

Add chunk size argument to load_domain_data
dimagi · Dec 13, 2023 · a8eb55c · a8eb55c
2 parents 1c666d8 + 43a15e2
commit a8eb55c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/corehq/apps/dump_reload/couch/load.py b/corehq/apps/dump_reload/couch/load.py
@@ -3,6 +3,7 @@
 
 from couchdbkit.exceptions import ResourceNotFound
 
+from corehq.apps.app_manager.models import Application, LinkedApplication, RemoteApp
 from corehq.apps.dump_reload.exceptions import DataExistsException
 from corehq.apps.dump_reload.interface import DataLoader
 from corehq.util.couch import (
@@ -53,7 +54,10 @@ def _get_db_for_doc_type(self, doc_type):
             if couch_db is None:
                 raise DocumentClassNotFound('No Document class with name "{}" could be found.'.format(doc_type))
             callback = LoaderCallback(self._success_counter, self.stdout)
-            db = IterDB(couch_db, new_edits=False, callback=callback)
+            chunksize = 100
+            if doc_type in [Application._doc_type, LinkedApplication._doc_type, RemoteApp._doc_type]:
+                chunksize = 1
+            db = IterDB(couch_db, new_edits=False, callback=callback, chunksize=chunksize)
             db.__enter__()
             self._dbs[doc_type] = db
         return self._dbs[doc_type]

diff --git a/corehq/util/couch.py b/corehq/util/couch.py
@@ -42,9 +42,8 @@ def _get_document_or_not_found_lite(cls, doc_id):
 def get_document_or_not_found(cls, domain, doc_id, additional_doc_types=None):
     allowed_doc_types = (additional_doc_types or []) + [cls.__name__]
     unwrapped = _get_document_or_not_found_lite(cls, doc_id)
-    if ((unwrapped.get('domain', None) != domain and
-         domain not in unwrapped.get('domains', [])) or
-        unwrapped['doc_type'] not in allowed_doc_types):
+    if ((unwrapped.get('domain', None) != domain and domain not in unwrapped.get('domains', []))
+       or unwrapped['doc_type'] not in allowed_doc_types):
         raise DocumentNotFound("Document {} of class {} not in domain {}!".format(
             doc_id,
             cls.__name__,
@@ -332,8 +331,8 @@ def _iter_update(doc_ids, try_num):
         if iter_db.error_ids:
             if try_num >= max_retries:
                 results.error_ids.update(iter_db.error_ids)
-                msg = ("The following documents did not correctly save:\n" +
-                       ", ".join(results.error_ids))
+                msg = ("The following documents did not correctly save:\n"
+                       + ", ".join(results.error_ids))
                 raise IterUpdateError(results, msg)
             else:
                 _iter_update(iter_db.error_ids, try_num + 1)
@@ -342,8 +341,8 @@ def _iter_update(doc_ids, try_num):
     if results.error_ids:
         msg = ("The following docs didn't correctly save.  Are you sure fn {} "
                "returned either None or an instance of DocUpdate?  Did you "
-               "change or remove the '_id' field?".format(fn.__name__) +
-               ", ".join(results.error_ids))
+               "change or remove the '_id' field?".format(fn.__name__)
+               + ", ".join(results.error_ids))
         raise IterUpdateError(results, msg)
 
     if verbose: