add dynamodb collection

layercake/layercake/dynamodb.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from ipaddress import IPv4Address
-from typing import Any
+from typing import Any, Type
 
 from aws_lambda_powertools import Logger
 from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
@@ -9,53 +9,66 @@ from botocore.exceptions import ClientError
 logger = Logger(__name__)
 
 
-def _serialize_python_types(obj: Any) -> str | dict | list:
-    match obj:
+def _serialize_python_type(value: Any) -> str | dict | list:
+    match value:
         case datetime():
-            return obj.isoformat()
+            return value.isoformat()
         case IPv4Address():
-            return str(obj)
+            return str(value)
         case list() | tuple():
-            return [_serialize_python_types(v) for v in obj]
+            return [_serialize_python_type(v) for v in value]
         case dict():
-            return {k: _serialize_python_types(v) for k, v in obj.items()}
+            return {k: _serialize_python_type(v) for k, v in value.items()}
         case _:
-            return obj
+            return value
 
 
-def serialize(obj: dict) -> dict:
+def serialize(value: dict) -> dict:
     serializer = TypeSerializer()
-    return {k: serializer.serialize(_serialize_python_types(v)) for k, v in obj.items()}
+    return {
+        k: serializer.serialize(_serialize_python_type(v)) for k, v in value.items()
+    }
 
 
-def deserialize(obj: dict) -> dict:
+def deserialize(value: dict) -> dict:
     deserializer = TypeDeserializer()
-    return {k: deserializer.deserialize(v) for k, v in obj.items()}
+    return {k: deserializer.deserialize(v) for k, v in value.items()}
 
 
 def Key(
-    val: str | tuple[str, ...],
+    keyparts: str | tuple[str, ...],
     *,
     prefix: str | None = None,
     delimiter: str = '#',
 ) -> str:
-    if not prefix and not isinstance(val, tuple):
-        return val
+    """Creates a composite key by joining string parts with a specified delimiter.
+    If a prefix is provided, it is added at the beginning of the key parts.
 
-    if isinstance(val, str):
-        val = (val,)
+    Example
+    -------
+    >>> Key(('abc', 'xyz'), prefix='examples', delimiter='#')
+    'examples#abc#xyz'
+    """
+
+    if not prefix and not isinstance(keyparts, tuple):
+        return keyparts
+
+    if isinstance(keyparts, str):
+        keyparts = (keyparts,)
 
     if prefix:
-        val = (prefix,) + val
+        keyparts = (prefix,) + keyparts
 
-    return delimiter.join(val)
+    return delimiter.join(keyparts)
 
 
-def KeyPair(pk: str, sk: str) -> dict[str, str]:
-    return {
-        'id': pk,
-        'sk': sk,
-    }
+class KeyPair(dict):
+    def __init__(self, pk: str, sk: str) -> None:
+        super().__init__(id=pk, sk=sk)
+
+    def __repr__(self) -> str:
+        pk, sk = self.values()
+        return f'KeyPair({pk!r}, {sk!r})'
 
 
 class TransactItems:
@@ -416,3 +429,42 @@ class DynamoDBPersistenceLayer:
             raise
         else:
             return True
+
+
+class DynamoDBCollection:
+    class MissingError(ValueError):
+        pass
+
+    def __init__(
+        self,
+        persistence_layer: DynamoDBPersistenceLayer,
+        exception_cls: Type[ValueError] = MissingError,
+    ) -> None:
+        if not issubclass(exception_cls, ValueError):
+            raise TypeError(
+                f'exception_cls must be a subclass of ValueError, got {exception_cls}'
+            )
+
+        self.persistence_layer = persistence_layer
+        self.exception_cls = exception_cls
+
+    def get_item(
+        self,
+        key: KeyPair,
+        path_spec: str | None = None,
+        raise_if_missing: bool = True,
+        default: Any = None,
+        delimiter: str = '#',
+    ) -> Any:
+        exc_cls = self.exception_cls
+        data = self.persistence_layer.get_item(key)
+
+        if raise_if_missing and not data:
+            raise exc_cls(f'Item with {key} not found.')
+
+        if path_spec and data:
+            from glom import glom
+
+            return glom(data, path_spec, default=default)
+
+        return data or default

layercake/layercake/jsonl.py

@@ -0,0 +1,18 @@
+import json
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Generator
+
+
+@contextmanager
+def readlines(path: Path | str) -> Generator[Any, None, None]:
+    """Return the lines from a JSON."""
+    if isinstance(path, str):
+        path = Path(path)
+
+    if not path.exists():
+        yield iter(())
+        return None
+
+    with open(path) as fp:
+        yield (json.loads(line) for line in fp)

layercake/tests/conftest.py

@@ -1,7 +1,44 @@
 import boto3
 import pytest
 
+import layercake.jsonl as jsonl
+from layercake.dynamodb import DynamoDBPersistenceLayer
+
+table_name = 'pytest'
+dynamodb_endpoint_url = 'http://127.0.0.1:8000'
+
+
+@pytest.fixture
+def dynamodb_client():
+    client = boto3.client('dynamodb', endpoint_url=dynamodb_endpoint_url)
+    client.create_table(
+        AttributeDefinitions=[
+            {'AttributeName': 'id', 'AttributeType': 'S'},
+            {'AttributeName': 'sk', 'AttributeType': 'S'},
+        ],
+        TableName=table_name,
+        KeySchema=[
+            {'AttributeName': 'id', 'KeyType': 'HASH'},
+            {'AttributeName': 'sk', 'KeyType': 'RANGE'},
+        ],
+        ProvisionedThroughput={
+            'ReadCapacityUnits': 123,
+            'WriteCapacityUnits': 123,
+        },
+    )
+
+    yield client
+
+    client.delete_table(TableName=table_name)
+
 
 @pytest.fixture()
-def dynamodb_client():
-    return boto3.client('dynamodb', endpoint_url='http://127.0.0.1:8000')
+def dynamodb_persistence_layer(dynamodb_client) -> DynamoDBPersistenceLayer:
+    return DynamoDBPersistenceLayer(table_name, dynamodb_client)
+
+
+@pytest.fixture()
+def dynamodb_seeds(dynamodb_client):
+    with jsonl.readlines('tests/seeds.jsonl') as lines:
+        for line in lines:
+            dynamodb_client.put_item(TableName=table_name, Item=line)

layercake/tests/seeds.jsonl

@@ -0,0 +1,2 @@
+{"update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "email_verified": {"BOOL": true}, "cognito:sub": {"S": "58efed8d-d276-41a8-8502-4ab8b5a6415e"}, "cpf": {"S": "07879819908"}, "sk": {"S": "0"}, "email": {"S": "sergio@somosbeta.com.br"}, "id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "name": {"S": "S\u00e9rgio Rafael de Siqueira"}, "last_login": {"S": "2024-02-08T20:53:45.818126-03:00"}, "tenant:org_id": {"L": [{"S": "cJtK9SsnJhKPyxESe7g3DG"}, {"S": "edp8njvgQuzNkLx2ySNfAD"}, {"S": "8TVSi5oACLxTiT8ycKPmaQ"}]}}
+{"email_verified": {"BOOL": true}, "update_date": {"S": "2024-02-08T16:42:33.776409-03:00"}, "create_date": {"S": "2019-03-25T00:00:00-03:00"}, "sk": {"S": "emails#sergio@somosbeta.com.br"}, "id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "email_primary": {"BOOL": true}, "mx_record_exists": {"BOOL": true}, "update_date": {"S": "2023-11-09T12:13:04.308986-03:00"}}

layercake/tests/test_dynamodb.py

@@ -5,6 +5,7 @@ import pytest
 from botocore.exceptions import ClientError
 
 from layercake.dynamodb import (
+    DynamoDBCollection,
     DynamoDBPersistenceLayer,
     Key,
     KeyPair,
@@ -37,35 +38,57 @@ def test_keypair():
     assert KeyPair('123', 'abc') == {'id': '123', 'sk': 'abc'}
 
 
-def test_transact_write_items(dynamodb_client):
-    user_layer = DynamoDBPersistenceLayer('pytest', dynamodb_client)
-    transact = TransactItems(user_layer.table_name)
+def test_transact_write_items(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    transact = TransactItems(dynamodb_persistence_layer.table_name)
+    transact.put(item=KeyPair('5OxmMjL-ujoR5IMGegQz', '0'))
+    transact.put(item=KeyPair('cpf', '07879819908'))
     transact.put(
-        item={
-            'id': '5OxmMjL-ujoR5IMGegQz',
-            'sk': '0',
-        }
-    )
-    transact.put(
-        item={
-            'id': 'cpf',
-            'sk': '07879819908',
-        }
-    )
-    transact.put(
-        item={
-            'id': 'email',
-            'sk': 'sergio@somosbeta.com.br',
-        },
+        item=KeyPair('email', 'sergio@somosbeta.com.br'),
         cond_expr='attribute_not_exists(sk)',
     )
     transact.put(
-        item={
-            'id': '5OxmMjL-ujoR5IMGegQz',
-            'sk': 'email:sergio@somosbeta.com.br',
-        },
+        item=KeyPair('5OxmMjL-ujoR5IMGegQz', 'emails#sergio@somosbeta.com.br'),
         cond_expr='attribute_not_exists(sk)',
     )
 
     with pytest.raises(ClientError):
-        user_layer.transact_write_items(transact)
+        dynamodb_persistence_layer.transact_write_items(transact)
+
+
+def test_collection(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+    tenant_item = collect.get_item(
+        key=KeyPair(
+            pk='5OxmMjL-ujoR5IMGegQz',
+            sk=Key('tenant'),
+        ),
+        raise_if_missing=False,
+        default={},
+    )
+    assert tenant_item == {}
+
+    email_item = collect.get_item(
+        key=KeyPair(
+            pk='5OxmMjL-ujoR5IMGegQz',
+            sk=Key('sergio@somosbeta.com.br', prefix='emails'),
+        ),
+        default={},
+    )
+    assert email_item == {
+        'email_verified': True,
+        'mx_record_exists': True,
+        'sk': 'emails#sergio@somosbeta.com.br',
+        'email_primary': True,
+        'id': '5OxmMjL-ujoR5IMGegQz',
+        'create_date': '2019-03-25T00:00:00-03:00',
+        'update_date': '2023-11-09T12:13:04.308986-03:00',
+    }
+
+    with pytest.raises(DynamoDBCollection.MissingError):
+        collect.get_item(key=KeyPair('5OxmMjL-ujoR5IMGegQz', 'notfound'))

layercake/tests/test_jsonl.py

@@ -0,0 +1,19 @@
+import tempfile
+from pathlib import Path
+
+import layercake.jsonl as jsonl
+
+
+def test_readlines():
+    with tempfile.NamedTemporaryFile() as fp:
+        fp.writelines([b'{}\n' for _ in range(4)])
+        fp.seek(0)
+
+        with jsonl.readlines(fp.name) as lines:
+            assert sum(1 for _ in lines) == 4
+
+    with jsonl.readlines(Path('notfound.jsonl')) as lines:
+        assert sum(1 for _ in lines) == 0
+
+    with jsonl.readlines(Path('notfound.jsonl')) as lines:
+        assert list(lines) == []