Merge pull request #118 from risenberg-cyberark/pydantic

feat: Advanced parser utility (pydantic)

Author: heitorlessa

Makefile

@@ -4,7 +4,7 @@ target:
 
 dev:
 	pip install --upgrade pip poetry pre-commit
-	poetry install
+	poetry install --extras "pydantic"
 	pre-commit install
 
 dev-docs:

aws_lambda_powertools/utilities/advanced_parser/__init__.py

@@ -0,0 +1,6 @@
+"""Advanced parser utility
+"""
+from .envelopes import Envelope, InvalidEnvelopeError, parse_envelope
+from .parser import parser
+
+__all__ = ["InvalidEnvelopeError", "Envelope", "parse_envelope", "parser"]

aws_lambda_powertools/utilities/advanced_parser/envelopes/__init__.py

@@ -0,0 +1,3 @@
+from .envelopes import Envelope, InvalidEnvelopeError, parse_envelope
+
+__all__ = ["InvalidEnvelopeError", "Envelope", "parse_envelope"]

aws_lambda_powertools/utilities/advanced_parser/envelopes/base.py

@@ -0,0 +1,38 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from pydantic import BaseModel, ValidationError
+
+logger = logging.getLogger(__name__)
+
+
+class BaseEnvelope(ABC):
+    def _parse_user_dict_schema(self, user_event: Dict[str, Any], schema: BaseModel) -> Any:
+        if user_event is None:
+            return None
+        logger.debug("parsing user dictionary schema")
+        try:
+            return schema(**user_event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception while extracting user custom schema")
+            raise
+
+    def _parse_user_json_string_schema(self, user_event: str, schema: BaseModel) -> Any:
+        if user_event is None:
+            return None
+        # this is used in cases where the underlying schema is not a Dict that can be parsed as baseModel
+        # but a plain string i.e SQS has plain string payload
+        if schema == str:
+            logger.debug("input is string, returning")
+            return user_event
+        logger.debug("trying to parse as json encoded string")
+        try:
+            return schema.parse_raw(user_event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception while extracting user custom schema")
+            raise
+
+    @abstractmethod
+    def parse(self, event: Dict[str, Any], schema: BaseModel):
+        return NotImplemented

aws_lambda_powertools/utilities/advanced_parser/envelopes/dynamodb.py

@@ -0,0 +1,31 @@
+import logging
+from typing import Any, Dict, List
+from typing_extensions import Literal
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.schemas import DynamoDBSchema
+
+logger = logging.getLogger(__name__)
+
+
+# returns a List of dictionaries which each contains two keys, "NewImage" and "OldImage".
+# The values are the parsed schema models. The images' values can also be None.
+# Length of the list is the record's amount in the original event.
+class DynamoDBEnvelope(BaseEnvelope):
+    def parse(self, event: Dict[str, Any], schema: BaseModel) -> List[Dict[Literal["NewImage", "OldImage"], BaseModel]]:
+        try:
+            parsed_envelope = DynamoDBSchema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input dynamodb stream event")
+            raise
+        output = []
+        for record in parsed_envelope.Records:
+            output.append(
+                {
+                    "NewImage": self._parse_user_dict_schema(record.dynamodb.NewImage, schema),
+                    "OldImage": self._parse_user_dict_schema(record.dynamodb.OldImage, schema),
+                }
+            )
+        return output

aws_lambda_powertools/utilities/advanced_parser/envelopes/envelopes.py

@@ -0,0 +1,42 @@
+import logging
+from enum import Enum
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.dynamodb import DynamoDBEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.event_bridge import EventBridgeEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.sqs import SqsEnvelope
+
+logger = logging.getLogger(__name__)
+
+
+"""Built-in envelopes"""
+
+
+class Envelope(str, Enum):
+    SQS = "sqs"
+    EVENTBRIDGE = "eventbridge"
+    DYNAMODB_STREAM = "dynamodb_stream"
+
+
+class InvalidEnvelopeError(Exception):
+    """Input envelope is not one of the Envelope enum values"""
+
+
+# enum to BaseEnvelope handler class
+__ENVELOPE_MAPPING = {
+    Envelope.SQS: SqsEnvelope,
+    Envelope.DYNAMODB_STREAM: DynamoDBEnvelope,
+    Envelope.EVENTBRIDGE: EventBridgeEnvelope,
+}
+
+
+def parse_envelope(event: Dict[str, Any], envelope: Envelope, schema: BaseModel):
+    envelope_handler: BaseEnvelope = __ENVELOPE_MAPPING.get(envelope)
+    if envelope_handler is None:
+        logger.exception("envelope must be an instance of Envelope enum")
+        raise InvalidEnvelopeError("envelope must be an instance of Envelope enum")
+    logger.debug(f"Parsing and validating event schema, envelope={str(envelope.value)}")
+    return envelope_handler().parse(event=event, schema=schema)

aws_lambda_powertools/utilities/advanced_parser/envelopes/event_bridge.py

@@ -0,0 +1,20 @@
+import logging
+from typing import Any, Dict
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.schemas import EventBridgeSchema
+
+logger = logging.getLogger(__name__)
+
+
+# returns a parsed BaseModel object according to schema type
+class EventBridgeEnvelope(BaseEnvelope):
+    def parse(self, event: Dict[str, Any], schema: BaseModel) -> BaseModel:
+        try:
+            parsed_envelope = EventBridgeSchema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input eventbridge event")
+            raise
+        return self._parse_user_dict_schema(parsed_envelope.detail, schema)

aws_lambda_powertools/utilities/advanced_parser/envelopes/sqs.py

@@ -0,0 +1,27 @@
+import logging
+from typing import Any, Dict, List, Union
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.utilities.advanced_parser.envelopes.base import BaseEnvelope
+from aws_lambda_powertools.utilities.advanced_parser.schemas import SqsSchema
+
+logger = logging.getLogger(__name__)
+
+
+# returns a list of parsed schemas of type BaseModel or plain string.
+# The record's body parameter is a string. However, it can also be a JSON encoded string which
+# can then be parsed into a BaseModel object.
+# Note that all records will be parsed the same way so if schema is str,
+# all the items in the list will be parsed as str and npt as JSON (and vice versa).
+class SqsEnvelope(BaseEnvelope):
+    def parse(self, event: Dict[str, Any], schema: Union[BaseModel, str]) -> List[Union[BaseModel, str]]:
+        try:
+            parsed_envelope = SqsSchema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input sqs event")
+            raise
+        output = []
+        for record in parsed_envelope.Records:
+            output.append(self._parse_user_json_string_schema(record.body, schema))
+        return output

aws_lambda_powertools/utilities/advanced_parser/parser.py

@@ -0,0 +1,68 @@
+import logging
+from typing import Any, Callable, Dict, Optional
+
+from pydantic import BaseModel, ValidationError
+
+from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+from aws_lambda_powertools.utilities.advanced_parser.envelopes import Envelope, parse_envelope
+
+logger = logging.getLogger(__name__)
+
+
+@lambda_handler_decorator
+def parser(
+    handler: Callable[[Dict, Any], Any],
+    event: Dict[str, Any],
+    context: Dict[str, Any],
+    schema: BaseModel,
+    envelope: Optional[Envelope] = None,
+) -> Any:
+    """Decorator to conduct advanced parsing & validation for lambda handlers events
+
+        As Lambda follows (event, context) signature we can remove some of the boilerplate
+        and also capture any exception any Lambda function throws as metadata.
+        event will be the parsed and passed as a BaseModel pydantic class of the input type "schema" 
+        to the lambda handler.
+        event will be extracted from the envelope in case envelope is not None.
+        In case envelope is None, the complete event is parsed to match the schema parameter BaseModel definition.
+        In case envelope is not None, first the event is parsed as the envelope's schema definition, and the user 
+        message is extracted and parsed again as the schema parameter's definition.
+
+        Example
+        -------
+        **Lambda function using validation decorator**
+
+            @parser(schema=MyBusiness, envelope=envelopes.EVENTBRIDGE)
+            def handler(event: MyBusiness , context: LambdaContext):
+                ...
+
+        Parameters
+        ----------
+        handler:  input for lambda_handler_decorator, wraps the handler lambda
+        event:    AWS event dictionary
+        context:  AWS lambda context
+        schema:   pydantic BaseModel class. This is the user data schema that will replace the event. 
+                  event parameter will be parsed and a new schema object will be created from it.
+        envelope: what envelope to extract the schema from, can be any AWS service that is currently 
+                  supported in the envelopes module. Can be None.
+
+        Raises
+        ------
+        err
+            TypeError - in case event is None
+            pydantic.ValidationError - event fails validation, either of the envelope
+    """
+    lambda_handler_name = handler.__name__
+    parsed_event = None
+    if envelope is None:
+        try:
+            logger.debug("Parsing and validating event schema, no envelope is used")
+            parsed_event = schema(**event)
+        except (ValidationError, TypeError):
+            logger.exception("Validation exception received from input event")
+            raise
+    else:
+        parsed_event = parse_envelope(event, envelope, schema)
+
+    logger.debug(f"Calling handler {lambda_handler_name}")
+    return handler(parsed_event, context)

aws_lambda_powertools/utilities/advanced_parser/schemas/__init__.py

@@ -0,0 +1,12 @@
+from .dynamodb import DynamoDBSchema, DynamoRecordSchema, DynamoScheme
+from .event_bridge import EventBridgeSchema
+from .sqs import SqsRecordSchema, SqsSchema
+
+__all__ = [
+    "DynamoDBSchema",
+    "EventBridgeSchema",
+    "DynamoScheme",
+    "DynamoRecordSchema",
+    "SqsSchema",
+    "SqsRecordSchema",
+]