Merge 7d6a223 into d0d21b5

dimagi · Jun 22, 2021 · 503e64c · 503e64c
2 parents d0d21b5 + 7d6a223
commit 503e64c
Showing 1 changed file with 15 additions and 10 deletions.
diff --git a/commcare_export/commcare_hq_client.py b/commcare_export/commcare_hq_client.py
@@ -126,7 +126,7 @@ def iterate(self, resource, paginator, params=None, checkpoint_manager=None):
         def iterate_resource(resource=resource, params=params):
             more_to_fetch = True
             last_batch_ids = set()
-            total_count = None
+            total_count = UNKNOWN_COUNT
             fetched = 0
             repeat_counter = 0
             last_params = None
@@ -141,35 +141,40 @@ def iterate_resource(resource=resource, params=params):
 
                 batch = self.get(resource, params)
                 last_params = copy.copy(params)
-                if not total_count or total_count == UNKNOWN_COUNT or fetched >= total_count:
-                    total_count = int(batch['meta']['total_count']) if batch['meta']['total_count'] else UNKNOWN_COUNT
+                batch_meta = batch['meta']
+                if total_count == UNKNOWN_COUNT or fetched >= total_count:
+                    if batch_meta.get('total_count'):
+                        total_count = int(batch_meta['total_count'])
+                    else:
+                        total_count = UNKNOWN_COUNT
                     fetched = 0
 
-                fetched += len(batch['objects'])
+                batch_objects = batch['objects']
+                fetched += len(batch_objects)
                 logger.debug('Received %s of %s', fetched, total_count)
-                if not batch['objects']:
+                if not batch_objects:
                     more_to_fetch = False
                 else:
                     got_new_data = False
-                    for obj in batch['objects']:
+                    for obj in batch_objects:
                         if obj['id'] not in last_batch_ids:
                             yield obj
                             got_new_data = True
 
-                    if batch['meta']['next']:
-                        last_batch_ids = {obj['id'] for obj in batch['objects']}
+                    if batch_meta.get('next'):
+                        last_batch_ids = {obj['id'] for obj in batch_objects}
                         params = paginator.next_page_params_from_batch(batch)
                         if not params:
                             more_to_fetch = False
                     else:
                         more_to_fetch = False
 
-                    limit = batch['meta'].get('limit')
+                    limit = batch_meta.get('limit')
                     if more_to_fetch:
                         repeated_last_page_of_non_counting_resource = (
                             not got_new_data
                             and total_count == UNKNOWN_COUNT
-                            and (limit and len(batch['objects']) < limit)
+                            and (limit and len(batch_objects) < limit)
                         )
                         more_to_fetch = not repeated_last_page_of_non_counting_resource