From 3fbc560f2c3d0c417edf42ed526ff81f2e15904f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9rgio=20Rafael=20Siqueira?=
 <osergiosiqueira@gmail.com>
Date: Fri, 21 Mar 2025 10:12:39 -0300
Subject: [PATCH] add ttl

---
 layercake/layercake/dynamodb.py  | 46 ++++++++++++++++++++++++++++++++
 layercake/tests/test_dynamodb.py | 41 +++++++++++++++++++++++-----
 2 files changed, 81 insertions(+), 6 deletions(-)

diff --git a/layercake/layercake/dynamodb.py b/layercake/layercake/dynamodb.py
index 185cbc3..34c550b 100644
--- a/layercake/layercake/dynamodb.py
+++ b/layercake/layercake/dynamodb.py
@@ -1,3 +1,4 @@
+import os
 from datetime import datetime
 from ipaddress import IPv4Address
 from typing import Any, Type
@@ -6,6 +7,10 @@ from aws_lambda_powertools import Logger
 from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
 from botocore.exceptions import ClientError
 
+from .dateutils import now, timestamp
+
+TZ = os.getenv('TZ', 'UTC')
+
 logger = Logger(__name__)
 
 
@@ -432,6 +437,20 @@ class DynamoDBPersistenceLayer:
 
 
 class DynamoDBCollection:
+    """
+    Example
+    -------
+        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'),
+            ),
+        )
+    """
+
     class MissingError(ValueError):
         pass
 
@@ -439,6 +458,7 @@ class DynamoDBCollection:
         self,
         persistence_layer: DynamoDBPersistenceLayer,
         exception_cls: Type[ValueError] = MissingError,
+        tz: str = TZ,
     ) -> None:
         if not issubclass(exception_cls, ValueError):
             raise TypeError(
@@ -447,6 +467,7 @@ class DynamoDBCollection:
 
         self.persistence_layer = persistence_layer
         self.exception_cls = exception_cls
+        self.tz = tz
 
     def get_item(
         self,
@@ -468,3 +489,28 @@ class DynamoDBCollection:
             return glom(data, path_spec, default=default)
 
         return data or default
+
+    def put_item(
+        self, key: KeyPair, ttl: int | datetime | None = None, **kwargs: Any
+    ) -> bool:
+        now_ = now(self.tz)
+
+        if isinstance(ttl, int):
+            kwargs.update(
+                {
+                    'ttl': ttl,
+                    'ttl_date': datetime.fromtimestamp(ttl, now_.tzinfo),
+                }
+            )
+
+        if isinstance(ttl, datetime):
+            kwargs.update(
+                {
+                    'ttl': timestamp(ttl),
+                    'ttl_date': ttl,
+                }
+            )
+
+        return self.persistence_layer.put_item(
+            item=key | {'create_date': now_} | kwargs
+        )
diff --git a/layercake/tests/test_dynamodb.py b/layercake/tests/test_dynamodb.py
index e0e2878..5f3431b 100644
--- a/layercake/tests/test_dynamodb.py
+++ b/layercake/tests/test_dynamodb.py
@@ -4,6 +4,7 @@ from ipaddress import IPv4Address
 import pytest
 from botocore.exceptions import ClientError
 
+from layercake.dateutils import ttl
 from layercake.dynamodb import (
     DynamoDBCollection,
     DynamoDBPersistenceLayer,
@@ -58,29 +59,30 @@ def test_transact_write_items(
         dynamodb_persistence_layer.transact_write_items(transact)
 
 
-def test_collection(
+def test_collection_get_item(
     dynamodb_seeds,
     dynamodb_persistence_layer: DynamoDBPersistenceLayer,
 ):
     collect = DynamoDBCollection(dynamodb_persistence_layer)
-    tenant_item = collect.get_item(
+    data_notfound = collect.get_item(
         key=KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
-            sk=Key('tenant'),
+            sk='tenant',
         ),
         raise_if_missing=False,
         default={},
     )
-    assert tenant_item == {}
+    assert data_notfound == {}
 
-    email_item = collect.get_item(
+    # This item was added from seeds
+    data = collect.get_item(
         key=KeyPair(
             pk='5OxmMjL-ujoR5IMGegQz',
             sk=Key('sergio@somosbeta.com.br', prefix='emails'),
         ),
         default={},
     )
-    assert email_item == {
+    assert data == {
         'email_verified': True,
         'mx_record_exists': True,
         'sk': 'emails#sergio@somosbeta.com.br',
@@ -92,3 +94,30 @@ def test_collection(
 
     with pytest.raises(DynamoDBCollection.MissingError):
         collect.get_item(key=KeyPair('5OxmMjL-ujoR5IMGegQz', 'notfound'))
+
+
+def test_collection_put_item(
+    dynamodb_persistence_layer: DynamoDBPersistenceLayer,
+):
+    collect = DynamoDBCollection(dynamodb_persistence_layer)
+
+    assert collect.put_item(
+        key=KeyPair(
+            '5OxmMjL-ujoR5IMGegQz',
+            Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
+        ),
+        name='Beta Educação',
+        ttl=ttl(days=3),
+    )
+
+    data = collect.get_item(
+        key=KeyPair(
+            pk='5OxmMjL-ujoR5IMGegQz',
+            sk=Key('6d1044d5-18c5-437c-9219-fc2ace7e5ebc', prefix='orgs'),
+        ),
+    )
+
+    assert data['sk'] == 'orgs#6d1044d5-18c5-437c-9219-fc2ace7e5ebc'
+    assert 'name' in data
+    assert 'ttl' in data
+    assert 'ttl_date' in data