Merge pull request #18 from RSS-Engineering/support_ttl

Support ttl field proper datetime format

Author: explorigin

README.md

@@ -90,6 +90,8 @@ NOTE: Rule complexity is limited by the querying capabilities of the backend.
 
 `global_indexes` - (optional) specify a mapping of index-name to tuple(partition_key).
 
+`ttl` - (optional) the name of the datetime-typed field that dynamo should consider to be the TTL field. PydantiCRUD will save this field as a float type instead of an ISO datetime string. This field only works properly with UTC-zoned datetime instances.
+
 ### SQLite (Python 3.7+)
 
 `database` - the filename of the database file for SQLite to use

pydanticrud/backends/dynamodb.py

@@ -1,7 +1,8 @@
 from typing import Optional, Set, Union, Dict, Any
 import logging
 import json
-from datetime import datetime
+from datetime import datetime, timezone
+from decimal import Decimal
 
 import boto3
 from boto3.dynamodb.conditions import Key, Attr
@@ -45,7 +46,7 @@ def expression_to_condition(expr, keys: set):
     if isinstance(expr, ast.NullExpression):
         return None, set()
     if isinstance(expr, ast.DatetimeExpression):
-        return _to_epoch_float(expr.value), set()
+        return _to_epoch_decimal(expr.value), set()
     if isinstance(expr, ast.StringExpression):
         return expr.value, set()
     if isinstance(expr, ast.FloatExpression):
@@ -70,17 +71,22 @@ def rule_to_boto_expression(rule: Rule, keys: Optional[Set[str]] = None):
     "bool": "BOOL",
 }
 
-EPOCH = datetime.utcfromtimestamp(0)
+EPOCH = datetime(1970, 1, 1, 0, 0)
 
 
-def _to_epoch_float(dt):
-    return (dt - EPOCH).total_seconds * 1000
+def _to_epoch_decimal(dt: datetime) -> Decimal:
+    """TTL fields must be stored as a float but boto only supports decimals."""
+    epock = EPOCH
+    if dt.tzinfo:
+        epock = epock.replace(tzinfo=timezone.utc)
+    return Decimal((dt - epock).total_seconds())
 
 
 SERIALIZE_MAP = {
     "number": str,  # float or decimal
     "string": str,
     "string:date-time": lambda d: d.isoformat(),
+    "string:ttl": lambda d: _to_epoch_decimal(d),
     "boolean": lambda d: 1 if d else 0,
     "object": json.dumps,
     "array": json.dumps,
@@ -96,7 +102,7 @@ def _to_epoch_float(dt):
 
 def chunk_list(lst, size):
     for i in range(0, len(lst), size):
-        yield lst[i:i + size]
+        yield lst[i : i + size]
 
 
 def index_definition(index_name, keys, gsi=False):
@@ -115,9 +121,10 @@ def index_definition(index_name, keys, gsi=False):
 
 
 class DynamoSerializer:
-    def __init__(self, schema):
+    def __init__(self, schema, ttl_field=None):
         self.properties = schema.get("properties")
         self.definitions = schema.get("definitions")
+        self.ttl_field = ttl_field
 
     def _get_type_possibilities(self, field_name) -> Set[tuple]:
         field_properties = self.properties.get(field_name)
@@ -142,7 +149,11 @@ def type_from_definition(definition_signature: Union[str, dict]) -> dict:
         return set([(t["type"], t.get("format", "")) for t in type_dicts])
 
     def _serialize_field(self, field_name, value):
-        field_types = self._get_type_possibilities(field_name)
+        if field_name == self.ttl_field:
+            field_types = {("string", "ttl")}
+        else:
+            field_types = self._get_type_possibilities(field_name)
+
         if value is not None:
             for t in field_types:
                 try:
@@ -205,9 +216,9 @@ def __init__(self, cls):
         cfg = cls.Config
         self.cls = cls
         self.schema = cls.schema()
-        self.serializer = DynamoSerializer(self.schema)
         self.hash_key = cfg.hash_key
         self.range_key = getattr(cfg, "range_key", None)
+        self.serializer = DynamoSerializer(self.schema, ttl_field=getattr(cfg, "ttl", None))
         self.table_name = cls.get_table_name()
 
         self.local_indexes = getattr(cfg, "local_indexes", {})
@@ -427,7 +438,9 @@ def batch_save(self, items: list) -> dict:
         # chunk list for size limit of 25 items to write using this batch_write operation refer below.
         # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/client/batch_write_item.html#:~:text=The%20BatchWriteItem%20operation,Data%20Types.
         for chunk in chunk_list(items, 25):
-            serialized_items = [self.serializer.serialize_record(item.dict(by_alias=True)) for item in chunk]
+            serialized_items = [
+                self.serializer.serialize_record(item.dict(by_alias=True)) for item in chunk
+            ]
             for serialized_item in serialized_items:
                 request_items[self.table_name].append({"PutRequest": {"Item": serialized_item}})
         try:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pydanticrud"
-version = "0.3.0"
+version = "0.4.0"
 description = "Supercharge your Pydantic models with CRUD methods and a pluggable backend"
 authors = ["Timothy Farrell <[email protected]>"]
 license = "MIT"

tests/test_dynamodb.py

@@ -1,6 +1,6 @@
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Any
 from decimal import Decimal
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from uuid import uuid4, UUID
 import random
 
@@ -27,6 +27,7 @@ class SimpleKeyModel(BaseModel):
     name: str
     total: float
     timestamp: datetime
+    expires: datetime
     sigfig: Decimal
     enabled: bool
     data: Dict[int, int] = None
@@ -36,6 +37,7 @@ class SimpleKeyModel(BaseModel):
     class Config:
         title = "ModelTitle123"
         hash_key = "name"
+        ttl = "expires"
         backend = DynamoDbBackend
         endpoint = "http://localhost:18002"
         global_indexes = {"by-id": ("id",)}
@@ -125,6 +127,7 @@ def simple_model_data_generator(**kwargs):
         name=random_unique_name(),
         total=round(random.random(), 9),
         timestamp=random_datetime(),
+        expires=(datetime.utcnow() + timedelta(seconds=random.randint(0, 10))).replace(tzinfo=timezone.utc),
         sigfig=Decimal(str(random.random())[:8]),
         enabled=random.choice((True, False)),
         data={random.randint(0, 1000): random.randint(0, 1000)},
@@ -309,6 +312,21 @@ def test_save_get_delete_simple(dynamo, simple_table):
         SimpleKeyModel.get(data["name"])
 
 
+def test_save_ttl_field_is_float(dynamo, simple_query_data):
+    """DynamoDB requires ttl fields to be a float in order to be successfully processed. Boto provides the ability to
+    set a float via a decimal (but not a float strangely)."""
+
+    key = simple_query_data[0]["name"]
+    table = SimpleKeyModel.__backend__.get_table()
+    resp = table.get_item(Key=SimpleKeyModel.__backend__._key_param_to_dict(key))
+    expires_value = resp["Item"]["expires"]
+    assert isinstance(expires_value, Decimal)
+    assert datetime.utcfromtimestamp(float(expires_value)).replace(tzinfo=timezone.utc) == simple_query_data[0]["expires"]
+
+    instance = SimpleKeyModel.get(key)
+    assert instance.expires == simple_query_data[0]["expires"]
+
+
 def test_query_with_hash_key_simple(dynamo, simple_query_data):
     res = SimpleKeyModel.query(Rule(f"name == '{simple_query_data[0]['name']}'"))
     res_data = {m.name: m.dict() for m in res}