feat(idempotency): support dataclasses & pydantic models payloads

aws_lambda_powertools/utilities/idempotency/base.py

@@ -21,6 +21,20 @@
 logger = logging.getLogger(__name__)
 
 
+def _prepare_data(data: Any) -> Any:
+    """Prepare data for json serialization.
+    This will convert dataclasses, pydantic models or event source data classes to a dict."""
+    if hasattr(data, "__dataclass_fields__"):
+        import dataclasses
+
+        return dataclasses.asdict(data)
+
+    if callable(getattr(data, "dict", None)):
+        return data.dict()
+
+    return getattr(data, "raw_event", data)
+
+
 class IdempotencyHandler:
     """
     Base class to orchestrate calls to persistence layer.
@@ -52,7 +66,7 @@ def __init__(
             Function keyword arguments
         """
         self.function = function
-        self.data = function_payload
+        self.data = _prepare_data(function_payload)
         self.fn_args = function_args
         self.fn_kwargs = function_kwargs
 

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -1,7 +1,6 @@
 """
 Persistence layers supporting idempotency
 """
-
 import datetime
 import hashlib
 import json
@@ -226,7 +225,6 @@ 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())
         return hashed_data.hexdigest()
 

tests/functional/idempotency/test_idempotency.py

@@ -8,9 +8,11 @@
 import jmespath
 import pytest
 from botocore import stub
+from pydantic import BaseModel
 
 from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEventV2, event_source
 from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer, IdempotencyConfig
+from aws_lambda_powertools.utilities.idempotency.base import _prepare_data
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
@@ -1069,3 +1071,40 @@ def test_invalid_dynamodb_persistence_layer():
         )
     # and raise a ValueError
     assert str(ve.value) == "key_attr [id] and sort_key_attr [id] cannot be the same!"
+
+
+def test_idempotent_function_dataclasses():
+    try:
+        # Scenario _prepare_data should convert a python dataclasses to a dict
+        from dataclasses import asdict, dataclass
+
+        @dataclass
+        class Foo:
+            name: str
+
+        expected_result = {"name": "Bar"}
+        data = Foo(name="Bar")
+        as_dict = _prepare_data(data)
+        assert as_dict == asdict(data)
+        assert as_dict == expected_result
+
+    except ModuleNotFoundError:
+        pass  # Python 3.6
+
+
+def test_idempotent_function_pydantic():
+    # Scenario _prepare_data should convert a pydantic to a dict
+    class Foo(BaseModel):
+        name: str
+
+    expected_result = {"name": "Bar"}
+    data = Foo(name="Bar")
+    as_dict = _prepare_data(data)
+    assert as_dict == data.dict()
+    assert as_dict == expected_result
+
+
+@pytest.mark.parametrize("data", [None, "foo", ["foo"], 1, True, {}])
+def test_idempotent_function_other(data):
+    # All other data types should be left as is
+    assert _prepare_data(data) == data