feat(event_handler): add Middleware support for REST Event Handler

aws_lambda_powertools/event_handler/middlewares/__init__.py

@@ -0,0 +1,3 @@
+from aws_lambda_powertools.event_handler.middlewares.base import BaseMiddlewareHandler, NextMiddleware
+
+__all__ = ["BaseMiddlewareHandler", "NextMiddleware"]

aws_lambda_powertools/event_handler/middlewares/base.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Generic
+
+from typing_extensions import Protocol
+
+from aws_lambda_powertools.event_handler.api_gateway import Response
+from aws_lambda_powertools.event_handler.types import EventHandlerInstance
+
+
+class NextMiddleware(Protocol):
+    def __call__(self, app: EventHandlerInstance) -> Response:
+        """Protocol for callback regardless of next_middleware(app), get_response(app) etc"""
+        ...
+
+    def __name__(self) -> str:  # noqa A003
+        """Protocol for name of the Middleware"""
+        ...
+
+
+class BaseMiddlewareHandler(Generic[EventHandlerInstance], ABC):
+    """Base implementation for Middlewares to run code before and after in a chain.
+
+
+    This is the middleware handler function where middleware logic is implemented.
+    The next middleware handler is represented by `next_middleware`, returning a Response object.
+
+    Examples
+    --------
+
+    **Correlation ID Middleware**
+
+    ```python
+    import requests
+
+    from aws_lambda_powertools import Logger
+    from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response
+    from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, NextMiddleware
+
+    app = APIGatewayRestResolver()
+    logger = Logger()
+
+
+    class CorrelationIdMiddleware(BaseMiddlewareHandler):
+        def __init__(self, header: str):
+            super().__init__()
+            self.header = header
+
+        def handler(self, app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
+            # BEFORE logic
+            request_id = app.current_event.request_context.request_id
+            correlation_id = app.current_event.get_header_value(
+                name=self.header,
+                default_value=request_id,
+            )
+
+            # Call next middleware or route handler ('/todos')
+            response = next_middleware(app)
+
+            # AFTER logic
+            response.headers[self.header] = correlation_id
+
+            return response
+
+
+    @app.get("/todos", middlewares=[CorrelationIdMiddleware(header="x-correlation-id")])
+    def get_todos():
+        todos: requests.Response = requests.get("https://jsonplaceholder.typicode.com/todos")
+        todos.raise_for_status()
+
+        # for brevity, we'll limit to the first 10 only
+        return {"todos": todos.json()[:10]}
+
+
+    @logger.inject_lambda_context
+    def lambda_handler(event, context):
+        return app.resolve(event, context)
+
+    ```
+
+    """
+
+    @abstractmethod
+    def handler(self, app: EventHandlerInstance, next_middleware: NextMiddleware) -> Response:
+        """
+        The Middleware Handler
+
+        Parameters
+        ----------
+        app: EventHandlerInstance
+            An instance of an Event Handler that implements ApiGatewayResolver
+        next_middleware: NextMiddleware
+            The next middleware handler in the chain
+
+        Returns
+        -------
+        Response
+            The response from the next middleware handler in the chain
+
+        """
+        raise NotImplementedError()
+
+    @property
+    def __name__(self) -> str:  # noqa A003
+        return str(self.__class__.__name__)
+
+    def __call__(self, app: EventHandlerInstance, next_middleware: NextMiddleware) -> Response:
+        """
+        The Middleware handler function.
+
+        Parameters
+        ----------
+        app: ApiGatewayResolver
+            An instance of an Event Handler that implements ApiGatewayResolver
+        next_middleware: NextMiddleware
+            The next middleware handler in the chain
+
+        Returns
+        -------
+        Response
+            The response from the next middleware handler in the chain
+        """
+        return self.handler(app, next_middleware)

aws_lambda_powertools/event_handler/middlewares/schema_validation.py

@@ -0,0 +1,124 @@
+import logging
+from typing import Dict, Optional
+
+from aws_lambda_powertools.event_handler.api_gateway import Response
+from aws_lambda_powertools.event_handler.exceptions import BadRequestError, InternalServerError
+from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, NextMiddleware
+from aws_lambda_powertools.event_handler.types import EventHandlerInstance
+from aws_lambda_powertools.utilities.validation import validate
+from aws_lambda_powertools.utilities.validation.exceptions import InvalidSchemaFormatError, SchemaValidationError
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaValidationMiddleware(BaseMiddlewareHandler):
+    """Middleware to validate API request and response against JSON Schema using the [Validation utility](https://docs.powertools.aws.dev/lambda/python/latest/utilities/validation/).
+
+    Examples
+    --------
+    **Validating incoming event**
+
+    ```python
+    import requests
+
+    from aws_lambda_powertools import Logger
+    from aws_lambda_powertools.event_handler import APIGatewayRestResolver, Response
+    from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler, NextMiddleware
+    from aws_lambda_powertools.event_handler.middlewares.schema_validation import SchemaValidationMiddleware
+
+    app = APIGatewayRestResolver()
+    logger = Logger()
+    json_schema_validation = SchemaValidationMiddleware(inbound_schema=INCOMING_JSON_SCHEMA)
+
+
+    @app.get("/todos", middlewares=[json_schema_validation])
+    def get_todos():
+        todos: requests.Response = requests.get("https://jsonplaceholder.typicode.com/todos")
+        todos.raise_for_status()
+
+        # for brevity, we'll limit to the first 10 only
+        return {"todos": todos.json()[:10]}
+
+
+    @logger.inject_lambda_context
+    def lambda_handler(event, context):
+        return app.resolve(event, context)
+    ```
+    """
+
+    def __init__(
+        self,
+        inbound_schema: Dict,
+        inbound_formats: Optional[Dict] = None,
+        outbound_schema: Optional[Dict] = None,
+        outbound_formats: Optional[Dict] = None,
+    ):
+        """See [Validation utility](https://docs.powertools.aws.dev/lambda/python/latest/utilities/validation/) docs for examples on all parameters.
+
+        Parameters
+        ----------
+        inbound_schema : Dict
+            JSON Schema to validate incoming event
+        inbound_formats : Optional[Dict], optional
+            Custom formats containing a key (e.g. int64) and a value expressed as regex or callback returning bool, by default None
+            JSON Schema to validate outbound event, by default None
+        outbound_formats : Optional[Dict], optional
+            Custom formats containing a key (e.g. int64) and a value expressed as regex or callback returning bool, by default None
+        """  # noqa: E501
+        super().__init__()
+        self.inbound_schema = inbound_schema
+        self.inbound_formats = inbound_formats
+        self.outbound_schema = outbound_schema
+        self.outbound_formats = outbound_formats
+
+    def bad_response(self, error: SchemaValidationError) -> Response:
+        message: str = f"Bad Response: {error.message}"
+        logger.debug(message)
+        raise BadRequestError(message)
+
+    def bad_request(self, error: SchemaValidationError) -> Response:
+        message: str = f"Bad Request: {error.message}"
+        logger.debug(message)
+        raise BadRequestError(message)
+
+    def bad_config(self, error: InvalidSchemaFormatError) -> Response:
+        logger.debug(f"Invalid Schema Format: {error}")
+        raise InternalServerError("Internal Server Error")
+
+    def handler(self, app: EventHandlerInstance, next_middleware: NextMiddleware) -> Response:
+        """Validates incoming JSON payload (body) against JSON Schema provided.
+
+        Parameters
+        ----------
+        app : EventHandlerInstance
+            An instance of an Event Handler
+        next_middleware : NextMiddleware
+            Callable to get response from the next middleware or route handler in the chain
+
+        Returns
+        -------
+        Response
+            It can return three types of response objects
+
+            - Original response: Propagates HTTP response returned from the next middleware if validation succeeds
+            - HTTP 400: Payload or response failed JSON Schema validation
+            - HTTP 500: JSON Schema provided has incorrect format
+        """
+        try:
+            validate(event=app.current_event.json_body, schema=self.inbound_schema, formats=self.inbound_formats)
+        except SchemaValidationError as error:
+            return self.bad_request(error)
+        except InvalidSchemaFormatError as error:
+            return self.bad_config(error)
+
+        result = next_middleware(app)
+
+        if self.outbound_formats is not None:
+            try:
+                validate(event=result.body, schema=self.inbound_schema, formats=self.inbound_formats)
+            except SchemaValidationError as error:
+                return self.bad_response(error)
+            except InvalidSchemaFormatError as error:
+                return self.bad_config(error)
+
+        return result

aws_lambda_powertools/event_handler/types.py

@@ -0,0 +1,5 @@
+from typing import TypeVar
+
+from aws_lambda_powertools.event_handler import ApiGatewayResolver
+
+EventHandlerInstance = TypeVar("EventHandlerInstance", bound=ApiGatewayResolver)

aws_lambda_powertools/shared/types.py

@@ -1,5 +1,14 @@
+import sys
 from typing import Any, Callable, Dict, List, TypeVar, Union
 
 AnyCallableT = TypeVar("AnyCallableT", bound=Callable[..., Any])  # noqa: VNE001
 # JSON primitives only, mypy doesn't support recursive tho
 JSONType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
+
+
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+__all__ = ["Protocol"]

aws_lambda_powertools/utilities/data_classes/common.py

@@ -1,7 +1,7 @@
 import base64
 import json
 from collections.abc import Mapping
-from typing import Any, Callable, Dict, Iterator, List, Optional
+from typing import Any, Callable, Dict, Iterator, List, Optional, overload
 
 from aws_lambda_powertools.shared.headers_serializer import BaseHeadersSerializer
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
@@ -156,7 +156,24 @@ def get_query_string_value(self, name: str, default_value: Optional[str] = None)
             default_value=default_value,
         )
 
-    # Maintenance: missing @overload to ensure return type is a str when default_value is set
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: str,
+        case_sensitive: Optional[bool] = False,
+    ) -> str:
+        ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: Optional[bool] = False,
+    ) -> Optional[str]:
+        ...
+
     def get_header_value(
         self,
         name: str,

aws_lambda_powertools/utilities/data_classes/vpc_lattice.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, overload
 
 from aws_lambda_powertools.shared.headers_serializer import (
     BaseHeadersSerializer,
@@ -91,6 +91,24 @@ def get_query_string_value(self, name: str, default_value: Optional[str] = None)
             default_value=default_value,
         )
 
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: str,
+        case_sensitive: Optional[bool] = False,
+    ) -> str:
+        ...
+
+    @overload
+    def get_header_value(
+        self,
+        name: str,
+        default_value: Optional[str] = None,
+        case_sensitive: Optional[bool] = False,
+    ) -> Optional[str]:
+        ...
+
     def get_header_value(
         self,
         name: str,