[CC-13705] Paginate ES report list

cc_dynamodb3/models.py

@@ -110,9 +110,14 @@ def _initial_data_to_dynamodb(cls, data):
         return dynamodb_data
 
     @classmethod
-    def all(cls):
-        for row, metadata in scan_all_in_table(cls.table()):
-            yield cls.from_row(row, metadata)
+    def all(cls, limit=None, paginate=False, exclusive_start_key=None):
+        if paginate:
+            for row, metadata, last_evaluated_key in scan_all_in_table(cls.table(), limit=limit, paginate=paginate,
+                                                                       exclusive_start_key=exclusive_start_key):
+                yield cls.from_row(row, metadata), last_evaluated_key
+        else:
+            for row, metadata in scan_all_in_table(cls.table()):
+                yield cls.from_row(row, metadata)
 
     @classmethod
     def paginated_query(cls, query_index=None, descending=False, limit=None, exclusive_start_key=None, filter_expression=None, **query_keys):

cc_dynamodb3/table.py

@@ -278,15 +278,22 @@ def scan_table(table_name_or_class, exclusive_start_key=None, **scan_kwargs):
 def _retrieve_all_matching(query_or_scan_func, *args, **kwargs):
     """Used by scan/query below."""
     limit = kwargs.pop('limit', None)
+    paginate = kwargs.pop('paginate', False)
     query_or_scan_kwargs = kwargs.copy()
+    if limit:
+        query_or_scan_kwargs['Limit'] = limit
+
     response = query_or_scan_func(*args, **query_or_scan_kwargs)
     total_found = 0
 
     # DynamoDB only returns up to 1MB of data per trip, so we need to keep querying or scanning.
     while True:
         metadata = response.get('ResponseMetadata', {})
         for row in response['Items']:
-            yield row, metadata
+            if paginate:
+                yield row, metadata, response.get('LastEvaluatedKey')
+            else:
+                yield row, metadata
             total_found += 1
             if limit and total_found == limit:
                 break

setup.py

@@ -10,7 +10,7 @@
         'schematics==1.1.1',
     ],
     tests_require=['pytest', 'mock', 'factory_boy', 'moto'],
-    version = '0.6.14',
+    version = '0.6.15',
     description = 'A dynamodb common configuration abstraction',
     author='Paul Craciunoiu',
     author_email='[email protected]',

tests/test_hash_only_model.py

@@ -106,4 +106,17 @@ def test_negative_timestamp():
 
     obj = HashOnlyModel.all().next()
     assert obj.created.year == long_ago.year
-    assert obj.item['created'] < 0
+    assert obj.item['created'] < 0
+
+
+def test_all_paginate():
+    HashOnlyModelFactory.create_table()
+    HashOnlyModelFactory(agency_subdomain='metzler', external_id=123)
+    HashOnlyModelFactory(agency_subdomain='metzler2', external_id=124)
+
+    obj, exclusive_start_key = HashOnlyModel.all(limit=1, paginate=True).next()
+    obj2, _ = HashOnlyModel.all(limit=1, paginate=True,
+                                exclusive_start_key=exclusive_start_key).next()
+
+    assert obj.external_id != obj2.external_id
+    assert {obj.external_id, obj2.external_id} == {123, 124}

tests/test_scan_query_all.py

@@ -34,3 +34,23 @@ def test_query_all_works_on_case_with_little_data():
 
     results = list(query_all_in_table(table, agency_id=1000))
     assert len(results) == 0
+
+
+def test_scan_all_paginate():
+    data = DYNAMODB_FIXTURES['nps_survey']
+    data_by_profile_id = {i['profile_id']: i for i in data}
+    table = mock_table_with_data('nps_survey', data)
+
+    results = list(scan_all_in_table(table, limit=1, paginate=True))
+    assert len(results) == 1
+
+    for result, metadata, last_evaluated_key in results:
+        item = data_by_profile_id[result.get('profile_id')]
+        assert item['agency_id'] == result.get('agency_id')
+        assert item['recommend_score'] == result.get('recommend_score')
+        assert item.get('favorite') == result.get('favorite')
+        assert last_evaluated_key
+
+    results2 = list(scan_all_in_table(table, limit=1, paginate=True,
+                                      exclusive_start_key=last_evaluated_key))
+    assert results2[0][0]['profile_id'] != results[0][0]['profile_id']