Ensure last_evaluated_key follows the same tuple-convention for primary keys.

Author: explorigin

pydanticrud/backends/dynamodb.py

@@ -120,7 +120,15 @@ def index_definition(index_name, keys, gsi=False):
 class DynamoIterableResult(IterableResult):
     def __init__(self, cls, result, serialized_items):
         super(DynamoIterableResult, self).__init__(cls, serialized_items, result.get("Count"))
-        self.last_evaluated_key = result.get("LastEvaluatedKey")
+
+        self.last_evaluated_key = None
+        lsk = result.get("LastEvaluatedKey")
+        if lsk:
+            _key = [lsk[cls.Config.hash_key]]
+            if cls.Config.range_key:
+                _key.append(lsk[cls.Config.range_key])
+            self.last_evaluated_key = tuple(_key)
+
         self.scanned_count = result["ScannedCount"]
 
 
@@ -288,7 +296,7 @@ def query(self,
               filter_expr: Optional[Rule] = None,
               limit: Optional[int] = None,
               exclusive_start_key: Optional[Dict[str, Any]] = None,
-              order: str = 'asc'
+              order: str = 'asc',
               ):
         table = self.get_table()
         f_expr, _ = rule_to_boto_expression(filter_expr) if filter_expr else (None, set())
@@ -298,7 +306,7 @@ def query(self,
         if limit:
             params["Limit"] = limit
         if exclusive_start_key:
-            params["ExclusiveStartKey"] = exclusive_start_key
+            params["ExclusiveStartKey"] = self._key_param_to_dict(exclusive_start_key)
         if f_expr:
             params["FilterExpression"] = f_expr
 

tests/test_dynamodb.py

@@ -211,9 +211,9 @@ def simple_query_data(simple_table):
 
 @pytest.fixture(scope="module")
 def complex_query_data(complex_table):
-    record_count = 20
+    record_count = 500
     presets = [dict()] * record_count
-    accounts = [str(uuid4()) for i in range(record_count//4)]
+    accounts = [str(uuid4()) for i in range(4)]
 
     data = [
         complex_model_data_generator(account=accounts[i % 4], **p)
@@ -386,6 +386,31 @@ def test_ordered_query_with_hash_key_complex(dynamo, complex_query_data, order):
     assert res_data == check_data
 
 
+@pytest.mark.parametrize('order', ('asc', 'desc'))
+def test_pagination_query_with_hash_key_complex(dynamo, complex_query_data, order):
+    page_size = 2
+    middle_record = complex_query_data[(len(complex_query_data)//2)]
+    query_rule = Rule(f"account == '{middle_record['account']}' and sort_date_key >= '{middle_record['sort_date_key']}'")
+    res = ComplexKeyModel.query(query_rule, order=order, limit=page_size)
+    res_data = [(m.account, m.sort_date_key) for m in res]
+    check_data = sorted([
+        (m["account"], m["sort_date_key"])
+        for m in complex_query_data
+        if m["account"] == middle_record['account'] and m["sort_date_key"] >= middle_record['sort_date_key']
+    ], reverse=order == 'desc')[:page_size]
+    assert res_data == check_data
+    assert res.last_evaluated_key == check_data[-1]
+
+    res = ComplexKeyModel.query(query_rule, order=order, limit=page_size, exclusive_start_key=res.last_evaluated_key)
+    res_data = [(m.account, m.sort_date_key) for m in res]
+    check_data = sorted([
+        (m["account"], m["sort_date_key"])
+        for m in complex_query_data
+        if m["account"] == middle_record['account'] and m["sort_date_key"] >= middle_record['sort_date_key']
+    ], reverse=order == 'desc')[page_size:page_size*2]
+    assert res_data == check_data
+
+
 def test_query_errors_with_nonprimary_key_complex(dynamo, complex_query_data):
     data_by_expires = complex_query_data[:]
     data_by_expires.sort(key=lambda d: d["expires"])