Fix exclusive_start_key getting lost in PageIterator (pynamodb#421)

drewisme · jpinner-lyft · commit 2a1060996f8d · 2017-11-28T12:41:10.000-08:00
diff --git a/pynamodb/pagination.py b/pynamodb/pagination.py
@@ -13,7 +13,7 @@ def __init__(self, operation, args, kwargs):
         self._args = args
         self._kwargs = kwargs
         self._first_iteration = True
-        self._last_evaluated_key = None
+        self._last_evaluated_key = kwargs.get('exclusive_start_key')
         self._total_scanned_count = 0
 
     def __iter__(self):
@@ -110,11 +110,8 @@ def next(self):
 
     @property
     def last_evaluated_key(self):
-        if self._first_iteration:
-            # Not started iterating yet: there cannot be a last_evaluated_key
-            return None
-
-        if self._index == self._count:
+        if self._first_iteration or self._index == self._count:
+            # Not started iterating yet: return `exclusive_start_key` if set, otherwise expect None; or,
             # Entire page has been consumed: last_evaluated_key is whatever DynamoDB returned
             # It may correspond to the current item, or it may correspond to an item evaluated but not returned.
             return self.page_iter.last_evaluated_key
diff --git a/pynamodb/tests/test_model.py b/pynamodb/tests/test_model.py
@@ -2224,6 +2224,32 @@ def test_query_limit_greater_than_available_items_and_page_size(self):
             self.assertEquals(results_iter.total_count, 30)
             self.assertEquals(results_iter.page_iter.total_scanned_count, 60)
 
+    def test_query_with_exclusive_start_key(self):
+        with patch(PATCH_METHOD) as req:
+            req.return_value = MODEL_TABLE_DATA
+            UserModel('foo', 'bar')
+
+        with patch(PATCH_METHOD) as req:
+            items = []
+            for idx in range(30):
+                item = copy.copy(GET_MODEL_ITEM_DATA.get(ITEM))
+                item['user_id'] = {STRING_SHORT: 'id-{0}'.format(idx)}
+                items.append(item)
+
+            req.side_effect = [
+                {'Count': 10, 'ScannedCount': 10, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': items[19]['user_id']}},
+            ]
+            results_iter = UserModel.query('foo', limit=10, page_size=10, last_evaluated_key={'user_id': items[9]['user_id']})
+            self.assertEquals(results_iter.last_evaluated_key, {'user_id': items[9]['user_id']})
+
+            results = list(results_iter)
+            self.assertEqual(len(results), 10)
+            self.assertEqual(len(req.mock_calls), 1)
+            self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
+            self.assertEquals(results_iter.last_evaluated_key, {'user_id': items[19]['user_id']})
+            self.assertEquals(results_iter.total_count, 10)
+            self.assertEquals(results_iter.page_iter.total_scanned_count, 10)
+
     def test_query(self):
         """
         Model.query
