diff --git a/layercake/layercake/dynamodb.py b/layercake/layercake/dynamodb.py index c44d811..084fd84 100644 --- a/layercake/layercake/dynamodb.py +++ b/layercake/layercake/dynamodb.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Self, Type, TypedDict from urllib.parse import quote, unquote from uuid import UUID +from boto3.dynamodb.table import BatchWriter from boto3.dynamodb.types import TypeDeserializer, TypeSerializer from botocore.exceptions import ClientError from glom import glom @@ -744,6 +745,41 @@ class DynamoDBPersistenceLayer: else: return True + def batch_writer( + self, + table_name: str | None = None, + overwrite_by_pkeys: list[str] = [], + ) -> BatchWriter: + """Create a batch writer object. + + This method creates a context manager for writing + objects to Amazon DynamoDB in batch. + + The batch writer will automatically handle buffering and sending items + in batches. In addition, the batch writer will also automatically + handle any unprocessed items and resend them as needed. All you need + to do is call ``put_item`` for any items you want to add, and + ``delete_item`` for any items you want to delete. + + Example usage:: + + with table.batch_writer() as batch: + for _ in range(1000000): + batch.put_item(Item={'HashKey': '...', + 'Otherstuff': '...'}) + # You can also delete_items in a batch. + batch.delete_item(Key={'HashKey': 'SomeHashKey'}) + + See Also + -------- + DynamoDB.Table.batch_writer https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/table/batch_writer.html#DynamoDB.Table.batch_writer + """ + return BatchWriter( + table_name=table_name or self.table_name, + client=self.dynamodb_client, + overwrite_by_pkeys=overwrite_by_pkeys, + ) + class MissingError(Exception): pass @@ -990,7 +1026,7 @@ class DynamoDBCollection: else: head, tail = {}, items - def _getin(pair: KeyPair, obj: dict) -> dict: + def _extract_sk_values(pair: KeyPair, obj: dict) -> dict: obj = omit((PK, SK), obj) sk = pair[SK] path_spec = getattr(sk, 'path_spec', None) @@ -1001,7 +1037,7 @@ class DynamoDBCollection: return glom(obj, path_spec) return obj - def _mapkey(pair: KeyPair) -> str: + def _map_key(pair: KeyPair) -> str: pk = pair[PK] sk = pair[SK] @@ -1016,7 +1052,9 @@ class DynamoDBCollection: return key.removeprefix(sk.remove_prefix or '') return head | { - _mapkey(pair): _getin(pair, obj) for pair, obj in zip(sortkeys, tail) if obj + _map_key(pair): _extract_sk_values(pair, obj) + for pair, obj in zip(sortkeys, tail) + if obj } def query( diff --git a/layercake/pyproject.toml b/layercake/pyproject.toml index 80a8ae4..d0cc5be 100644 --- a/layercake/pyproject.toml +++ b/layercake/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "layercake" -version = "0.3.5" +version = "0.4.0" description = "Packages shared dependencies to optimize deployment and ensure consistency across functions." readme = "README.md" authors = [