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

aws_lambda_powertools/utilities/idempotency/__init__.py

@@ -11,4 +11,10 @@
 
 from .idempotency import IdempotencyConfig, idempotent, idempotent_function
 
-__all__ = ("DynamoDBPersistenceLayer", "BasePersistenceLayer", "idempotent", "idempotent_function", "IdempotencyConfig")
+__all__ = (
+    "DynamoDBPersistenceLayer",
+    "BasePersistenceLayer",
+    "idempotent",
+    "idempotent_function",
+    "IdempotencyConfig",
+)

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,17 @@ 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 +182,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 +192,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 +219,8 @@ 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()
+        return self.output_serializer.from_dict(response_dict) if response_dict else None
 
     def _get_function_response(self):
         try:
@@ -226,7 +239,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/idempotency.py

@@ -14,6 +14,7 @@
 from aws_lambda_powertools.utilities.idempotency.persistence.base import (
     BasePersistenceLayer,
 )
+from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 logger = logging.getLogger(__name__)
@@ -85,6 +86,7 @@ def idempotent_function(
     data_keyword_argument: str,
     persistence_store: BasePersistenceLayer,
     config: Optional[IdempotencyConfig] = None,
+    output_serializer: Optional[BaseIdempotencySerializer] = None,
 ) -> Any:
     """
     Decorator to handle idempotency of any function
@@ -99,7 +101,9 @@ def idempotent_function(
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
         Configuration
-
+    output_serializer: Optional[BaseIdempotencySerializer]
+            Serializer to transform the data to and from a dictionary.
+            If not supplied, no serialization is done via the NoOpSerializer
     Examples
     --------
     **Processes an order in an idempotent manner**
@@ -124,6 +128,7 @@ 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,
             ),
         )
 
@@ -147,6 +152,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/base.py

@@ -0,0 +1,19 @@
+"""
+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:
+        pass
+
+    @abstractmethod
+    def from_dict(self, data: Dict) -> Any:
+        pass

aws_lambda_powertools/utilities/idempotency/serialization/custom_dict.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/no_op.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/pydantic.py

@@ -0,0 +1,30 @@
+from typing import Dict, Type
+
+from pydantic import BaseModel
+
+from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+
+Model = BaseModel
+
+
+class PydanticSerializer(BaseIdempotencySerializer):
+    def __init__(self, model: Type[Model]):
+        """
+        Parameters
+        ----------
+        model: Model
+            A Pydantic model of the type to transform
+        """
+        self.__model: Type[Model] = model
+
+    def to_dict(self, data: Model) -> 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) -> Model:
+        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)

docs/utilities/idempotency.md

@@ -152,6 +152,22 @@ 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
+
+By supplying an output serializer, you can control the return type of the function, allowing cleaner integration with the rest of your code base.
+
+=== "Using A Custom Type (Dataclasses)"
+
+    ```python hl_lines="3-8 26-28 32-35 42 48"
+    --8<-- "examples/idempotency/src/working_with_idempotent_function_dataclass_output_serializer.py"
+    ```
+
+=== "Using Pydantic"
+
+    ```python hl_lines="3-8 23-24 32 35"
+    --8<-- "examples/idempotency/src/working_with_idempotent_function_pydantic_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.

examples/idempotency/src/working_with_idempotent_function_dataclass_output_serializer.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    IdempotencyConfig,
+    idempotent_function,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.custom_dict import CustomDictSerializer
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+dynamodb = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
+config = IdempotencyConfig(event_key_jmespath="order_id")  # see Choosing a payload subset section
+
+
+@dataclass
+class OrderItem:
+    sku: str
+    description: str
+
+
+@dataclass
+class Order:
+    item: OrderItem
+    order_id: int
+
+
+@dataclass
+class OrderOutput:
+    order_id: int
+
+
+order_output_serializer: CustomDictSerializer = CustomDictSerializer(
+    to_dict=lambda x: x.asdict(),
+    from_dict=lambda x: OrderOutput(**x),
+)
+
+
+@idempotent_function(
+    data_keyword_argument="order",
+    config=config,
+    persistence_store=dynamodb,
+    output_serializer=order_output_serializer,
+)
+def process_order(order: Order) -> OrderOutput:
+    return OrderOutput(order_id=order.order_id)
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    config.register_lambda_context(context)  # see Lambda timeouts section
+    order_item = OrderItem(sku="fake", description="sample")
+    order = Order(item=order_item, order_id=1)
+
+    # `order` parameter must be called as a keyword argument to work
+    process_order(order=order)

examples/idempotency/src/working_with_idempotent_function_pydantic_output_serializer.py

@@ -0,0 +1,44 @@
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    IdempotencyConfig,
+    idempotent_function,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.pydantic import PydanticSerializer
+from aws_lambda_powertools.utilities.parser import BaseModel
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+dynamodb = DynamoDBPersistenceLayer(table_name="IdempotencyTable")
+config = IdempotencyConfig(event_key_jmespath="order_id")  # see Choosing a payload subset section
+
+
+class OrderItem(BaseModel):
+    sku: str
+    description: str
+
+
+class Order(BaseModel):
+    item: OrderItem
+    order_id: int
+
+
+class OrderOutput(BaseModel):
+    order_id: int
+
+
+@idempotent_function(
+    data_keyword_argument="order",
+    config=config,
+    persistence_store=dynamodb,
+    output_serializer=PydanticSerializer(model=OrderOutput),
+)
+def process_order(order: Order) -> OrderOutput:
+    return OrderOutput(order_id=order.order_id)
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    config.register_lambda_context(context)  # see Lambda timeouts section
+    order_item = OrderItem(sku="fake", description="sample")
+    order = Order(item=order_item, order_id=1)
+
+    # `order` parameter must be called as a keyword argument to work
+    process_order(order=order)