From a98c4ec1370043606c8c6bc58b4b5bd75e7056a2 Mon Sep 17 00:00:00 2001 From: Thor Whalen <1906276+thorwhalen@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:12:48 +0000 Subject: [PATCH 1/2] feat: paginate scan/query; batch items()/values() through wrappers; _primary_key hook (#4) Co-Authored-By: Claude Sonnet 5.5 --- dynamodol/base.py | 176 +++++++++++++++++++++++++++-- dynamodol/partition_query.py | 29 ++--- tests/test_offline.py | 36 +++++- tests/test_pagination_and_bulk.py | 179 ++++++++++++++++++++++++++++++ 4 files changed, 390 insertions(+), 30 deletions(-) create mode 100644 tests/test_pagination_and_bulk.py diff --git a/dynamodol/base.py b/dynamodol/base.py index 18d782e..f21c07c 100644 --- a/dynamodol/base.py +++ b/dynamodol/base.py @@ -7,9 +7,11 @@ from functools import wraps from lazyprop import lazyprop from typing import Any, Tuple -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, ItemsView, Mapping, ValuesView -from dol import KvReader, KvPersister, Store +from contextlib import contextmanager + +from dol import KvReader, KvPersister, Store, inner_most_key class NoSuchKeyError(KeyError): @@ -146,6 +148,96 @@ def decimal_to_float(x): return x +def paginated(call, /, **kwargs): + """Yield the ``Items`` of a ``scan``/``query`` across *all* its pages. + + DynamoDB returns at most ~1 MB per call and says where it stopped + (``LastEvaluatedKey``); a single call silently truncates a bigger table. + """ + while True: + response = call(**kwargs) + yield from response["Items"] + last_key = response.get("LastEvaluatedKey") + if not last_key: + return + kwargs = {**kwargs, "ExclusiveStartKey": last_key} + + +def paginated_count(call, /, **kwargs) -> int: + """The ``Count`` of a ``Select='COUNT'`` ``scan``/``query`` summed over all its pages.""" + total = 0 + while True: + response = call(**kwargs) + total += response["Count"] + last_key = response.get("LastEvaluatedKey") + if not last_key: + return total + kwargs = {**kwargs, "ExclusiveStartKey": last_key} + + +BATCH_GET_MAX_KEYS = 100 # DynamoDB's limit for one BatchGetItem call + + +def _chunks(items, size): + items = list(items) + for i in range(0, len(items), size): + yield items[i : i + size] + + +def _leaf_reader_of(store): + """The ``DynamoDbBaseReader`` under any number of ``dol`` wrappers (``None`` if none).""" + seen = set() + while store is not None and id(store) not in seen: + if isinstance(store, DynamoDbBaseReader): + return store + seen.add(id(store)) + store = getattr(store, "store", None) + + +class _BatchedViewMixin: + """Iterate ``(key, value)`` pairs of the caller's store, a batch of keys per backend call. + + ``Mapping`` builds ``items()``/``values()`` from ``store[key]`` for every key, i.e. one + ``GetItem`` per row. Here the keys of a batch are resolved to the reader's own keys + (``dol.inner_most_key``), their items are read by one ``BatchGetItem``, and only while + ``store[key]`` runs are lookups of those keys answered from them. The lookups still go + through the caller's own store, so key/value codecs, ``filt_iter`` scopes and the + reader's prefix/partition scoping keep applying (i2mint/dynamodol#4); nothing is + cached while the caller's code runs between pairs. A key the batch did not return is + looked up with ``GetItem`` as usual. + """ + + def _pairs(self): + mapping = self._mapping + reader = _leaf_reader_of(mapping) + if reader is None: + for key in mapping: + yield key, mapping[key] + return + for keys in _chunks(mapping, BATCH_GET_MAX_KEYS): + inner_keys = ( + keys + if mapping is reader + else [inner_most_key(mapping, key, default=None) for key in keys] + ) + items = reader._batch_get(inner_keys) + for key in keys: + with reader._using_items(items): + value = mapping[key] + yield key, value + + +class _BatchedItemsView(_BatchedViewMixin, ItemsView): + def __iter__(self): + yield from self._pairs() + + +class _BatchedValuesView(_BatchedViewMixin, ValuesView): + def __iter__(self): + for _, value in self._pairs(): + yield value + + class _RemovedMethod: """Attribute that raises an ``AttributeError`` naming the method's replacement. @@ -316,7 +408,15 @@ def _keys_values_expression(self): ), } - def __getitem__(self, k): + ItemsView = _BatchedItemsView + ValuesView = _BatchedValuesView + _no_such_key_error = NoSuchKeyError # what a missing key raises (subclasses: their own) + + def _primary_key(self, k) -> dict: + """The item's primary key (``{key field: value}``) that the store key ``k`` names. + + Raises ``NoSuchKeyError`` for a key of the wrong shape. + """ try: _k = k if isinstance(k, str): @@ -325,28 +425,82 @@ def __getitem__(self, k): "If a sort key is defined, object keys must be tuples." ) _k = (k,) - _k = {att: key for att, key in zip(self.key_fields, _k)} + return {att: key for att, key in zip(self.key_fields, _k)} except (ValueError, TypeError) as e: raise NoSuchKeyError(f"Key not found: {k}") from e - item = get_item_or_raise(self.table, _k, k, **self._values_expression) - return self.format_get_item(item) + + def _ident(self, primary_key): + return tuple(primary_key[f] for f in self.key_fields) + + _prefetched_items = None # {identity: item} while a view looks up the items it just read + + def _batch_get(self, keys) -> dict: + """``{identity: item}`` for the store keys ``keys``, by ``BatchGetItem`` (<= 100 per call). + + Keys that name no item, or that the backend does not return, are simply absent + (they are looked up one by one later). + """ + wanted = {} + for k in keys: + try: + primary_key = self._primary_key(k) + wanted[self._ident(primary_key)] = primary_key + except (NoSuchKeyError, KeyError, TypeError): + continue + items = {} + for part in _chunks(wanted.values(), BATCH_GET_MAX_KEYS): + request = {"Keys": part, **self._keys_values_expression} + try: + response = self.db.batch_get_item(RequestItems={self.table_name: request}) + except Exception as e: + code = (getattr(e, "response", None) or {}).get("Error", {}).get("Code") + if code in KEY_CANNOT_NAME_AN_ITEM_ERROR_CODES: + continue + raise + for item in response.get("Responses", {}).get(self.table_name, []): + items[self._ident(item)] = item + return items + + @contextmanager + def _using_items(self, items): + previous, self._prefetched_items = self._prefetched_items, items + try: + yield + finally: + self._prefetched_items = previous + + def _item_of_primary_key(self, primary_key, k): + if self._prefetched_items: + item = self._prefetched_items.get(self._ident(primary_key)) + if item is not None: + return item + return get_item_or_raise( + self.table, + primary_key, + k, + error_cls=self._no_such_key_error, + **self._values_expression, + ) + + def __getitem__(self, k): + primary_key = self._primary_key(k) + return self.format_get_item(self._item_of_primary_key(primary_key, k)) def _scan_records(self, projection): """The raw records for a ``projection`` (an ``Expression*`` kwargs dict).""" - return self.table.scan(**projection)["Items"] + return list(paginated(self.table.scan, **projection)) iter_items = _RemovedMethod("iter_items", "dynamodol.scan_items(store)") iter_values = _RemovedMethod("iter_values", "dynamodol.scan_values(store)") def __iter__(self): # This is extremely inefficient and should not be used with large tables in production - response = self.table.scan(**self._keys_expression) - yield from (self.format_get_key(d) for d in response["Items"]) + for record in paginated(self.table.scan, **self._keys_expression): + yield self.format_get_key(record) def __len__(self): # This is extremely inefficient and should not be used with large tables in production - response = self.table.scan(Select="COUNT") - return response["Count"] + return paginated_count(self.table.scan, Select="COUNT") @staticmethod @wraps(get_db) diff --git a/dynamodol/partition_query.py b/dynamodol/partition_query.py index 292a8d4..6dc8680 100644 --- a/dynamodol/partition_query.py +++ b/dynamodol/partition_query.py @@ -5,6 +5,8 @@ from typing import Any from dynamodol.base import ( + paginated, + paginated_count, DynamoDbBaseReader, DynamoDbBasePersister, db_defaults, @@ -175,15 +177,16 @@ def filter_kwargs(self): return result def _scan_records(self, projection): - return self.table.query(**self.filter_kwargs, **projection)["Items"] + return list(paginated(self.table.query, **self.filter_kwargs, **projection)) def __iter__(self): - response = self.table.query(**self.filter_kwargs, **self._keys_expression) - yield from (self.format_get_key(d) for d in response["Items"]) + for record in paginated( + self.table.query, **self.filter_kwargs, **self._keys_expression + ): + yield self.format_get_key(record) def __len__(self): - response = self.table.query(**self.filter_kwargs, Select="COUNT") - return response["Count"] + return paginated_count(self.table.query, **self.filter_kwargs, Select="COUNT") @dataclass @@ -206,6 +209,7 @@ class DynamoDbPartitionReader(DynamoDbQueryReader): """ partition: str = field(default=None) + _no_such_key_error = NoSuchKeyError # this module's, as before def __post_init__(self): DynamoDbQueryReader.__post_init__(self) @@ -221,10 +225,8 @@ def format_get_key(self, item): """TODO: replace with _id_of_key, etc.""" return item[self.sort_key] - def __getitem__(self, k): - key = {self.partition_key: self.partition, self.sort_key: k} - item = get_item_or_raise(self.table, key, k, error_cls=NoSuchKeyError) - return self.format_get_item(item) + def _primary_key(self, k) -> dict: + return {self.partition_key: self.partition, self.sort_key: k} @dataclass @@ -253,13 +255,14 @@ def format_get_key(self, item): """TODO: replace with _id_of_key, etc.""" return item[self.sort_key][len(self.prefix) :] - def __getitem__(self, k): + def _primary_key(self, k) -> dict: try: - key = {self.partition_key: self.partition, self.sort_key: self.prefix + k} + return { + self.partition_key: self.partition, + self.sort_key: self.prefix + k, + } except TypeError as e: # a non-str key cannot carry the prefix raise NoSuchKeyError(f"Key not found: {k}") from e - item = get_item_or_raise(self.table, key, k, error_cls=NoSuchKeyError) - return self.format_get_item(item) class DynamoDbPartitionPersister(DynamoDbBasePersister, DynamoDbPartitionReader): diff --git a/tests/test_offline.py b/tests/test_offline.py index 806cf6f..dec8ed7 100644 --- a/tests/test_offline.py +++ b/tests/test_offline.py @@ -35,16 +35,19 @@ class FakeTable: ``delete_item`` so backend failures can be exercised without a live table. """ - def __init__(self, items=(), *, delete_error=None, get_error=None): + def __init__(self, items=(), *, delete_error=None, get_error=None, page_size=None): self.items = [dict(item) for item in items] self.delete_error = delete_error self.get_error = get_error + self.page_size = page_size # None: one page; else DynamoDB-style pagination + self.calls = [] # (operation, kwargs) of every scan/query/get_item/batch call @staticmethod def _matches(item, key): return all(item.get(k) == v for k, v in key.items()) def get_item(self, Key, **kwargs): + self.calls.append(("get_item", Key)) if self.get_error is not None: raise self.get_error for item in self.items: @@ -53,9 +56,16 @@ def get_item(self, Key, **kwargs): return {} def scan(self, **kwargs): + self.calls.append(("scan", kwargs)) + start = kwargs.get("ExclusiveStartKey", {}).get("offset", 0) + end = len(self.items) if not self.page_size else start + self.page_size + page = self.items[start:end] + response = {} + if end < len(self.items): + response["LastEvaluatedKey"] = {"offset": end} if kwargs.get("Select") == "COUNT": - return {"Count": len(self.items)} - return {"Items": [dict(item) for item in self.items]} + return {**response, "Count": len(page)} + return {**response, "Items": [dict(item) for item in page]} def query(self, **kwargs): return self.scan(**kwargs) @@ -92,9 +102,23 @@ def create_table(self, **kwargs): def Table(self, table_name): return self._table - -def _mk_persister(items=(), *, delete_error=None, get_error=None): - table = FakeTable(items, delete_error=delete_error, get_error=get_error) + def batch_get_item(self, RequestItems): + ((name, request),) = RequestItems.items() + assert len(request["Keys"]) <= 100, "BatchGetItem takes at most 100 keys" + self._table.calls.append(("batch_get_item", request["Keys"])) + found = [ + dict(item) + for key in request["Keys"] + for item in self._table.items + if self._table._matches(item, key) + ] + return {"Responses": {name: found}, "UnprocessedKeys": {}} + + +def _mk_persister(items=(), *, delete_error=None, get_error=None, page_size=None): + table = FakeTable( + items, delete_error=delete_error, get_error=get_error, page_size=page_size + ) store = DynamoDbBasePersister( db=FakeDb(table), table_name=TABLE_NAME, diff --git a/tests/test_pagination_and_bulk.py b/tests/test_pagination_and_bulk.py new file mode 100644 index 0000000..7b17dde --- /dev/null +++ b/tests/test_pagination_and_bulk.py @@ -0,0 +1,179 @@ +"""Pagination of scan/query, and batched ``items()``/``values()`` (i2mint/dynamodol#4).""" + +import pytest +from dol import KeyCodecs, Pipe, filt_iter, wrap_kvs + +import dynamodol +from dynamodol import ( + DynamoDbBasePersister, + DynamoDbPartitionPersister, + DynamoDbPrefixReader, +) + +from test_offline import DATA_FIELDS, KEY_FIELDS, TABLE_NAME, FakeDb, FakeTable + +N = 250 + + +def _rows(n=N): + return [{"key": f"k{i:04d}", "value": f"v{i}"} for i in range(n)] + + +def _base(rows=None, *, page_size=None): + table = FakeTable(_rows() if rows is None else rows, page_size=page_size) + store = DynamoDbBasePersister( + db=FakeDb(table), table_name=TABLE_NAME, key_fields=KEY_FIELDS, data_fields=DATA_FIELDS + ) + return store, table + + +def _prefix_reader(n=N, *, page_size=None): + rows = [ + {"pkey": "part1", "skey": f"v1/alice/{i:04d}", "value": f"A{i}"} for i in range(n) + ] + [{"pkey": "part1", "skey": f"v1/bob/{i:04d}", "value": f"B{i}"} for i in range(n)] + table = FakeTable(rows, page_size=page_size) + store = DynamoDbPrefixReader( + db=FakeDb(table), + table_name=TABLE_NAME, + key_fields=("pkey", "skey"), + data_fields=DATA_FIELDS, + partition="part1", + prefix="v1/", + ) + return store, table + + +def _count(table, op): + return sum(1 for name, _ in table.calls if name == op) + + +# ---- pagination ------------------------------------------------------------------- + + +def test_scan_reads_every_page(): + store, _ = _base(page_size=40) + assert len(list(store)) == N == len(store) + assert len(list(dynamodol.scan_items(store))) == N + assert len(list(dynamodol.scan_values(store))) == N + + +def test_query_reads_every_page(): + store, table = _prefix_reader(page_size=60) + assert len(list(store)) == 2 * N == len(store) # the fake query does not filter + + +def test_partition_persister_pages(): + table = FakeTable( + [{"pkey": "p", "skey": f"s{i:03d}", "value": "x"} for i in range(30)], page_size=7 + ) + store = DynamoDbPartitionPersister( + db=FakeDb(table), + table_name=TABLE_NAME, + key_fields=("pkey", "skey"), + data_fields=DATA_FIELDS, + partition="p", + ) + assert len(list(store)) == 30 == len(store) + + +# ---- batched items()/values() ------------------------------------------------------ + + +def test_bare_items_and_values_use_batches_not_one_get_per_key(): + store, table = _base() + expected = {r["key"]: r["value"] for r in _rows()} + assert dict(store.items()) == expected + assert sorted(store.values()) == sorted(expected.values()) + assert _count(table, "get_item") == 0 + assert _count(table, "batch_get_item") <= 2 * 3 # 250 keys = 3 calls, twice + + +@pytest.mark.parametrize( + "mk", + [ + lambda store: wrap_kvs(store, key_of_id=str.upper, id_of_key=str.lower), + lambda store: wrap_kvs(store, obj_of_data=lambda v: v.upper()), + lambda store: filt_iter(store, filt=lambda k: k.endswith("7")), + lambda store: Pipe( + filt_iter(filt=lambda k: k.endswith("7")), + wrap_kvs(key_of_id=str.upper, id_of_key=str.lower), + )(store), + ], + ids=["key codec", "value codec", "filt_iter", "filt + codec"], +) +def test_wrapped_items_are_correct_and_batched(mk): + store, table = _base() + view = mk(store) + per_key = {k: view[k] for k in view} # the one-get-per-key answer + table.calls.clear() + assert dict(view.items()) == per_key + assert sorted(view.values()) == sorted(per_key.values()) + assert _count(table, "get_item") == 0 and per_key + + +def test_prefix_reader_scope_is_applied_through_batches(): + store, table = _prefix_reader() + alice = Pipe(filt_iter.prefixes("alice/"), KeyCodecs.prefixed("alice/"))(store) + table.calls.clear() + got = dict(alice.items()) + assert len(got) == N and got["0007"] == "A7" + assert all(v.startswith("A") for v in got.values()) # bob's rows never come back + assert _count(table, "get_item") == 0 + + +def test_an_item_the_batch_did_not_return_is_looked_up_singly(): + store, table = _base(rows=_rows(5)) + original = table.items + + def batch_drops_k2(RequestItems): + ((name, request),) = RequestItems.items() + found = [dict(i) for key in request["Keys"] for i in original if table._matches(i, key)] + return {"Responses": {name: [i for i in found if i["key"] != "k0002"]}} + + store.db.batch_get_item = batch_drops_k2 + assert dict(store.items()) == {r["key"]: r["value"] for r in _rows(5)} + assert ("get_item", {"key": "k0002"}) in table.calls + + +def test_a_key_deleted_after_the_batch_is_a_snapshot_but_plain_lookups_are_exact(): + store, table = _base(rows=_rows(5)) + for key, _ in store.items(): # the batch was read when the iteration began + if key == "k0000": + table.items = [i for i in table.items if i["key"] != "k0003"] + with pytest.raises(KeyError): + store["k0003"] # a fresh lookup is exact + + +def test_batch_is_only_active_during_the_views_own_lookup(): + store, table = _base(rows=_rows(5)) + it = iter(store.items()) + next(it) + assert store._prefetched_items is None # nothing cached while the caller's code runs + table.items = [dict(i, value="CHANGED") for i in table.items] + assert store["k0001"] == "CHANGED" # plain lookups are never served from the batch + it.close() + + +def test_batch_get_validation_error_falls_back_to_plain_lookups(): + from test_offline import _client_error + + store, table = _base(rows=_rows(5)) + + def refuse(RequestItems): + raise _client_error("ValidationException", "BatchGetItem") + + store.db.batch_get_item = refuse + assert dict(store.items()) == {r["key"]: r["value"] for r in _rows(5)} + + +def test_batch_get_backend_error_propagates(): + from test_offline import _client_error + + store, _ = _base(rows=_rows(5)) + + def throttled(RequestItems): + raise _client_error("ProvisionedThroughputExceededException", "BatchGetItem") + + store.db.batch_get_item = throttled + with pytest.raises(Exception, match="Throughput"): + list(store.items()) From d18dd97fb2b81155d67042f4848d5cee156c93cc Mon Sep 17 00:00:00 2001 From: Thor Whalen <1906276+thorwhalen@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:20:14 +0000 Subject: [PATCH 2/2] fix: lazy key batching; thread-local batch (review of #8) Co-Authored-By: Claude Sonnet 5.5 --- dynamodol/base.py | 22 ++++++++++++++++++---- tests/test_pagination_and_bulk.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/dynamodol/base.py b/dynamodol/base.py index f21c07c..6031172 100644 --- a/dynamodol/base.py +++ b/dynamodol/base.py @@ -9,6 +9,7 @@ from typing import Any, Tuple from collections.abc import Iterable, ItemsView, Mapping, ValuesView +import threading from contextlib import contextmanager from dol import KvReader, KvPersister, Store, inner_most_key @@ -179,9 +180,15 @@ def paginated_count(call, /, **kwargs) -> int: def _chunks(items, size): - items = list(items) - for i in range(0, len(items), size): - yield items[i : i + size] + """Consecutive lists of at most ``size`` of the (lazily consumed) ``items``.""" + chunk = [] + for item in items: + chunk.append(item) + if len(chunk) >= size: + yield chunk + chunk = [] + if chunk: + yield chunk def _leaf_reader_of(store): @@ -432,7 +439,14 @@ def _primary_key(self, k) -> dict: def _ident(self, primary_key): return tuple(primary_key[f] for f in self.key_fields) - _prefetched_items = None # {identity: item} while a view looks up the items it just read + @property + def _prefetched_items(self): + """``{identity: item}`` while *this thread's* view looks up the items it just read.""" + return getattr(self.__dict__.setdefault("_local", threading.local()), "items", None) + + @_prefetched_items.setter + def _prefetched_items(self, items): + self.__dict__.setdefault("_local", threading.local()).items = items def _batch_get(self, keys) -> dict: """``{identity: item}`` for the store keys ``keys``, by ``BatchGetItem`` (<= 100 per call). diff --git a/tests/test_pagination_and_bulk.py b/tests/test_pagination_and_bulk.py index 7b17dde..44fb5eb 100644 --- a/tests/test_pagination_and_bulk.py +++ b/tests/test_pagination_and_bulk.py @@ -177,3 +177,33 @@ def throttled(RequestItems): store.db.batch_get_item = throttled with pytest.raises(Exception, match="Throughput"): list(store.items()) + + +def test_items_are_read_lazily_not_all_keys_and_a_full_count_up_front(): + store, table = _base(rows=_rows(6), page_size=1) + next(iter(store.items())) + scans = [c for c in table.calls if c[0] == "scan"] + assert len(scans) <= 6 and all(kw.get("Select") != "COUNT" for _, kw in scans) + + +def test_a_batch_in_one_thread_is_invisible_to_others_and_never_left_behind(): + import threading + + store, table = _base(rows=_rows(3)) + entered, release, seen = threading.Event(), threading.Event(), {} + + def hold(): # thread A: inside its lookup, holding its batch + with store._using_items({("k0000",): {"key": "k0000", "value": "STALE"}}): + entered.set() + release.wait(5) + + t = threading.Thread(target=hold) + t.start() + entered.wait(5) + seen["other_thread"] = store["k0000"] # must not be served from A's batch + release.set() + t.join() + assert seen["other_thread"] == "v0" and store._prefetched_items is None + del store["k0000"] # (the fake's put_item matches whole items, so replace by delete+put) + store["k0000"] = "NEW" + assert store["k0000"] == "NEW"