add partition key to env var

layercake/layercake/__init__.py

@@ -1,2 +1 @@
-def hello() -> str:
+
-    return "Hello from layercake!"

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(
+            KeyPair(
-                pk='5OxmMjL-ujoR5IMGegQz',
+                'b3511b5a-cb32-4833-a373-f8223f2088d4',
-                sk=Key('sergio@somosbeta.com.br', prefix='emails'),
+                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,
+        }

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"

layercake/pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+env =
+    PARTITION_KEY=id

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"}]}}
+{"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"}]}}
-{"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": "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"}}

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():
+def test_composekey():
-    assert Key(('122', 'abc'), prefix='schedules') == 'schedules#122#abc'
+    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

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"