feat(idempotency): add support to custom serialization/deserialization on idempotency decorator (#2951)

Co-authored-by: arad.yaron <arad.yaron@cyberark.com>
Co-authored-by: Leandro Damascena <lcdama@amazon.pt>
Co-authored-by: Ruben Fonseca <rubefons@amazon.com>

Author: 3 people

aws_lambda_powertools/utilities/idempotency/base.py

@@ -18,6 +18,12 @@
     BasePersistenceLayer,
     DataRecord,
 )
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencySerializer,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.no_op import (
+    NoOpSerializer,
+)
 
 MAX_RETRIES = 2
 logger = logging.getLogger(__name__)
@@ -51,6 +57,7 @@ def __init__(
         function_payload: Any,
         config: IdempotencyConfig,
         persistence_store: BasePersistenceLayer,
+        output_serializer: Optional[BaseIdempotencySerializer] = None,
         function_args: Optional[Tuple] = None,
         function_kwargs: Optional[Dict] = None,
     ):
@@ -65,12 +72,16 @@ def __init__(
             Idempotency Configuration
         persistence_store : BasePersistenceLayer
             Instance of persistence layer to store idempotency records
+        output_serializer: Optional[BaseIdempotencySerializer]
+            Serializer to transform the data to and from a dictionary.
+            If not supplied, no serialization is done via the NoOpSerializer
         function_args: Optional[Tuple]
             Function arguments
         function_kwargs: Optional[Dict]
             Function keyword arguments
         """
         self.function = function
+        self.output_serializer = output_serializer or NoOpSerializer()
         self.data = deepcopy(_prepare_data(function_payload))
         self.fn_args = function_args
         self.fn_kwargs = function_kwargs
@@ -170,7 +181,7 @@ def _get_idempotency_record(self) -> Optional[DataRecord]:
 
         return data_record
 
-    def _handle_for_status(self, data_record: DataRecord) -> Optional[Dict[Any, Any]]:
+    def _handle_for_status(self, data_record: DataRecord) -> Optional[Any]:
         """
         Take appropriate action based on data_record's status
 
@@ -180,8 +191,9 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Dict[Any, Any]
 
         Returns
         -------
-        Optional[Dict[Any, Any]
+        Optional[Any]
             Function's response previously used for this idempotency key, if it has successfully executed already.
+            In case an output serializer is configured, the response is deserialized.
 
         Raises
         ------
@@ -206,8 +218,10 @@ def _handle_for_status(self, data_record: DataRecord) -> Optional[Dict[Any, Any]
                 f"Execution already in progress with idempotency key: "
                 f"{self.persistence_store.event_key_jmespath}={data_record.idempotency_key}",
             )
-
-        return data_record.response_json_as_dict()
+        response_dict: Optional[dict] = data_record.response_json_as_dict()
+        if response_dict is not None:
+            return self.output_serializer.from_dict(response_dict)
+        return None
 
     def _get_function_response(self):
         try:
@@ -226,7 +240,8 @@ def _get_function_response(self):
 
         else:
             try:
-                self.persistence_store.save_success(data=self.data, result=response)
+                serialized_response: dict = self.output_serializer.to_dict(response) if response else None
+                self.persistence_store.save_success(data=self.data, result=serialized_response)
             except Exception as save_exception:
                 raise IdempotencyPersistenceLayerError(
                     "Failed to update record state to success in idempotency store",

aws_lambda_powertools/utilities/idempotency/exceptions.py

@@ -71,3 +71,15 @@ class IdempotencyKeyError(BaseError):
     """
     Payload does not contain an idempotent key
     """
+
+
+class IdempotencyModelTypeError(BaseError):
+    """
+    Model type does not match expected payload output
+    """
+
+
+class IdempotencyNoSerializationModelError(BaseError):
+    """
+    No model was supplied to the serializer
+    """

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -4,7 +4,8 @@
 import functools
 import logging
 import os
-from typing import Any, Callable, Dict, Optional, cast
+from inspect import isclass
+from typing import Any, Callable, Dict, Optional, Type, Union, cast
 
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
 from aws_lambda_powertools.shared import constants
@@ -14,6 +15,10 @@
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 logger = logging.getLogger(__name__)
@@ -85,6 +90,7 @@ def idempotent_function(
     data_keyword_argument: str,
     persistence_store: BasePersistenceLayer,
     config: Optional[IdempotencyConfig] = None,
+    output_serializer: Optional[Union[BaseIdempotencySerializer, Type[BaseIdempotencyModelSerializer]]] = None,
 ) -> Any:
     """
     Decorator to handle idempotency of any function
@@ -99,6 +105,11 @@ def idempotent_function(
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
         Configuration
+    output_serializer: Optional[Union[BaseIdempotencySerializer, Type[BaseIdempotencyModelSerializer]]]
+            Serializer to transform the data to and from a dictionary.
+            If not supplied, no serialization is done via the NoOpSerializer.
+            In case a serializer of type inheriting BaseIdempotencyModelSerializer is given,
+            the serializer is derived from the function return type.
 
     Examples
     --------
@@ -124,9 +135,14 @@ def process_order(customer_id: str, order: dict, **kwargs):
                 data_keyword_argument=data_keyword_argument,
                 persistence_store=persistence_store,
                 config=config,
+                output_serializer=output_serializer,
             ),
         )
 
+    if isclass(output_serializer) and issubclass(output_serializer, BaseIdempotencyModelSerializer):
+        # instantiate an instance of the serializer class
+        output_serializer = output_serializer.instantiate(function.__annotations__.get("return", None))
+
     config = config or IdempotencyConfig()
 
     @functools.wraps(function)
@@ -147,6 +163,7 @@ def decorate(*args, **kwargs):
             function_payload=payload,
             config=config,
             persistence_store=persistence_store,
+            output_serializer=output_serializer,
             function_args=args,
             function_kwargs=kwargs,
         )

aws_lambda_powertools/utilities/idempotency/serialization/__init__.py

@@ -0,0 +1,47 @@
+"""
+Serialization for supporting idempotency
+"""
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class BaseIdempotencySerializer(ABC):
+    """
+    Abstract Base Class for Idempotency serialization layer, supporting dict operations.
+    """
+
+    @abstractmethod
+    def to_dict(self, data: Any) -> Dict:
+        raise NotImplementedError("Implementation of to_dict is required")
+
+    @abstractmethod
+    def from_dict(self, data: Dict) -> Any:
+        raise NotImplementedError("Implementation of from_dict is required")
+
+
+class BaseIdempotencyModelSerializer(BaseIdempotencySerializer):
+    """
+    Abstract Base Class for Idempotency serialization layer, for using a model as data object representation.
+    """
+
+    @classmethod
+    @abstractmethod
+    def instantiate(cls, model_type: Any) -> BaseIdempotencySerializer:
+        """
+        Creates an instance of a serializer based on a provided model type.
+        In case the model_type is unknown, None will be sent as `model_type`.
+        It's on the implementer to verify that:
+        - None is handled correctly
+        - A model type not matching the expected types is handled
+
+        Parameters
+        ----------
+        model_type: Any
+            The model type to instantiate the class for
+
+        Returns
+        -------
+        BaseIdempotencySerializer
+            Instance of the serializer class
+        """
+        pass

aws_lambda_powertools/utilities/idempotency/serialization/base.py

@@ -0,0 +1,23 @@
+from typing import Any, Callable, Dict
+
+from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+
+
+class CustomDictSerializer(BaseIdempotencySerializer):
+    def __init__(self, to_dict: Callable[[Any], Dict], from_dict: Callable[[Dict], Any]):
+        """
+        Parameters
+        ----------
+        to_dict: Callable[[Any], Dict]
+            A function capable of transforming the saved data object representation into a dictionary
+        from_dict: Callable[[Dict], Any]
+            A function capable of transforming the saved dictionary into the original data object representation
+        """
+        self.__to_dict: Callable[[Any], Dict] = to_dict
+        self.__from_dict: Callable[[Dict], Any] = from_dict
+
+    def to_dict(self, data: Any) -> Dict:
+        return self.__to_dict(data)
+
+    def from_dict(self, data: Dict) -> Any:
+        return self.__from_dict(data)

aws_lambda_powertools/utilities/idempotency/serialization/custom_dict.py

@@ -0,0 +1,43 @@
+from dataclasses import asdict, is_dataclass
+from typing import Any, Dict, Type
+
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyModelTypeError,
+    IdempotencyNoSerializationModelError,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
+
+DataClass = Any
+
+
+class DataclassSerializer(BaseIdempotencyModelSerializer):
+    """
+    A serializer class for transforming data between dataclass objects and dictionaries.
+    """
+
+    def __init__(self, model: Type[DataClass]):
+        """
+        Parameters
+        ----------
+        model: Type[DataClass]
+            A dataclass type to be used for serialization and deserialization
+        """
+        self.__model: Type[DataClass] = model
+
+    def to_dict(self, data: DataClass) -> Dict:
+        return asdict(data)
+
+    def from_dict(self, data: Dict) -> DataClass:
+        return self.__model(**data)
+
+    @classmethod
+    def instantiate(cls, model_type: Any) -> BaseIdempotencySerializer:
+        if model_type is None:
+            raise IdempotencyNoSerializationModelError("No serialization model was supplied")
+
+        if not is_dataclass(model_type):
+            raise IdempotencyModelTypeError("Model type is not inherited of dataclass type")
+        return cls(model=model_type)

aws_lambda_powertools/utilities/idempotency/serialization/dataclass.py

@@ -0,0 +1,18 @@
+from typing import Dict
+
+from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+
+
+class NoOpSerializer(BaseIdempotencySerializer):
+    def __init__(self):
+        """
+        Parameters
+        ----------
+        Default serializer, does not transform data
+        """
+
+    def to_dict(self, data: Dict) -> Dict:
+        return data
+
+    def from_dict(self, data: Dict) -> Dict:
+        return data

aws_lambda_powertools/utilities/idempotency/serialization/no_op.py

@@ -0,0 +1,47 @@
+from typing import Any, Dict, Type
+
+from pydantic import BaseModel
+
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyModelTypeError,
+    IdempotencyNoSerializationModelError,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
+
+
+class PydanticSerializer(BaseIdempotencyModelSerializer):
+    """Pydantic serializer for idempotency models"""
+
+    def __init__(self, model: Type[BaseModel]):
+        """
+        Parameters
+        ----------
+        model: Model
+            Pydantic model to be used for serialization
+        """
+        self.__model: Type[BaseModel] = model
+
+    def to_dict(self, data: BaseModel) -> Dict:
+        if callable(getattr(data, "model_dump", None)):
+            # Support for pydantic V2
+            return data.model_dump()  # type: ignore[unused-ignore,attr-defined]
+        return data.dict()
+
+    def from_dict(self, data: Dict) -> BaseModel:
+        if callable(getattr(self.__model, "model_validate", None)):
+            # Support for pydantic V2
+            return self.__model.model_validate(data)  # type: ignore[unused-ignore,attr-defined]
+        return self.__model.parse_obj(data)
+
+    @classmethod
+    def instantiate(cls, model_type: Any) -> BaseIdempotencySerializer:
+        if model_type is None:
+            raise IdempotencyNoSerializationModelError("No serialization model was supplied")
+
+        if not issubclass(model_type, BaseModel):
+            raise IdempotencyModelTypeError("Model type is not inherited from pydantic BaseModel")
+
+        return cls(model=model_type)

aws_lambda_powertools/utilities/idempotency/serialization/pydantic.py

@@ -152,6 +152,45 @@ When using `idempotent_function`, you must tell us which keyword parameter in yo
     --8<-- "examples/idempotency/src/working_with_idempotent_function_pydantic.py"
     ```
 
+#### Output serialization
+
+The default return of the `idempotent_function` decorator is a JSON object, but you can customize the function's return type by utilizing the `output_serializer` parameter. The output serializer supports any JSON serializable data, **Python Dataclasses** and **Pydantic Models**.
+
+!!! info "When using the `output_serializer` parameter, the data will continue to be stored in DynamoDB as a JSON object."
+
+Working with Pydantic Models:
+
+=== "Explicitly passing the Pydantic model type"
+
+    ```python hl_lines="6 24 25 32 35 44"
+    --8<-- "examples/idempotency/src/working_with_pydantic_explicitly_output_serializer.py"
+    ```
+=== "Deducing the Pydantic model type from the return type annotation"
+
+    ```python hl_lines="6 24 25 32 36 45"
+    --8<-- "examples/idempotency/src/working_with_pydantic_deduced_output_serializer.py"
+    ```
+
+Working with Python Dataclasses:
+
+=== "Explicitly passing the model type"
+
+    ```python hl_lines="8 27-29 36 39 48"
+    --8<-- "examples/idempotency/src/working_with_dataclass_explicitly_output_serializer.py"
+    ```
+
+=== "Deducing the model type from the return type annotation"
+
+    ```python hl_lines="8 27-29 36 40 49"
+    --8<-- "examples/idempotency/src/working_with_dataclass_deduced_output_serializer.py"
+    ```
+
+=== "Using A Custom Type (Dataclasses)"
+
+    ```python hl_lines="9 33 37 41-44 51 54"
+    --8<-- "examples/idempotency/src/working_with_idempotent_function_custom_output_serializer.py"
+    ```
+
 #### Batch integration
 
 You can can easily integrate with [Batch utility](batch.md){target="_blank"} via context manager. This ensures that you process each record in an idempotent manner, and guard against a [Lambda timeout](#lambda-timeouts) idempotent situation.