Improve first iteration logic (#1121)

Backporting #1101 to 5.x branch to allow handling an exception raised when retrieving the first item.

Author: ikonst

docs/release_notes.rst

@@ -3,6 +3,28 @@
 Release Notes
 =============
 
+v5.3.3
+----------
+* Fix :py:class:`~pynamodb.pagination.PageIterator` and :py:class:`~pynamodb.pagination.ResultIterator`
+  to allow recovery from an exception when retrieving the first item (#1101).
+
+  .. code-block:: python
+
+    results = MyModel.query('hash_key')
+    while True:
+        try:
+            item = next(results)
+        except StopIteration:
+            break
+        except pynamodb.exceptions.QueryError as ex:
+            if ex.cause_response_code == 'ThrottlingException':
+                time.sleep(1)  # for illustration purposes only
+            else:
+                raise
+        else:
+            handle_item(item)
+
+
 v5.3.2
 ----------
 * Prevent ``typing_tests`` from being installed into site-packages (#1118)

pynamodb/__init__.py

@@ -7,4 +7,4 @@
 """
 __author__ = 'Jharrod LaFon'
 __license__ = 'MIT'
-__version__ = '5.3.2'
+__version__ = '5.3.3'

pynamodb/pagination.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Callable, Dict, Iterable, Iterator, TypeVar, Optional
+from typing import Any, Callable, Dict, Iterable, Iterator, Optional, TypeVar
 
 from pynamodb.constants import (CAMEL_COUNT, ITEMS, LAST_EVALUATED_KEY, SCANNED_COUNT,
                                 CONSUMED_CAPACITY, TOTAL, CAPACITY_UNITS)
@@ -90,8 +90,8 @@ def __init__(
         self._operation = operation
         self._args = args
         self._kwargs = kwargs
-        self._first_iteration = True
         self._last_evaluated_key = kwargs.get('exclusive_start_key')
+        self._is_last_page = False
         self._total_scanned_count = 0
         self._rate_limiter = None
         if rate_limit:
@@ -102,18 +102,17 @@ def __iter__(self) -> Iterator[_T]:
         return self
 
     def __next__(self) -> _T:
-        if self._last_evaluated_key is None and not self._first_iteration:
+        if self._is_last_page:
             raise StopIteration()
 
-        self._first_iteration = False
-
         self._kwargs['exclusive_start_key'] = self._last_evaluated_key
 
         if self._rate_limiter:
             self._rate_limiter.acquire()
             self._kwargs['return_consumed_capacity'] = TOTAL
         page = self._operation(*self._args, settings=self._settings, **self._kwargs)
         self._last_evaluated_key = page.get(LAST_EVALUATED_KEY)
+        self._is_last_page = self._last_evaluated_key is None
         self._total_scanned_count += page[SCANNED_COUNT]
 
         if self._rate_limiter:
@@ -170,10 +169,11 @@ def __init__(
         settings: OperationSettings = OperationSettings.default,
     ) -> None:
         self.page_iter: PageIterator = PageIterator(operation, args, kwargs, rate_limit, settings)
-        self._first_iteration = True
         self._map_fn = map_fn
         self._limit = limit
         self._total_count = 0
+        self._index = 0
+        self._count = 0
 
     def _get_next_page(self) -> None:
         page = next(self.page_iter)
@@ -189,10 +189,6 @@ def __next__(self) -> _T:
         if self._limit == 0:
             raise StopIteration
 
-        if self._first_iteration:
-            self._first_iteration = False
-            self._get_next_page()
-
         while self._index == self._count:
             self._get_next_page()
 
@@ -209,7 +205,7 @@ def next(self) -> _T:
 
     @property
     def last_evaluated_key(self) -> Optional[Dict[str, Dict[str, Any]]]:
-        if self._first_iteration or self._index == self._count:
+        if 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.

tests/test_model.py

@@ -1378,6 +1378,33 @@ def test_query_with_exclusive_start_key(self):
             self.assertEqual(results_iter.total_count, 10)
             self.assertEqual(results_iter.page_iter.total_scanned_count, 10)
 
+    def test_query_with_failure(self):
+        items = [
+            {
+                **GET_MODEL_ITEM_DATA[ITEM],
+                'user_id': {
+                    STRING: f'id-{idx}'
+                },
+            }
+            for idx in range(30)
+        ]
+
+        with patch(PATCH_METHOD) as req:
+            req.side_effect = [
+                Exception('bleep-bloop'),
+                {'Count': 10, 'ScannedCount': 10, 'Items': items[0:10], 'LastEvaluatedKey': {'user_id': items[10]['user_id']}},
+            ]
+            results_iter = UserModel.query('foo', limit=10, page_size=10)
+
+            with pytest.raises(Exception, match='bleep-bloop'):
+                next(results_iter)
+
+            first_item = next(results_iter)
+            assert first_item.user_id == 'id-0'
+
+            second_item = next(results_iter)
+            assert second_item.user_id == 'id-1'
+
     def test_query(self):
         """
         Model.query