Skip to content

Commit 01d0b59

Browse files
jpinner-lyftjpinner-lyft
authored
Return the correct last_evaluated_key for limited queries/scans. Fixes pynamodb#406 (pynamodb#410)
1 parent 1828bda commit 01d0b59

File tree

5 files changed

+82
-18
lines changed

5 files changed

+82
-18
lines changed

‎pynamodb/connection/base.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,22 @@ def hash_keyname(self):
9595
break
9696
return self._hash_keyname
9797

98+
def get_key_names(self, index_name=None):
99+
"""
100+
Returns the names of the primary key attributes and index key attributes (if index_name is specified)
101+
"""
102+
key_names = [self.hash_keyname]
103+
if self.range_keyname:
104+
key_names.append(self.range_keyname)
105+
if index_name is not None:
106+
index_hash_keyname = self.get_index_hash_keyname(index_name)
107+
if index_hash_keyname not in key_names:
108+
key_names.append(index_hash_keyname)
109+
index_range_keyname = self.get_index_range_keyname(index_name)
110+
if index_range_keyname is not None and index_range_keyname not in key_names:
111+
key_names.append(index_range_keyname)
112+
return key_names
113+
98114
def get_index_hash_keyname(self, index_name):
99115
"""
100116
Returns the name of the hash key for a given index

‎pynamodb/connection/table.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,12 @@ def __init__(self,
2828
max_retry_attempts=max_retry_attempts,
2929
base_backoff_ms=base_backoff_ms)
3030

31+
def get_meta_table(self, refresh=False):
32+
"""
33+
Returns a MetaTable
34+
"""
35+
return self.connection.get_meta_table(self.table_name, refresh=refresh)
36+
3137
def delete_item(self, hash_key,
3238
range_key=None,
3339
condition=None,

‎pynamodb/pagination.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,16 @@ def __next__(self):
3535
def next(self):
3636
return self.__next__()
3737

38+
@property
39+
def key_names(self):
40+
# If the current page has a last_evaluated_key, use it to determine key attributes
41+
if self._last_evaluated_key:
42+
return self._last_evaluated_key.keys()
43+
44+
# Use the table meta data to determine the key attributes
45+
table_meta = self._operation.im_self.get_meta_table()
46+
return table_meta.get_key_names(self._kwargs.get('index_name'))
47+
3848
@property
3949
def page_size(self):
4050
return self._kwargs.get('limit')
@@ -100,7 +110,20 @@ def next(self):
100110

101111
@property
102112
def last_evaluated_key(self):
103-
return self.page_iter.last_evaluated_key
113+
if self._first_iteration:
114+
# Not started iterating yet: there cannot be a last_evaluated_key
115+
return None
116+
117+
if self._index == self._count:
118+
# Entire page has been consumed: last_evaluated_key is whatever DynamoDB returned
119+
# It may correspond to the current item, or it may correspond to an item evaluated but not returned.
120+
return self.page_iter.last_evaluated_key
121+
122+
# In the middle of a page of results: reconstruct a last_evaluated_key from the current item
123+
# The operation should be resumed starting at the last item returned, not the last item evaluated.
124+
# This can occur if the 'limit' is reached in the middle of a page.
125+
item = self._items[self._index - 1]
126+
return dict((key, item[key]) for key in self.page_iter.key_names)
104127

105128
@property
106129
def total_count(self):

‎pynamodb/tests/test_base_connection.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66
import six
77
from pynamodb.compat import CompatTestCase as TestCase
88
from pynamodb.connection import Connection
9+
from pynamodb.connection.base import MetaTable
910
from botocore.vendored import requests
1011
from pynamodb.exceptions import (VerboseClientError,
1112
TableError, DeleteError, UpdateError, PutError, GetError, ScanError, QueryError, TableDoesNotExist)
12-
from pynamodb.constants import DEFAULT_REGION, UNPROCESSED_ITEMS, STRING_SHORT, BINARY_SHORT, DEFAULT_ENCODING
13+
from pynamodb.constants import (
14+
DEFAULT_REGION, UNPROCESSED_ITEMS, STRING_SHORT, BINARY_SHORT, DEFAULT_ENCODING, TABLE_KEY)
1315
from pynamodb.expressions.operand import Path
1416
from pynamodb.tests.data import DESCRIBE_TABLE_DATA, GET_ITEM_DATA, LIST_TABLE_DATA
1517
from pynamodb.tests.deep_eq import deep_eq
@@ -25,6 +27,23 @@
2527
PATCH_METHOD = 'pynamodb.connection.Connection._make_api_call'
2628

2729

30+
class MetaTableTestCase(TestCase):
31+
"""
32+
Tests for the meta table class
33+
"""
34+
35+
def setUp(self):
36+
self.meta_table = MetaTable(DESCRIBE_TABLE_DATA.get(TABLE_KEY))
37+
38+
def test_get_key_names(self):
39+
key_names = self.meta_table.get_key_names()
40+
self.assertEqual(key_names, ["ForumName", "Subject"])
41+
42+
def test_get_key_names_index(self):
43+
key_names = self.meta_table.get_key_names("LastPostIndex")
44+
self.assertEqual(key_names, ["ForumName", "Subject", "LastPostDateTime"])
45+
46+
2847
class ConnectionTestCase(TestCase):
2948
"""
3049
Tests for the base connection class

‎pynamodb/tests/test_model.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2125,9 +2125,9 @@ def test_query_limit_less_than_available_items_multiple_page(self):
21252125
items.append(item)
21262126

21272127
req.side_effect = [
2128-
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': 'x'},
2129-
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': 'y'},
2130-
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': 'z'},
2128+
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': {'user_id': 'x'}},
2129+
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': 'y'}},
2130+
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': {'user_id': 'z'}},
21312131
]
21322132
results_iter = UserModel.query('foo', limit=25)
21332133
results = list(results_iter)
@@ -2136,7 +2136,7 @@ def test_query_limit_less_than_available_items_multiple_page(self):
21362136
self.assertEquals(req.mock_calls[0][1][1]['Limit'], 25)
21372137
self.assertEquals(req.mock_calls[1][1][1]['Limit'], 25)
21382138
self.assertEquals(req.mock_calls[2][1][1]['Limit'], 25)
2139-
self.assertEquals(results_iter.last_evaluated_key, 'z')
2139+
self.assertEquals(results_iter.last_evaluated_key, {'user_id': items[24]['user_id']})
21402140
self.assertEquals(results_iter.total_count, 30)
21412141
self.assertEquals(results_iter.page_iter.total_scanned_count, 60)
21422142

@@ -2153,9 +2153,9 @@ def test_query_limit_less_than_available_and_page_size(self):
21532153
items.append(item)
21542154

21552155
req.side_effect = [
2156-
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': 'x'},
2157-
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': 'y'},
2158-
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': 'z'},
2156+
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': {'user_id': 'x'}},
2157+
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': 'y'}},
2158+
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': {'user_id': 'x'}},
21592159
]
21602160
results_iter = UserModel.query('foo', limit=25, page_size=10)
21612161
results = list(results_iter)
@@ -2164,7 +2164,7 @@ def test_query_limit_less_than_available_and_page_size(self):
21642164
self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
21652165
self.assertEquals(req.mock_calls[1][1][1]['Limit'], 10)
21662166
self.assertEquals(req.mock_calls[2][1][1]['Limit'], 10)
2167-
self.assertEquals(results_iter.last_evaluated_key, 'z')
2167+
self.assertEquals(results_iter.last_evaluated_key, {'user_id': items[24]['user_id']})
21682168
self.assertEquals(results_iter.total_count, 30)
21692169
self.assertEquals(results_iter.page_iter.total_scanned_count, 60)
21702170

@@ -2181,8 +2181,8 @@ def test_query_limit_greater_than_available_items_multiple_page(self):
21812181
items.append(item)
21822182

21832183
req.side_effect = [
2184-
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': 'x'},
2185-
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': 'y'},
2184+
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': {'user_id': 'x'}},
2185+
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': 'y'}},
21862186
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30]},
21872187
]
21882188
results_iter = UserModel.query('foo', limit=50)
@@ -2209,8 +2209,8 @@ def test_query_limit_greater_than_available_items_and_page_size(self):
22092209
items.append(item)
22102210

22112211
req.side_effect = [
2212-
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': 'x'},
2213-
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': 'y'},
2212+
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': {'user_id': 'x'}},
2213+
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': 'y'}},
22142214
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30]},
22152215
]
22162216
results_iter = UserModel.query('foo', limit=50, page_size=10)
@@ -2545,9 +2545,9 @@ def test_scan_limit_with_page_size(self):
25452545
items.append(item)
25462546

25472547
req.side_effect = [
2548-
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': 'x'},
2549-
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': 'y'},
2550-
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': 'z'},
2548+
{'Count': 10, 'ScannedCount': 20, 'Items': items[:10], 'LastEvaluatedKey': {'user_id': 'x'}},
2549+
{'Count': 10, 'ScannedCount': 20, 'Items': items[10:20], 'LastEvaluatedKey': {'user_id': 'y'}},
2550+
{'Count': 10, 'ScannedCount': 20, 'Items': items[20:30], 'LastEvaluatedKey': {'user_id': 'z'}},
25512551
]
25522552
results_iter = UserModel.scan(limit=25, page_size=10)
25532553
results = list(results_iter)
@@ -2556,7 +2556,7 @@ def test_scan_limit_with_page_size(self):
25562556
self.assertEquals(req.mock_calls[0][1][1]['Limit'], 10)
25572557
self.assertEquals(req.mock_calls[1][1][1]['Limit'], 10)
25582558
self.assertEquals(req.mock_calls[2][1][1]['Limit'], 10)
2559-
self.assertEquals(results_iter.last_evaluated_key, 'z')
2559+
self.assertEquals(results_iter.last_evaluated_key, {'user_id': items[24]['user_id']})
25602560
self.assertEquals(results_iter.total_count, 30)
25612561
self.assertEquals(results_iter.page_iter.total_scanned_count, 60)
25622562

0 commit comments

Comments
 (0)