dynamodb_repository.py

from __future__ import annotations

import boto3
from functools import lru_cache
import dataclasses as dc
import uuid


def _reconstruct_params(cls, params: dict) -> dict:
    if not hasattr(cls, '__dc_field_names'):
        cls.__dc_field_names = tuple(f.name for f in dc.fields(cls))

    return {k:v for k,v in params.items() if k in cls.__dc_field_names}
  

@dc.dataclass(frozen=True)
class MyEntity():
    Id: uuid.UUID = dc.field(compare=False)
    name: str = dc.field()
    role: str = dc.field()
    
    @classmethod
    def reconstruct(cls, **kwargs) -> MyEntity:
        kwargs = _reconstruct_params(cls, kwargs)
        kwargs["Id"] = uuid.UUID(kwargs["Id"])
        return cls(**kwargs)

    def to_dict(self) -> dict:
        data = dc.asdict(self)

        return {
          **data,
          "Id": str(data["Id"]),
        }


class MyEntityRepository:
    def __init__(self, table_name: str, group_index_name: str) -> None:
        assert table_name, "expecting table name"
        assert group_index_name, "expecting index name"

        self._table_name = table_name
        self._group_index_name = group_index_name

    @lru_cache(maxsize=1, typed=True)
    def _table(self):
        return boto3.resource("dynamodb").Table(self._table_name)

    def load(self, param_name: str):
        kwargs = {
            "IndexName": self._group_index_name,
            "KeyConditionExpression": "#N = :N",
            "ExpressionAttributeNames": {
                "#N": "ParamName",
            },
            "ExpressionAttributeValues": {":N": param_name},
        }

        result = self._table().query(**kwargs)

        items = [MyEntity.reconstruct(**item) for item in result.get("Items", [])]

        return items

    def store(self, entries: List[MyEntity]):
        assert len(entries), 'expecting non-empty list of entries to store'

        with self._table().batch_writer() as batch:
            for entry in entries:
                item = {k:v for k,v in entry.to_dict().items() if v is not None}

                batch.put_item(Item=item)

    def delete(self, ids: List[str]):
        assert len(ids), 'expecting non-empty list of entries to delete'

        with self._table().batch_writer() as batch:
            for _id in ids:
                batch.delete_item(Key={'Id': _id})