feat(idempotency): support dataclasses & pydantic models payloads

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -1,7 +1,6 @@
 """
 Persistence layers supporting idempotency
 """
-
 import datetime
 import hashlib
 import json
@@ -105,6 +104,21 @@ def response_json_as_dict(self) -> dict:
         return json.loads(self.response_data)
 
 
+def _prepare_data(data: Any) -> Any:
+    # Convert pydantic as a dict
+    _dict = getattr(data, "dict", None)
+    if callable(_dict):
+        return _dict()
+
+    # Convert dataclasses as a dict
+    if hasattr(data, "__dataclass_fields__"):
+        import dataclasses
+
+        return dataclasses.asdict(data)
+
+    return getattr(data, "raw_event", data)
+
+
 class BasePersistenceLayer(ABC):
     """
     Abstract Base Class for Idempotency persistence layer.
@@ -226,8 +240,7 @@ def _generate_hash(self, data: Any) -> str:
             Hashed representation of the provided data
 
         """
-        data = getattr(data, "raw_event", data)  # could be a data class depending on decorator order
-        hashed_data = self.hash_function(json.dumps(data, cls=Encoder, sort_keys=True).encode())
+        hashed_data = self.hash_function(json.dumps(_prepare_data(data), cls=Encoder, sort_keys=True).encode())
         return hashed_data.hexdigest()
 
     def _validate_payload(self, data: Dict[str, Any], data_record: DataRecord) -> None:

tests/functional/idempotency/test_idempotency.py

@@ -20,7 +20,7 @@
     IdempotencyValidationError,
 )
 from aws_lambda_powertools.utilities.idempotency.idempotency import idempotent, idempotent_function
-from aws_lambda_powertools.utilities.idempotency.persistence.base import BasePersistenceLayer, DataRecord
+from aws_lambda_powertools.utilities.idempotency.persistence.base import BasePersistenceLayer, DataRecord, _prepare_data
 from aws_lambda_powertools.utilities.validation import envelopes, validator
 from tests.functional.idempotency.conftest import serialize
 from tests.functional.utils import load_event
@@ -1057,3 +1057,35 @@ def two(data):
     assert one(data=mock_event) == "one"
     assert two(data=mock_event) == "two"
     assert len(persistence_store.table.method_calls) == 4
+
+
+def test_idempotent_function_dataclasses():
+    try:
+        # Scenario _prepare_data should allow for python dataclasses
+        from dataclasses import asdict, dataclass
+
+        @dataclass
+        class Foo:
+            name: str
+
+        data = Foo(name="Bar")
+        as_dict = _prepare_data(data)
+        assert asdict(data) == as_dict
+        assert as_dict == {"name": "Bar"}
+
+    except ModuleNotFoundError:
+        # Python 3.6
+        pass
+
+
+def test_idempotent_function_pydantic():
+    # Scenario _prepare_data should allow for pydantic
+    from pydantic import BaseModel
+
+    class Foo(BaseModel):
+        name: str
+
+    data = Foo(name="Bar")
+    as_dict = _prepare_data(data)
+    assert as_dict == data.dict()
+    assert as_dict == {"name": "Bar"}