add ttl

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
+        )