From fa4c39f59922a3bd9124b9755353f5d84860f80f Mon Sep 17 00:00:00 2001 From: Paul Craciunoiu Date: Fri, 14 Oct 2016 15:04:16 -0500 Subject: [PATCH] [CC-13705] Paginate ES report list --- cc_dynamodb3/models.py | 11 ++++++++--- cc_dynamodb3/table.py | 9 ++++++++- setup.py | 2 +- tests/test_hash_only_model.py | 15 ++++++++++++++- tests/test_scan_query_all.py | 20 ++++++++++++++++++++ 5 files changed, 51 insertions(+), 6 deletions(-) diff --git a/cc_dynamodb3/models.py b/cc_dynamodb3/models.py index f9998b6..fc519bd 100644 --- a/cc_dynamodb3/models.py +++ b/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): diff --git a/cc_dynamodb3/table.py b/cc_dynamodb3/table.py index 6822963..842ebf5 100644 --- a/cc_dynamodb3/table.py +++ b/cc_dynamodb3/table.py @@ -278,7 +278,11 @@ 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 @@ -286,7 +290,10 @@ def _retrieve_all_matching(query_or_scan_func, *args, **kwargs): 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 diff --git a/setup.py b/setup.py index d2f086c..1c2e4d8 100644 --- a/setup.py +++ b/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='pcraciunoiu@clearcareonline.com', diff --git a/tests/test_hash_only_model.py b/tests/test_hash_only_model.py index 6bf8137..fd51c36 100644 --- a/tests/test_hash_only_model.py +++ b/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 \ No newline at end of file + 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} diff --git a/tests/test_scan_query_all.py b/tests/test_scan_query_all.py index 9160b78..5211fd3 100644 --- a/tests/test_scan_query_all.py +++ b/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']