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

[walmsles]

Issue number: #953 #1955 #1696

Summary

Adds middleware support to event_handlers for REST API Gateway resolvers (all types). First pass POC implementation following a similar implementation style to Chalice and Starlette

• need to add documentation
• functional tests in place
• Reference middleware implementations provided for CORSMiddleware, CacheControlMiddleware as an alternative to existing ResponseBuilder implementation - existing ResponseBuilder implementation unchanged

Changes

• BaseRouter route functions have additional middleware parameter expecting a list of Callable middleware handlers
• Added special case registered_api_middleware, which is added as the last middleware handler internally (maintain backward compatibility) - without this function, the signature for REST routes would change to match the middleware signature, which feels like a breaking change. This special middleware also ensures all middleware will get a Response object as the response, which is essential for being able to write middleware (no type guessing).
• Added example Core middleware as Class based middleware - CORSMiddleware, CacheControlMiddleware
• Middleware can be any Python Callable
• Capability for adding global middleware (all routes) and route-specific middleware. Global middleware is always executed first since it represents cross-cutting concerns; anything else should be route specific.
• No ordering gymnastics - Middleware is applied in Order of definition, router-middleware + route-middleware

User experience

Users can now add middleware either for a specific route or centrally for ALL routes:

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 NextMiddleware

app = APIGatewayRestResolver()
logger = Logger()


def inject_correlation_id(app: APIGatewayRestResolver, next_middleware: NextMiddleware) -> Response:
    request_id = app.current_event.request_context.request_id

    # Use API Gateway REST API request ID if caller didn't include a correlation ID
    correlation_id = app.current_event.headers.get("x-correlation-id", request_id)

    # Inject correlation ID in shared context and Logger
    app.append_context(correlation_id=correlation_id)
    logger.set_correlation_id(request_id)

    # Get response from next middleware OR /todos route
    result = next_middleware(app)

    # Include Correlation ID in the response back to caller
    result.headers["x-correlation-id"] = correlation_id
    return result


@app.get("/todos", middlewares=[inject_correlation_id])
def get_todos():
    todos: 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)

Technical Debt

The following elements are listed as required technical debt to avoid breaking changes for existing REST API resolver REST users.

1. registered_api_adapter - this middleware adapter will ensure the existing API route functions will be called in the same way they are defined today, and it MUST be the last middleware handler in the Middleware call stack - LOTS of comments and context around where this happens when building the middleware call stack.
2. To avoid circular module references (because api_gateway.py module is full of self-references) unable to add proper Type Hinting to registered_api_adapter in aws_lambda_powertools/event_handler/middlewares/registered_api.py, but also wanted to place the middleware handler in the middleware folder where it belongs - could relocate this into api_gateway.py module (since is internal to powertools).
3. Return type of MiddlewareStackWrapper.__call__ is Union[Dict, Tuple, Response]. Ideally this would be consistent as Response - but not possible since API routes can return one of those 3 types - unsure if that could ever change but feel worth calling it out.

Checklist

If your change doesn't seem to apply, please leave them unchecked.

• Meet tenets criteria
• I have performed a self-review of this change
• Changes have been tested
• Changes are documented
• PR title follows conventional commit semantics

Is this a breaking change?

No, not a breaking change (I hope)

Checklist:

• Migration process documented
• Implement warnings (if it can live side by side)

Acknowledgment

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Disclaimer: We value your time and bandwidth. As such, any pull requests created on non-triaged issues might not be successful.

[heitorlessa]

This is so exciting!! I've raised a few questions on areas that weren't immediately clear to me, tiny improvements, and one around implementing a short-circuiting for middlewares to return early.

I'll review tests later. Thank you so much for taking this on @walmsles

[walmsles]

Question: I am thinking of relocating the internal registered_api_adapter into the api_gateway.py module and enforce typing of middleware to be inherited from the BaseMiddlewareHandler - I feel enforcing class usage for middleware allows for easier construction and configuration of middleware handlers. Much simpler to construct and reason about IMO.

[heitorlessa]

go for it!

[walmsles]

@heitorlessa have improved the BaseMiddlewareHandler class and implemented SchemaValidationMiddleware (based on the existing Powertools validate function). Enforcing Class based middleware is not possible due to the cyclic nature of api_gateway.py and middleware dependencies - Leaving it as "Callable" will suit a wider audience of devs and is required for the registered_api_adapter to work.

Let me know of any updates required - The code is done and now the hard part - documentation

[heitorlessa]

Thanks a lot!! I'll look into the changes tomorrow :)

[codecov-commenter]

Codecov Report

Patch coverage: 89.30% and project coverage change: -0.16% ⚠️

Comparison is base (d3fa713) 96.51% compared to head (20c344a) 96.36%.
Report is 3 commits behind head on develop.

❗ Current head 20c344a differs from pull request most recent head 7fbb132. Consider uploading reports for the commit 7fbb132 to get more accurate results

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2917      +/-   ##
===========================================
- Coverage    96.51%   96.36%   -0.16%     
===========================================
  Files          180      184       +4     
  Lines         7909     8050     +141     
  Branches      1491     1506      +15     
===========================================
+ Hits          7633     7757     +124     
- Misses         221      236      +15     
- Partials        55       57       +2     

Files Changed	Coverage Δ
...ols/event_handler/middlewares/schema_validation.py	76.74% <76.74%> (ø)
...ambda_powertools/event_handler/middlewares/base.py	88.88% <88.88%> (ø)
aws_lambda_powertools/event_handler/api_gateway.py	98.76% <94.18%> (-1.24%)	⬇️
...a_powertools/event_handler/middlewares/__init__.py	100.00% <100.00%> (ø)
aws_lambda_powertools/event_handler/types.py	100.00% <100.00%> (ø)
aws_lambda_powertools/shared/types.py	100.00% <100.00%> (ø)
...lambda_powertools/utilities/data_classes/common.py	100.00% <100.00%> (ø)
...a_powertools/utilities/data_classes/vpc_lattice.py	93.47% <100.00%> (ø)

☔ View full report in Codecov by Sentry.

📢 Have feedback on the report? Share it here.

[heitorlessa]

thanks a lot for those changes! Managed to spend more quality time to debug and understand this now - made some initial asks to start to be ready for prod.

Adding tasks so we get this to the finish line - we can divide and conquer (let me know the ones you want to pick up)

• Change docstrings to Numpy (our style)
• Add logger.debug statements in key areas to improve observability
• Include current middleware name and next in line in the call chain (suggested)
• Include debugging for anyone printing get_response to find out the call chain
• Address comments and/or suggestions
• Add missing tests for middlewares with Router
• Add missing tests for middlewares short circuiting via HTTP Errors
• Improve early return test by having a third middleware to ensure it doesn't get called
• Make attributes private unless we need to expose them (if so, docstrings)
• Update docstrings with examples where relevant
• Expand docstring for registered_api_adapter to provide an example on breaking change (hard to get ull picture if glancing through without debugging)
• Update docstring on how middlewares stack is built with an ascii diagram (who takes precedence, etc. - monodraw might be the best tool)
• Consider enriching exceptions with registered middlewares to ease debugging (which middlewares were registered? what was the call chain?) -- best effort
• Create diagram for happy path - Middleware at route level and router level (use()) (@walmsles WIP)
• Create diagram for happy path - Middleware defined in both main and split routes (@walmsles WIP)
• Create diagram for early return - short circuiting (@walmsles WIP)
• Document patterns for authoring middlewares - being a good citizen
• Document anti-patterns (e.g., avoid discarding responses in first middleware unless intentional)
• Document design in alignment with Chalice

[walmsles]

thanks a lot for those changes! Managed to spend more quality time to debug and understand this now - made some initial asks to start to be ready for prod.

Let's tick off some of these items :-).

Adding tasks so we get this to the finish line - we can divide and conquer (let me know the ones you want to pick up)

• Change docstrings to Numpy (our style)
• Add logger.debug statements in key areas to improve observability
• Include current middleware name and next in line in the call chain (suggested)
• Include debugging for anyone printing get_response to find out the call chain
• Address comments and/or suggestions
• Add missing tests for middlewares with Router
• Add missing tests for middlewares short circuiting via HTTP Errors
• Improve early return test by having a third middleware to ensure it doesn't get called
• Make attributes private unless we need to expose them (if so, docstrings)
• Update docstrings with examples where relevant
• Expand docstring for registered_api_adapter to provide an example on breaking change (hard to get full picture if glancing through without debugging)
• Update docstring on how middlewares stack is built with an ascii diagram (who takes precedence, etc. -- monodraw might be the best tool)
• Consider enriching exceptions with registered middlewares to ease debugging (which middlewares were registered? what was the call chain?) -- best effort
• Create diagram for happy path - Middleware at route level and router level (use())
• Create diagram for happy path - Middleware defined in both main and split routes
• Create diagram for early return - short circuiting
• Document patterns for authoring middlewares - being a good citizen
• Document anti-patterns (e.g., avoid discarding responses in first middleware unless intentional)
• Document design in alignment with Chalice

@heitorlessa lets discuss enriching exceptions techniques to assist with best-effort debugging
Started on some diagrams for the documentation side of things.

[walmsles]

@heitorlessa last round of commits added:

• additional debug in middleware code
• array tracking frame processing
• Additional code to support separate Router instances (via include_router)
• Added source draw.io document (advise if source is okay) for current WIP diagrams for middleware docs
• improved testing - Router, additional middleware in early-return tests, test for HTTPError return (exception handling)

Still to do:

• Change docstrings to Numpy (our style)
• Add logger.debug statements in key areas to improve observability
• Include current middleware name and next in line in the call chain (suggested)
• Include debugging for anyone printing get_response to find out the call chain
• Address comments and/or suggestions
• Add missing tests for middlewares with Router
• Add missing tests for middlewares short circuiting via HTTP Errors
• Improve early return test by having a third middleware to ensure it doesn't get called
• Make attributes private unless we need to expose them (if so, docstrings)
• Update docstrings with examples where relevant
• Expand docstring for registered_api_adapter to provide an example on breaking change (hard to get full picture if glancing through without debugging)
• Update docstring on how middlewares stack is built with an ascii diagram (who takes precedence, etc. -- monodraw might be the best tool)
• Consider enriching exceptions with registered middlewares to ease debugging (which middlewares were registered? what was the call chain?) -- best effort
• Create diagram for happy path - Middleware at route level and router level (use()) (@walmsles WIP)
• Create diagram for happy path - Middleware defined in both main and split routes (@walmsles WIP)
• Create diagram for early return - short circuiting (@walmsles WIP)
• Document patterns for authoring middlewares - being a good citizen
• Document anti-patterns (e.g., avoid discarding responses in first middleware unless intentional)
• Document design in alignment with Chalice

• Unsure how to tackle docstring ascii diagram that makes sense in low-fi like that - happy for ideas
• Need help with exception enrichment idea - happy for inputs here.
• I feel a thorough review of the functionality is in order and feedback - need help with debug places - I added some, but there may be more I am missing.
• I need input on Doc Strings for DX - this needs more eyes.
• I plan to put together a starting point on documentation this week.

[heitorlessa]

Quick notes through peer reviewing:

• Correct _executed_middleware as it's not including the following middleware that might have returned early/raised exception
• We no longer need the try/catch within middleware stack - bubble that up since we're catching all of it within API Processor
• Let's stick with the alternative image in drawio for Middlewares
• We need to add labels/legend to the boxes: Request, Response, Uncalled, Exception
• Investigate how to provide the middleware stack info during an exception
• Missing test for unhandled and non-HTTP exception e.g., raise RuntimeError("oops") inside any middleware
  • This will help clarify how we can make debugging easier for customers

[walmsles]

Quick notes through peer reviewing:

• Correct _executed_middleware as it's not including the following middleware that might have returned early/raised exception

Renamed to _push_middleware_stack_frame, relocated to MiddlewareFrame.__call__() method (where it needs to be).

• We no longer need the try/catch within middleware stack - bubble that up since we're catching all of it within API Processor

Removed - less code is better code!

• Let's stick with the alternative image in drawio for Middlewares
• We need to add labels/legend to the boxes: Request, Response, Uncalled, Exception
• Investigate how to provide the middleware stack info during an exception

@heitorlessa - please think about including processed stack frames into he debug exception response - I have struggled with getting stack trace into JSON formattable string to get something close to:

{
    "body": { 
        "stack_trace": ".... stack trace",
        "middleware_stack": [
              "Middleware 'SchemaValidationMiddleware'.  Next call chain is: SchemaValidationMiddleware -> NextMiddleware"
         ]
    }
}

But it is just not working ☹️

[rubenfonseca]

I'm planning to use this middleware to build automatic Pydantic validation based on route and the function params. In order for this to work, I need to have access to the Route that was matched during my middleware. Do you know if there's a way to do it the way this is implemented, or can we consider explosing that to the Middlewares?

[walmsles]

I'm planning to use this middleware to build automatic Pydantic validation based on route and the function params. In order for this to work, I need to have access to the Route that was matched during my middleware. Do you know if there's a way to do it the way this is implemented, or can we consider explosing that to the Middlewares?

Hi @rubenfonseca - the only state passed between middleware is the Resolver (Router) object itself so middleware has access to the "app" instance created by the developer. The actual "Route" that is being resolved is not currently referenced by the Resolver (which does seem unusual).

I feel the best way to achieve this would be to add a reference to the resolved Route object to the additional context - which is already managed between invocations and feels like the right place for this type of reference since it is Context related.

Have pushed a change adding "powertools_route" to the additional context - let me know your thoughts.

[heitorlessa]

Looking at the docs <3

[heitorlessa]

I ran out of time as I explored variations of the docs.

My mental model now is that we could adjust the pictures for a transparent background, and start with an upfront simpler image to demonstrate what this is. This will allow customers already familiar with that model to simply skip to what code is needed to get it done.

I think we can break the existing content in subsections so we progressively teach readers one thing at a time without bloating - build upon each knowledge, vary structure so it's a more lively reading.

I'll get back to it next Monday as I return from PTO :)

---

stateDiagram
    direction LR

    EventHandler: Event Handler
    Before: Before logic
    Next: Next `get_response()`
    MiddlewareLoop: Middleware loop
    After: Response
    Response: Final response

    EventHandler --> Middleware
    state MiddlewareLoop {
        direction LR
        Middleware --> Before
        Before --> Next
        Next --> Middleware: More middlewares?
        Next --> After
    }
    EventHandler --> Response: Middlewares complete?

Loading

from typing import Any, Callable

from aws_lambda_powertools import Logger
from aws_lambda_powertools.event_handler import Response, APIGatewayRestResolver
from aws_lambda_powertools.event_handler.api_gateway import BaseRouter

import requests

app = APIGatewayRestResolver()
logger = Logger()


# need a better first example, just to show completeness
def log_incoming_request(app: BaseRouter, get_response: Callable[..., Any], **context_args) -> Response:
    # Do Before processing here
    logger.info("Received request...", path=app.current_event.path)  # (1)!

    # Get Next response
    result = get_response(app, **context_args)  # (2)

    # Do After processing here

    # return the response
    return result  # (3)


@app.get("/todos", middlewares=[log_incoming_request])
def get_todos():
    todos: 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]}

[heitorlessa]

Two tests for the middleware order to ensure we don't introduce breaking change. I'll call out that in the docs now so a customer can decide whether to use the default or reverse order.

Appending to context logic is the same to ensure we are only testing the order in which .include_router is called.

---

@pytest.mark.parametrize(
    "app, event",
    [
        (ApiGatewayResolver(proxy_type=ProxyEventType.APIGatewayProxyEvent), API_REST_EVENT),
        (APIGatewayRestResolver(), API_REST_EVENT),
        (APIGatewayHttpResolver(), API_RESTV2_EVENT),
    ],
)
def test_api_gateway_middleware_order_with_include_router_last(app: EventHandlerInstance, event):
    # GIVEN two global middlewares: one for App and one for Router
    router = Router()

    def global_app_middleware(app: EventHandlerInstance, next_middleware: NextMiddleware):
        middleware_order: List[str] = router.context.get("middleware_order", [])
        middleware_order.append("app")

        app.append_context(middleware_order=middleware_order)
        return next_middleware(app)

    def global_router_middleware(router: EventHandlerInstance, next_middleware: NextMiddleware):
        middleware_order: List[str] = router.context.get("middleware_order", [])
        middleware_order.append("router")

        router.append_context(middleware_order=middleware_order)
        return next_middleware(app)

    @router.get("/my/path")
    def dummy_route():
        middleware_order = app.context["middleware_order"]

        assert middleware_order[0] == "app"
        assert middleware_order[1] == "router"

        return Response(status_code=200, body="works!")

    # WHEN App global middlewares are registered first
    # followed by include_router

    router.use([global_router_middleware])  # mimics App importing Router
    app.include_router(router)
    app.use([global_app_middleware])

    # THEN resolving a request should start processing global Router middlewares first
    # due to insertion order
    result = app(event, {})

    assert result["statusCode"] == 200

---

@pytest.mark.parametrize(
    "app, event",
    [
        (ApiGatewayResolver(proxy_type=ProxyEventType.APIGatewayProxyEvent), API_REST_EVENT),
        (APIGatewayRestResolver(), API_REST_EVENT),
        (APIGatewayHttpResolver(), API_RESTV2_EVENT),
    ],
)
def test_api_gateway_middleware_order_with_include_router_first(app: EventHandlerInstance, event):
    # GIVEN two global middlewares: one for App and one for Router
    router = Router()

    def global_app_middleware(app: EventHandlerInstance, next_middleware: NextMiddleware):
        middleware_order: List[str] = router.context.get("middleware_order", [])
        middleware_order.append("app")

        app.append_context(middleware_order=middleware_order)
        return next_middleware(app)

    def global_router_middleware(router: EventHandlerInstance, next_middleware: NextMiddleware):
        middleware_order: List[str] = router.context.get("middleware_order", [])
        middleware_order.append("router")

        router.append_context(middleware_order=middleware_order)
        return next_middleware(app)

    @router.get("/my/path")
    def dummy_route():
        middleware_order = app.context["middleware_order"]

        assert middleware_order[0] == "router"
        assert middleware_order[1] == "app"

        return Response(status_code=200, body="works!")

    # WHEN App include router middlewares first
    # followed by App global middlewares registration

    router.use([global_router_middleware])  # mimics App importing Router
    app.include_router(router)

    app.use([global_app_middleware])

    # THEN resolving a request should start processing global Router middlewares first
    # due to insertion order
    result = app(event, {})

    assert result["statusCode"] == 200

[walmsles]

Two tests for the middleware order to ensure we don't introduce breaking change. I'll call out that in the docs now so a customer can decide whether to use the default or reverse order.

👌

Appending to context logic is the same to ensure we are only testing the order in which .include_router is called.

It is important to make it clear that the merging of global middleware happens when include_router is called. Does that make sense, or is it harder to understand?

For clarity, it may be easier to:

1. Store router middlewares during include_router in a separate included_router_middleware list internally
2. Actual merging then happens during _build_middleware_stack (which occurs just prior to processing the route for the first time).
3. Merge order would then be set: Parent Router Globals + Included Router Globals (in order of include) + Route specific middleware

Is the above simpler to reason about? The current implementation does the merging when include_router is called, so if you have:

    router.use([global_router_middleware])  # mimics App importing Router

    app.use([global_app_middleware])
    app.include_router(router)
    app.use([global_app_middleware_two])

You end up with the middleware order: global_app_middleware, global_router_middleware, global_app_middleware_two
Basically - in the order they are used or included on the app instance.

The alternative I outlined above would result in the following order: global_app_middleware, global_app_middleware_two, global_router_middleware

It is simpler to reason about but also out of order regarding the process of constructing the app routes - there are definite pros/cons for each.

[heitorlessa]

Merging after CI checks complete 🎉 Finally over @walmsles -- THANK YOU so much for all the time you've put into this to make it a reality.

I've cleaned up the last bits, added mentions of Middleware for Router, and am leaving this to another PR so I can get started on re:Invent stuff.

• allow .use() to accept single middleware or a list of middleware
• create a TypedDict for internal info we're adding to make it easier to access it (and document)

[sonarqubecloud]

Kudos, SonarCloud Quality Gate passed!   

0 Bugs
0 Vulnerabilities
0 Security Hotspots
0 Code Smells

No Coverage information
0.6% Duplication

[leandrodamascena]

Hey there, @heitorlessa and @walmsles! I must say, it was a blast watching this PR come together and seeing how you tackled all those decisions about code and UX best practices. Seriously impressive work!

[github-actions]

This is now released under 2.24.0 version!