Merge pull request #10 from RSS-Engineering/support-pagination-dynamodb

Support pagination in DynamoDB

Author: explorigin

README.md

@@ -53,12 +53,15 @@ NOTE: Rule complexity is limited by the querying capabilities of the backend.
 
 ### DynamoDB
 
-`query(query_expr: Optional[Rule], filter_expr: Optional[Rule])` - Providing a
-  `query_expr` parameter will try to apply the keys of the expression to an
-  existing index. Providing a `filter_expr` parameter will filter the results of
+`query(query_expr: Optional[Rule], filter_expr: Optional[Rule], limit: Optional[str], exclusive_start_key: Optional[tuple[Any]], order: str = 'asc'`
+  - Providing a `query_expr` parameter will try to apply the keys of the expression to an
+  existing index.
+  - Providing a `filter_expr` parameter will filter the results of
   a passed `query_expr` or run a dynamodb `scan` if no `query_expr` is passed.
-  An empty call to `query()` will return the scan results (and be resource
+  - An empty call to `query()` will return the scan results (and be resource
   intensive).
+  - Providing a `limit` parameter will limit the number of results. If more results remain, the returned dataset will have an `last_evaluated_key` property that can be passed to `exclusive_start_key` to continue with the next page.
+  - Providing `order='desc'` will return the result set in descending order. This is not available for query calls that "scan" dynamodb.
 
 ## Backend Configuration Members
 

pydanticrud/backends/dynamodb.py

@@ -1,4 +1,4 @@
-from typing import Optional, Set
+from typing import Optional, Set, Any, Dict
 import logging
 import json
 from datetime import datetime
@@ -9,6 +9,7 @@
 from botocore.exceptions import ClientError
 from rule_engine import Rule, ast, types
 
+from ..main import IterableResult
 from ..exceptions import DoesNotExist, ConditionCheckFailed
 
 log = logging.getLogger(__name__)
@@ -116,9 +117,25 @@ def index_definition(index_name, keys, gsi=False):
     return schema
 
 
+class DynamoIterableResult(IterableResult):
+    def __init__(self, cls, result, serialized_items):
+        super(DynamoIterableResult, self).__init__(cls, serialized_items, result.get("Count"))
+
+        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"]
+
+
 class Backend:
     def __init__(self, cls):
         cfg = cls.Config
+        self.cls = cls
         self.schema = cls.schema()
         self.hash_key = cfg.hash_key
         self.range_key = getattr(cfg, 'range_key', None)
@@ -274,11 +291,22 @@ def exists(self):
         except ClientError:
             return False
 
-    def query(self, query_expr: Optional[Rule] = None, filter_expr: Optional[Rule] = None):
+    def query(self,
+              query_expr: Optional[Rule] = None,
+              filter_expr: Optional[Rule] = None,
+              limit: Optional[int] = None,
+              exclusive_start_key: Optional[tuple[Any]] = None,
+              order: str = 'asc',
+              ):
         table = self.get_table()
         f_expr, _ = rule_to_boto_expression(filter_expr) if filter_expr else (None, set())
 
         params = {}
+
+        if limit:
+            params["Limit"] = limit
+        if exclusive_start_key:
+            params["ExclusiveStartKey"] = self._key_param_to_dict(exclusive_start_key)
         if f_expr:
             params["FilterExpression"] = f_expr
 
@@ -291,6 +319,9 @@ def query(self, query_expr: Optional[Rule] = None, filter_expr: Optional[Rule] =
             index_name = self._get_best_index(keys_used)
             params["KeyConditionExpression"] = q_expr
 
+            if order != 'asc':
+                params["ScanIndexForward"] = False
+
             if index_name:
                 params["IndexName"] = index_name
             elif not keys_used.issubset({self.hash_key, self.range_key}):
@@ -305,16 +336,18 @@ def query(self, query_expr: Optional[Rule] = None, filter_expr: Optional[Rule] =
             except DynamoDBNeedsKeyConditionError:
                 raise ConditionCheckFailed("Non-key attributes are not valid in the query expression. Use filter "
                                            "expression")
-
         else:
+            if order != 'asc':
+                raise ConditionCheckFailed("Scans do not support reverse order.")
+
             try:
                 resp = table.scan(**params)
             except ClientError as e:
                 if e.response["Error"]["Code"] == "ResourceNotFoundException":
                     return []
                 raise e
 
-        return [self._deserialize_record(rec) for rec in resp["Items"]]
+        return DynamoIterableResult(self.cls, resp, (self._deserialize_record(rec) for rec in resp["Items"]))
 
     def get(self, key):
         _key = self._key_param_to_dict(key)

pydanticrud/main.py

@@ -12,6 +12,26 @@ def __new__(mcs, name, bases, namespace, **kwargs):
         return cls
 
 
+class IterableResult:
+    def __init__(self, cls, records, count=None):
+        self.records = [cls.parse_obj(i) for i in records]
+        self.count = count  # None indicates "unknown"
+
+        self._current_index = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            member = self.records[self._current_index]
+            self._current_index += 1
+            return member
+        except IndexError:
+            self._current_index = 0
+            raise StopIteration
+
+
 class BaseModel(PydanticBaseModel, metaclass=CrudMetaClass):
     @classmethod
     def initialize(cls):
@@ -28,7 +48,9 @@ def exists(cls) -> bool:
     @classmethod
     def query(cls, *args, **kwargs):
         res = cls.__backend__.query(*args, **kwargs)
-        return [cls.parse_obj(i) for i in res]
+        if not isinstance(res, IterableResult):
+            res = IterableResult(cls, res)
+        return res
 
     @classmethod
     def get(cls, *args, **kwargs):

tests/test_dynamodb.py

@@ -211,8 +211,14 @@ def simple_query_data(simple_table):
 
 @pytest.fixture(scope="module")
 def complex_query_data(complex_table):
-    presets = [dict()] * 20
-    data = [datum for datum in [complex_model_data_generator(**i) for i in presets]]
+    record_count = 500
+    presets = [dict()] * record_count
+    accounts = [str(uuid4()) for i in range(4)]
+
+    data = [
+        complex_model_data_generator(account=accounts[i % 4], **p)
+        for i, p in enumerate(presets)
+    ]
     for datum in data:
         ComplexKeyModel.parse_obj(datum).save()
     try:
@@ -234,6 +240,7 @@ def alias_query_data(alias_table):
         for datum in data:
             AliasKeyModel.delete(datum["name"])
 
+
 @pytest.fixture(scope="module")
 def nested_query_data(nested_table):
     presets = [dict()] * 5
@@ -249,7 +256,7 @@ def nested_query_data(nested_table):
 
 
 @pytest.fixture(scope="module")
-def nested_query_data_optional(nested_table):
+def nested_query_data_empty_ticket(nested_table):
     presets = [dict()] * 5
     data = [datum for datum in [nested_model_data_generator(include_ticket=False, **i) for i in presets]]
     for datum in data:
@@ -282,6 +289,14 @@ def test_query_with_hash_key_simple(dynamo, simple_query_data):
     assert res_data == {simple_query_data[0]["name"]: simple_query_data[0]}
 
 
+def test_scan_errors_with_order(dynamo, simple_query_data):
+    data_by_timestamp = simple_query_data[:]
+    data_by_timestamp.sort(key=lambda d: d["timestamp"])
+    with pytest.raises(ConditionCheckFailed,
+                       match=r"Scans do not support reverse order."):
+        SimpleKeyModel.query(order='desc')
+
+
 def test_query_errors_with_nonprimary_key_simple(dynamo, simple_query_data):
     data_by_timestamp = simple_query_data[:]
     data_by_timestamp.sort(key=lambda d: d["timestamp"])
@@ -347,13 +362,55 @@ def test_query_with_hash_key_complex(dynamo, complex_query_data):
     res_data = {(m.account, m.sort_date_key): m.dict() for m in res}
     assert res_data == {(record["account"], record["sort_date_key"]): record}
 
-    # Check that it works regardless of order
+    # Check that it works regardless of query attribute order
     res = ComplexKeyModel.query(
         Rule(f"sort_date_key == '{record['sort_date_key']}' and account == '{record['account']}'"))
     res_data = {(m.account, m.sort_date_key): m.dict() for m in res}
     assert res_data == {(record["account"], record["sort_date_key"]): record}
 
 
+@pytest.mark.parametrize('order', ('asc', 'desc'))
+def test_ordered_query_with_hash_key_complex(dynamo, complex_query_data, order):
+    middle_record = complex_query_data[(len(complex_query_data)//2)]
+    res = ComplexKeyModel.query(
+        Rule(f"account == '{middle_record['account']}' and sort_date_key >= '{middle_record['sort_date_key']}'"),
+        order=order
+    )
+    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')
+
+    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"])
@@ -381,15 +438,13 @@ def test_query_scan_complex(dynamo, complex_query_data):
 
 
 def test_query_with_nested_model(dynamo, nested_query_data):
-    data_by_expires = nested_model_data_generator()
-    res = NestedModel.query(filter_expr=Rule(f"expires <= '{data_by_expires['expires']}'"))
+    res = NestedModel.query()
     res_data = [m.ticket for m in res]
     assert any(elem is not None for elem in res_data)
 
 
-def test_query_with_nested_model_optional(dynamo, nested_query_data_optional):
-    data_by_expires = nested_model_data_generator(include_ticket=False)
-    res = NestedModel.query(filter_expr=Rule(f"expires <= '{data_by_expires['expires']}'"))
+def test_query_with_nested_model_optional(dynamo, nested_query_data_empty_ticket):
+    res = NestedModel.query()
     res_data = [m.ticket for m in res]
     assert any(elem is None for elem in res_data)
 
@@ -404,6 +459,7 @@ def test_query_alias_save(dynamo):
     except Exception as e:
         raise pytest.fail("Failed to save Alias model!")
 
+
 def test_get_alias_model_data(dynamo, alias_query_data):
     data = alias_model_data_generator()
     res = AliasKeyModel.get(alias_query_data[0]['name'])