Improve first iteration logic (#1101)

Improving the first iteration logic in `PageIterator` and `ResultIterator`.

Both iterators had a `_first_iteration` state attribute that could be replaced with better logic:
- In `PageIterator`, it can be replaced with `_is_last_page` which
  - is exception-safe (only mutate after we do I/O), and
  - (IMO) more directly relates to the state we're handling
- In `ResultsIterator`, we can do without `_first_iteration` (but we should initialize `_index` and `_count`)

Inspired by #1059.

Author: ikonst

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

@@ -1347,6 +1347,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