Merge pull request #2 from aradyaron/feat/idempotency-output-serializer-with-model

Feat/idempotency output serializer with model

Author: aradyaron

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 type
+    """
+
+
+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,7 +15,10 @@
 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.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
 logger = logging.getLogger(__name__)
@@ -86,7 +90,7 @@ def idempotent_function(
     data_keyword_argument: str,
     persistence_store: BasePersistenceLayer,
     config: Optional[IdempotencyConfig] = None,
-    output_serializer: Optional[BaseIdempotencySerializer] = None,
+    output_serializer: Optional[Union[BaseIdempotencySerializer, Type[BaseIdempotencyModelSerializer]]] = None,
 ) -> Any:
     """
     Decorator to handle idempotency of any function
@@ -101,9 +105,11 @@ def idempotent_function(
         Instance of BasePersistenceLayer to store data
     config: IdempotencyConfig
         Configuration
-    output_serializer: Optional[BaseIdempotencySerializer]
+    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
+            If not supplied, no serialization is done via the NoOpSerializer.
+            In case a serializer of type inheriting BaseIdempotencyModelSerializer] is given,
+            the serializer is deduced from the function return type.
     Examples
     --------
     **Processes an order in an idempotent manner**
@@ -131,6 +137,9 @@ def process_order(customer_id: str, order: dict, **kwargs):
                 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()
 

aws_lambda_powertools/utilities/idempotency/serialization/base.py

@@ -17,3 +17,31 @@ def to_dict(self, data: Any) -> Dict:
     @abstractmethod
     def from_dict(self, data: Dict) -> Any:
         pass
+
+
+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:
+        """
+        Instantiate the serializer from the given model type.
+        In case there is the model_type is unknown, None will be sent to the method.
+        It's on the implementer to verify that:
+        - None is handled
+        - 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/dataclass.py

@@ -0,0 +1,39 @@
+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):
+    def __init__(self, model: Type[DataClass]):
+        """
+        Parameters
+        ----------
+        model: Model
+            A Pydantic model of the type to transform
+        """
+        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/pydantic.py

@@ -1,13 +1,20 @@
-from typing import Dict, Type
+from typing import Any, Dict, Type
 
 from pydantic import BaseModel
 
-from aws_lambda_powertools.utilities.idempotency.serialization.base import BaseIdempotencySerializer
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyModelTypeError,
+    IdempotencyNoSerializationModelError,
+)
+from aws_lambda_powertools.utilities.idempotency.serialization.base import (
+    BaseIdempotencyModelSerializer,
+    BaseIdempotencySerializer,
+)
 
 Model = BaseModel
 
 
-class PydanticSerializer(BaseIdempotencySerializer):
+class PydanticSerializer(BaseIdempotencyModelSerializer):
     def __init__(self, model: Type[Model]):
         """
         Parameters
@@ -28,3 +35,13 @@ def from_dict(self, data: Dict) -> Model:
             # 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)