diff --git a/layercake/layercake/__init__.py b/layercake/layercake/__init__.py index c64f24a..8b13789 100644 --- a/layercake/layercake/__init__.py +++ b/layercake/layercake/__init__.py @@ -1,2 +1 @@ -def hello() -> str: - return "Hello from layercake!" + diff --git a/layercake/layercake/dynamodb.py b/layercake/layercake/dynamodb.py index 24426fd..c16ccfa 100644 --- a/layercake/layercake/dynamodb.py +++ b/layercake/layercake/dynamodb.py @@ -1,7 +1,8 @@ import os +from abc import ABC, abstractmethod from datetime import datetime from ipaddress import IPv4Address -from typing import Any, Type +from typing import Any, Type, TypedDict from aws_lambda_powertools import Logger from boto3.dynamodb.types import TypeDeserializer, TypeSerializer @@ -10,6 +11,11 @@ from botocore.exceptions import ClientError from .dateutils import now, timestamp TZ = os.getenv('TZ', 'UTC') +DELIMITER = os.getenv('DELIMITER', '#') +LIMIT = int(os.getenv('LIMIT', 25)) +PARTITION_KEY = os.getenv('PARTITION_KEY', 'pk') +SORT_KEY = os.getenv('SORT_KEY', 'sk') + logger = Logger(__name__) @@ -40,11 +46,11 @@ def deserialize(value: dict) -> dict: return {k: deserializer.deserialize(v) for k, v in value.items()} -def Key( +def ComposeKey( keyparts: str | tuple[str, ...], *, prefix: str | None = None, - delimiter: str = '#', + delimiter: str = DELIMITER, ) -> str: """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. @@ -67,14 +73,49 @@ def Key( return delimiter.join(keyparts) -class KeyPair(dict): +class PrimaryKey(ABC, dict): + @abstractmethod + def expr_attr_name(self) -> dict: ... + + @abstractmethod + def expr_attr_values(self) -> dict: ... + + +class PartitionKey(PrimaryKey): + """Represents a partition key for DynamoDB queries""" + + def __init__(self, pk: str) -> None: + super().__init__(**{PARTITION_KEY: pk}) + + def expr_attr_name(self) -> dict: + return {'#pk': PARTITION_KEY} + + def expr_attr_values(self) -> dict: + return {':pk': self[PARTITION_KEY]} + + +class KeyPair(PrimaryKey): + """Represents a composite key (partition key and sort key) for DynamoDB queries""" + def __init__(self, pk: str, sk: str) -> None: - super().__init__(id=pk, sk=sk) + super().__init__(**{PARTITION_KEY: pk, SORT_KEY: sk}) def __repr__(self) -> str: pk, sk = self.values() return f'KeyPair({pk!r}, {sk!r})' + def expr_attr_name(self) -> dict: + return { + '#pk': PARTITION_KEY, + '#sk': SORT_KEY, + } + + def expr_attr_values(self) -> dict: + return { + ':pk': self[PARTITION_KEY], + ':sk': self[SORT_KEY], + } + class TransactItems: """ @@ -440,20 +481,40 @@ class DynamoDBCollection: """ Example ------- - Get an item using a composed sort key + **Get an item using a composed sort key** collect = DynamoDBCollection(...) collect.get_item( - key=KeyPair( - pk='5OxmMjL-ujoR5IMGegQz', - sk=Key('sergio@somosbeta.com.br', prefix='emails'), + KeyPair( + 'b3511b5a-cb32-4833-a373-f8223f2088d4', + ComposeKey('sergio@somosbeta.com.br', prefix='emails'), ), ) + + **Get items using a composed partition key** + + collect = DynamoDBCollection(...) + collect.get_items( + PartitionKey( + ComposeKey('b3511b5a-cb32-4833-a373-f8223f2088d4', prefix='logs') + ), + ) + + **Get items using a key pair** + + collect = DynamoDBCollection(...) + collect.get_items( + KeyPair('b3511b5a-cb32-4833-a373-f8223f2088d4', 'emails), + ) """ class MissingError(ValueError): pass + class PaginatedResult(TypedDict): + items: list[dict] + last_key: str | None + def __init__( self, persistence_layer: DynamoDBPersistenceLayer, @@ -533,4 +594,37 @@ class DynamoDBCollection: expr_attr_values=expr_attr_values, ) - def get_items(self): ... + def get_items( + self, + key: PartitionKey | KeyPair, + *, + expr_attr_name: dict = {}, + expr_attr_values: dict = {}, + start_key: str | None = None, + filter_expr: str | None = None, + index_forward: bool = False, + limit: int = LIMIT, + ) -> PaginatedResult: + key_cond_expr = ( + '#pk = :pk AND begins_with(#sk, :sk)' + if isinstance(key, KeyPair) + else '#pk = :pk' + ) + + expr_attr_name.update(key.expr_attr_name()) + expr_attr_values.update(key.expr_attr_values()) + + response = self.persistence_layer.query( + key_cond_expr=key_cond_expr, + expr_attr_name=expr_attr_name, + expr_attr_values=expr_attr_values, + filter_expr=filter_expr, + index_forward=index_forward, + limit=limit, + # start_key=start_key if start_key else {}, + ) + + return { + 'items': response['items'], + 'last_key': response['last_key'] if 'last_key' in response else None, + } diff --git a/layercake/pyproject.toml b/layercake/pyproject.toml index 83b1948..33a5dcb 100644 --- a/layercake/pyproject.toml +++ b/layercake/pyproject.toml @@ -23,7 +23,12 @@ dependencies = [ ] [dependency-groups] -dev = ["pytest>=8.3.5", "pytest-cov>=6.0.0", "ruff>=0.11.1"] +dev = [ + "pytest>=8.3.5", + "pytest-cov>=6.0.0", + "pytest-env>=1.1.5", + "ruff>=0.11.1", +] [tool.pytest.ini_options] addopts = "--cov --cov-report html -v" diff --git a/layercake/pytest.ini b/layercake/pytest.ini new file mode 100644 index 0000000..47b8b40 --- /dev/null +++ b/layercake/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +env = + PARTITION_KEY=id diff --git a/layercake/tests/seeds.jsonl b/layercake/tests/seeds.jsonl index 1602661..4b76320 100644 --- a/layercake/tests/seeds.jsonl +++ b/layercake/tests/seeds.jsonl @@ -1,2 +1,4 @@ -{"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"}} +{"id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "0"}, "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"}, "email": {"S": "sergio@somosbeta.com.br"}, "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"}]}} +{"id": {"S": "5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "emails#sergio@somosbeta.com.br"}, "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"}, "email_primary": {"BOOL": true}, "mx_record_exists": {"BOOL": true}, "update_date": {"S": "2023-11-09T12:13:04.308986-03:00"}} +{"id": {"S": "logs#5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "2024-02-08T16:42:33.776409-03:00"}, "action": {"S": "OPEN_EMAIL"}} +{"id": {"S": "logs#5OxmMjL-ujoR5IMGegQz"}, "sk": {"S": "2019-03-25T00:00:00-03:00"}, "action": {"S": "CLICK_EMAIL"}} diff --git a/layercake/tests/test_dynamodb.py b/layercake/tests/test_dynamodb.py index dff1909..01fcf37 100644 --- a/layercake/tests/test_dynamodb.py +++ b/layercake/tests/test_dynamodb.py @@ -6,10 +6,11 @@ from botocore.exceptions import ClientError from layercake.dateutils import ttl from layercake.dynamodb import ( + ComposeKey, DynamoDBCollection, DynamoDBPersistenceLayer, - Key, KeyPair, + PartitionKey, TransactItems, serialize, ) @@ -31,12 +32,22 @@ def test_serialize(): } -def test_key(): - assert Key(('122', 'abc'), prefix='schedules') == 'schedules#122#abc' +def test_composekey(): + assert ComposeKey(('122', 'abc'), prefix='schedules') == 'schedules#122#abc' + assert ComposeKey(('122', 'abc')) == '122#abc' + assert ComposeKey('122') == '122' + + +def test_partitionkey(): + assert PartitionKey('123') == {'id': '123'} + assert PartitionKey('123').expr_attr_name() == {'#pk': 'id'} + assert PartitionKey('123').expr_attr_values() == {':pk': '123'} def test_keypair(): assert KeyPair('123', 'abc') == {'id': '123', 'sk': 'abc'} + assert KeyPair('123', 'abc').expr_attr_name() == {'#pk': 'id', '#sk': 'sk'} + assert KeyPair('123', 'abc').expr_attr_values() == {':pk': '123', ':sk': 'abc'} def test_transact_write_items( @@ -65,7 +76,7 @@ def test_collection_get_item( ): collect = DynamoDBCollection(dynamodb_persistence_layer) data_notfound = collect.get_item( - key=KeyPair( + KeyPair( pk='5OxmMjL-ujoR5IMGegQz', sk='tenant', ), @@ -76,9 +87,9 @@ def test_collection_get_item( # This item was added from seeds data = collect.get_item( - key=KeyPair( + KeyPair( pk='5OxmMjL-ujoR5IMGegQz', - sk=Key('sergio@somosbeta.com.br', prefix='emails'), + sk=ComposeKey('sergio@somosbeta.com.br', prefix='emails'), ), default={}, ) @@ -102,18 +113,18 @@ def test_collection_put_item( collect = DynamoDBCollection(dynamodb_persistence_layer) assert collect.put_item( - key=KeyPair( + KeyPair( '5OxmMjL-ujoR5IMGegQz', - Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'), + ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'), ), name='Beta Educação', ttl=ttl(days=3), ) data = collect.get_item( - key=KeyPair( + KeyPair( pk='5OxmMjL-ujoR5IMGegQz', - sk=Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'), + sk=ComposeKey('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'), ), ) @@ -131,8 +142,23 @@ def test_collection_delete_item( # This item was added from seeds assert collect.delete_item( - key=KeyPair( + KeyPair( '5OxmMjL-ujoR5IMGegQz', - Key('sergio@somsbeta.com.br', prefix='emails'), + ComposeKey('sergio@somsbeta.com.br', prefix='emails'), ) ) + + +def test_collection_get_items( + dynamodb_seeds, + dynamodb_persistence_layer: DynamoDBPersistenceLayer, +): + collect = DynamoDBCollection(dynamodb_persistence_layer) + + # This item was added from seeds + data = collect.get_items( + PartitionKey( + ComposeKey('5OxmMjL-ujoR5IMGegQz', prefix='logs'), + ), + ) + assert len(data['items']) == 2 diff --git a/layercake/uv.lock b/layercake/uv.lock index 182d071..dcba7c5 100644 --- a/layercake/uv.lock +++ b/layercake/uv.lock @@ -403,6 +403,7 @@ dependencies = [ dev = [ { name = "pytest" }, { name = "pytest-cov" }, + { name = "pytest-env" }, { name = "ruff" }, ] @@ -426,6 +427,7 @@ requires-dist = [ dev = [ { name = "pytest", specifier = ">=8.3.5" }, { name = "pytest-cov", specifier = ">=6.0.0" }, + { name = "pytest-env", specifier = ">=1.1.5" }, { name = "ruff", specifier = ">=0.11.1" }, ] @@ -620,6 +622,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/3b/48e79f2cd6a61dbbd4807b4ed46cb564b4fd50a76166b1c4ea5c1d9e2371/pytest_cov-6.0.0-py3-none-any.whl", hash = "sha256:eee6f1b9e61008bd34975a4d5bab25801eb31898b032dd55addc93e96fcaaa35", size = 22949 }, ] +[[package]] +name = "pytest-env" +version = "1.1.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/31/27f28431a16b83cab7a636dce59cf397517807d247caa38ee67d65e71ef8/pytest_env-1.1.5.tar.gz", hash = "sha256:91209840aa0e43385073ac464a554ad2947cc2fd663a9debf88d03b01e0cc1cf", size = 8911 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/b8/87cfb16045c9d4092cfcf526135d73b88101aac83bc1adcf82dfb5fd3833/pytest_env-1.1.5-py3-none-any.whl", hash = "sha256:ce90cf8772878515c24b31cd97c7fa1f4481cd68d588419fd45f10ecaee6bc30", size = 6141 }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0"