improve get items

layercake/Makefile

@@ -6,3 +6,9 @@ build:
 
 deploy: export build
 	sam deploy --debug
+
+pytest:
+	uv run pytest
+
+htmlcov: pytest
+	uv run python -m http.server 80 -d htmlcov

layercake/layercake/dynamodb.py

@@ -14,7 +14,7 @@ from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
 from botocore.exceptions import ClientError
 from glom import glom
 
-from .dateutils import now, timestamp
+from .dateutils import timestamp
 from .funcs import omit, pick
 
 TZ = os.getenv('TZ', 'UTC')
@@ -142,14 +142,13 @@ if TYPE_CHECKING:
         ----------
         sk: str
             The sort key value.
-        table_name: str, optional
-            Optional name of the table associated with the sort key.
         path_spec: str, optional
             Optional specification for nested data extraction.
+        remove_prefix: str, optional
+            Optional prefix to remove from the key when forming the result dict.
         """
 
         sk: str
-        table_name: str | None = None
         path_spec: str | None = None
         remove_prefix: str | None = None
 else:
@@ -163,17 +162,16 @@ else:
         ----------
         sk: str
             The sort key value.
-        table_name: str, optional
-            Optional name of the table associated with the sort key.
         path_spec: str, optional
             Optional specification for nested data extraction.
+        remove_prefix: str, optional
+            Optional prefix to remove from the key when forming the result dict.
         """
 
         def __new__(
             cls,
             sk: str,
             *,
-            table_name: str | None = None,
             path_spec: str | None = None,
             remove_prefix: str | None = None,
         ) -> str:
@@ -183,37 +181,16 @@ else:
             self,
             sk: str,
             *,
-            table_name: str | None = None,
             path_spec: str | None = None,
             remove_prefix: str | None = None,
         ) -> None:
             # __init__ is used to store the parameters for later reference.
             # For immutable types like str, __init__ cannot change the instance's value.
             self.sk = sk
-            self.table_name = table_name
             self.path_spec = path_spec
             self.remove_prefix = remove_prefix
 
 
-@dataclass
-class TransactKey:
-    """
-    Example
-    -------
-        TransactKey('e9bb7dc6-c7b2-4d34-8931-d298353758ec')
-            + SortKey('0')
-            + SortKey('tenant')
-    """
-
-    pk: str
-    sk: tuple[SortKey, ...] = ()
-
-    def __add__(self, sk: SortKey) -> 'TransactKey':
-        if not isinstance(sk, SortKey):
-            raise TypeError('Can only add a SortKey to a TransactKey')
-        return TransactKey(pk=self.pk, sk=self.sk + (sk,))
-
-
 class Key(ABC, dict):
     @abstractmethod
     def expr_attr_name(self) -> dict: ...
@@ -274,6 +251,41 @@ class KeyPair(Key):
 
         return cls(*pair)
 
+    def __add__(self, other: Self) -> 'KeyChain':
+        return KeyChain(pairs=(self, other))
+
+
+@dataclass(frozen=True)
+class KeyChain:
+    pairs: tuple[KeyPair, ...] = ()
+
+    def __add__(self, other: KeyPair) -> 'KeyChain':
+        if not isinstance(other, KeyPair):
+            raise TypeError('Can only add a KeyPair to a KeyChain')
+
+        return KeyChain(pairs=self.pairs + (other,))
+
+
+@dataclass(frozen=True)
+class TransactKey:
+    """
+    Example
+    -------
+        TransactKey('e9bb7dc6-c7b2-4d34-8931-d298353758ec')
+            + SortKey('0')
+            + SortKey('tenant')
+    """
+
+    pk: str
+    pairs: tuple[KeyPair, ...] = ()
+
+    def __add__(self, other: SortKey) -> 'TransactKey':
+        if not isinstance(other, SortKey):
+            raise TypeError('Can only add a SortKey to a TransactKey')
+
+        pair = KeyPair(self.pk, other)
+        return TransactKey(pk=self.pk, pairs=self.pairs + (pair,))
+
 
 class TransactionCanceledReason(TypedDict):
     code: str
@@ -481,6 +493,10 @@ class DynamoDBPersistenceLayer:
         self.table_name = table_name
         self.dynamodb_client = dynamodb_client
 
+    @property
+    def collect(self) -> 'DynamoDBCollection':
+        return DynamoDBCollection(self)
+
     def query(
         self,
         *,
@@ -810,13 +826,10 @@ class DynamoDBCollection:
         bool
             True if the operation is successful, False otherwise.
         """
-        now_ = now(self.tz)
-
         if isinstance(ttl, int):
             kwargs.update(
                 {
                     'ttl': ttl,
-                    'ttl_date': datetime.fromtimestamp(ttl, now_.tzinfo),
                 }
             )
 
@@ -824,7 +837,6 @@ class DynamoDBCollection:
             kwargs.update(
                 {
                     'ttl': timestamp(ttl),
-                    'ttl_date': ttl,
                 }
             )
 
@@ -865,7 +877,7 @@ class DynamoDBCollection:
 
     def get_items(
         self,
-        key: TransactKey,
+        key: TransactKey | KeyChain,
         flatten_top: bool = True,
     ) -> dict[str, Any]:
         """Get multiple items via a transaction based on the provided TransactKey.
@@ -899,19 +911,16 @@ class DynamoDBCollection:
         """
 
         # If no sort key is provided, the query is skipped
-        if not key.sk:
+        if not key.pairs:
             return {}
 
         table_name = self.persistence_layer.table_name
-        sortkeys = key.sk[1:] if flatten_top else key.sk
+        sortkeys = key.pairs[1:] if flatten_top else key.pairs
         transact = TransactItems(table_name)
 
-        # Add a get operation for each sort key for the transaction
+        # Add a get operation for each key for the transaction
-        for sk in key.sk:
+        for pair in key.pairs:
-            transact.get(
+            transact.get(key=pair)
-                key=KeyPair(key.pk, sk),
-                table_name=sk.table_name,
-            )
 
         items = self.persistence_layer.transact_get_items(transact)
 
@@ -920,20 +929,29 @@ class DynamoDBCollection:
         else:
             head, tail = {}, items
 
-        def _getin(sk: SortKey, v: dict) -> dict:
+        def _getin(pair: KeyPair, v: dict) -> dict:
             v = omit((PK, SK), v)
+            sk = pair[SK]
+            path_spec = getattr(sk, 'path_spec', None)
 
-            if sk.path_spec:
+            if path_spec:
                 from glom import glom
 
-                return glom(v, sk.path_spec)
+                return glom(v, path_spec)
             return v
 
-        def _removeprefix(sk: SortKey) -> str:
+        def _removeprefix(pair: KeyPair) -> str:
-            return sk.removeprefix(sk.remove_prefix) if sk.remove_prefix else sk
+            sk = pair[SK]
+
+            if not isinstance(sk, SortKey):
+                return pair[PK]
+
+            return sk.removeprefix(sk.remove_prefix or '')
 
         return head | {
-            _removeprefix(k): _getin(k, item) for k, item in zip(sortkeys, tail) if item
+            _removeprefix(pair): _getin(pair, item)
+            for pair, item in zip(sortkeys, tail)
+            if item
         }
 
     def query(

layercake/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "layercake"
-version = "0.2.21"
+version = "0.3.0"
 description = "Packages shared dependencies to optimize deployment and ensure consistency across functions."
 readme = "README.md"
 authors = [
@@ -32,16 +32,22 @@ dev = [
     "jsonlines>=4.0.0",
     "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"
 
+[tool.ruff]
+target-version = "py311"
+src = ["app"]
+
 [tool.ruff.format]
 quote-style = "single"
 
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

layercake/pytest.ini

@@ -1,5 +0,0 @@
-[pytest]
-env =
-    DYNAMODB_ENDPOINT_URL=http://127.0.0.1:8000
-    DYNAMODB_PARTITION_KEY=id
-    DYNAMODB_SORT_KEY=sk

layercake/tests/conftest.py

@@ -4,13 +4,18 @@ import boto3
 import jsonlines
 import pytest
 
-from layercake.dynamodb import DynamoDBPersistenceLayer
+PYTEST_TABLE_NAME = 'pytest'
+DYNAMODB_ENDPOINT_URL = 'http://localhost:8000'
+PK = 'id'
+SK = 'sk'
 
-PYTEST_TABLE_NAME = os.getenv('PYTEST_TABLE_NAME', 'pytest')
+
-# Check `pytest.ini` for more details.
+# https://docs.pytest.org/en/7.1.x/reference/reference.html#pytest.hookspec.pytest_configure
-DYNAMODB_ENDPOINT_URL = os.getenv('DYNAMODB_ENDPOINT_URL')
+def pytest_configure():
-PK = os.getenv('DYNAMODB_PARTITION_KEY', 'pk')
+    os.environ['TZ'] = 'America/Sao_Paulo'
-SK = os.getenv('DYNAMODB_SORT_KEY', 'sk')
+    os.environ['DYNAMODB_PARTITION_KEY'] = PK
+    os.environ['DYNAMODB_SORT_KEY'] = SK
+    os.environ['PYTEST_TABLE_NAME'] = PYTEST_TABLE_NAME
 
 
 @pytest.fixture
@@ -38,7 +43,9 @@ def dynamodb_client():
 
 
 @pytest.fixture()
-def dynamodb_persistence_layer(dynamodb_client) -> DynamoDBPersistenceLayer:
+def dynamodb_persistence_layer(dynamodb_client):
+    from layercake.dynamodb import DynamoDBPersistenceLayer
+
     return DynamoDBPersistenceLayer(PYTEST_TABLE_NAME, dynamodb_client)
 
 

layercake/tests/seeds.jsonl

@@ -1,5 +1,7 @@
 {"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": "cpf"}, "sk": {"S": "07879819908"}, "user_id": {"S": "5OxmMjL-ujoR5IMGegQz"}}
+{"id": {"S": "email"}, "sk": {"S": "osergiosiqueira@gmail.com"}, "user_id": {"S": "5OxmMjL-ujoR5IMGegQz"}}
 {"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"}}
 {"id": {"S": "cJtK9SsnJhKPyxESe7g3DG"}, "sk": {"S": "0"}, "name": {"S": "EDUSEG"}, "cnpj": {"S": "15608435000190"}, "email": {"S": "org+15608435000190@users.noreply.betaeducacao.com.br"}}

layercake/tests/test_dynamodb.py

@@ -119,7 +119,7 @@ def test_collection_get_item(
     dynamodb_seeds,
     dynamodb_persistence_layer: DynamoDBPersistenceLayer,
 ):
-    collect = DynamoDBCollection(dynamodb_persistence_layer)
+    collect = dynamodb_persistence_layer.collect
     data_notfound = collect.get_item(
         KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
@@ -157,6 +157,24 @@ def test_collection_get_item(
         )
 
 
+def test_collection_get_item_path_spec(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = dynamodb_persistence_layer.collect
+
+    # This data was added from seeds
+    data = collect.get_item(
+        KeyPair(
+            pk='5OxmMjL-ujoR5IMGegQz',
+            sk=ComposeKey('sergio@somosbeta.com.br', prefix='emails'),
+        ),
+        'mx_record_exists',
+        default={},
+    )
+    assert data
+
+
 def test_collection_put_item(
     dynamodb_persistence_layer: DynamoDBPersistenceLayer,
 ):
@@ -181,7 +199,6 @@ def test_collection_put_item(
     assert data['sk'] == 'orgs#6d1044d5-18c5-437c-9219-fc2ace7e5ebc'
     assert 'name' in data
     assert 'ttl' in data
-    assert 'ttl_date' in data
 
 
 def test_collection_delete_item(
@@ -291,3 +308,66 @@ def test_collection_get_items_unflatten(
         },
         'payment_policy': {'due_days': Decimal('90')},
     }
+
+
+def test_collection_get_items_pair(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+    doc = collect.get_items(
+        KeyPair('5OxmMjL-ujoR5IMGegQz', '0')
+        + KeyPair('cpf', '07879819908')
+        + KeyPair('email', 'osergiosiqueira@gmail.com')
+    )
+
+    assert doc == {
+        'tenant:org_id': [
+            'cJtK9SsnJhKPyxESe7g3DG',
+            'edp8njvgQuzNkLx2ySNfAD',
+            '8TVSi5oACLxTiT8ycKPmaQ',
+        ],
+        'email_verified': True,
+        'last_login': '2024-02-08T20:53:45.818126-03:00',
+        'sk': '0',
+        'cpf': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
+        'name': 'Sérgio Rafael de Siqueira',
+        'id': '5OxmMjL-ujoR5IMGegQz',
+        'create_date': '2019-03-25T00:00:00-03:00',
+        'cognito:sub': '58efed8d-d276-41a8-8502-4ab8b5a6415e',
+        'update_date': '2024-02-08T16:42:33.776409-03:00',
+        'email': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
+    }
+
+
+def test_collection_get_items_pair_unflatten(
+    dynamodb_seeds,
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+    doc = collect.get_items(
+        KeyPair('5OxmMjL-ujoR5IMGegQz', '0')
+        + KeyPair('cpf', '07879819908')
+        + KeyPair('email', 'osergiosiqueira@gmail.com'),
+        flatten_top=False,
+    )
+
+    assert doc == {
+        '5OxmMjL-ujoR5IMGegQz': {
+            'tenant:org_id': [
+                'cJtK9SsnJhKPyxESe7g3DG',
+                'edp8njvgQuzNkLx2ySNfAD',
+                '8TVSi5oACLxTiT8ycKPmaQ',
+            ],
+            'email_verified': True,
+            'last_login': '2024-02-08T20:53:45.818126-03:00',
+            'cpf': '07879819908',
+            'name': 'Sérgio Rafael de Siqueira',
+            'create_date': '2019-03-25T00:00:00-03:00',
+            'cognito:sub': '58efed8d-d276-41a8-8502-4ab8b5a6415e',
+            'update_date': '2024-02-08T16:42:33.776409-03:00',
+            'email': 'sergio@somosbeta.com.br',
+        },
+        'cpf': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
+        'email': {'user_id': '5OxmMjL-ujoR5IMGegQz'},
+    }

layercake/uv.lock

@@ -589,7 +589,7 @@ wheels = [
 
 [[package]]
 name = "layercake"
-version = "0.2.20"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "arnparse" },
@@ -616,7 +616,6 @@ dev = [
     { name = "jsonlines" },
     { name = "pytest" },
     { name = "pytest-cov" },
-    { name = "pytest-env" },
     { name = "ruff" },
 ]
 
@@ -646,7 +645,6 @@ dev = [
     { name = "jsonlines", specifier = ">=4.0.0" },
     { 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" },
 ]
 
@@ -973,18 +971,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/36/3b/48e79f2cd6a61dbbd4807b4ed46cb564b4fd50a76166b1c4ea5c1d9e2371/pytest_cov-6.0.0-py3-none-any.whl", hash = "sha256:eee6f1b9e61008bd34975a4d5bab25801eb31898b032dd55addc93e96fcaaa35", size = 22949, upload-time = "2024-10-29T20:13:33.215Z" },
 ]
 
-[[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, upload-time = "2024-09-17T22:39:18.566Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/de/b8/87cfb16045c9d4092cfcf526135d73b88101aac83bc1adcf82dfb5fd3833/pytest_env-1.1.5-py3-none-any.whl", hash = "sha256:ce90cf8772878515c24b31cd97c7fa1f4481cd68d588419fd45f10ecaee6bc30", size = 6141, upload-time = "2024-09-17T22:39:16.942Z" },
-]
-
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"