feat(event-handler): context support to share data between routers (#1567)

Author: heitorlessa

aws_lambda_powertools/event_handler/api_gateway.py

@@ -237,6 +237,7 @@ def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dic
 class BaseRouter(ABC):
     current_event: BaseProxyEvent
     lambda_context: LambdaContext
+    context: dict
 
     @abstractmethod
     def route(
@@ -383,6 +384,14 @@ def lambda_handler(event, context):
         """
         return self.route(rule, "PATCH", cors, compress, cache_control)
 
+    def append_context(self, **additional_context):
+        """Append key=value data as routing context"""
+        self.context.update(**additional_context)
+
+    def clear_context(self):
+        """Resets routing context"""
+        self.context.clear()
+
 
 class ApiGatewayResolver(BaseRouter):
     """API Gateway and ALB proxy resolver
@@ -448,6 +457,7 @@ def __init__(
             env=os.getenv(constants.EVENT_HANDLER_DEBUG_ENV, "false"), choice=debug
         )
         self._strip_prefixes = strip_prefixes
+        self.context: Dict = {}  # early init as customers might add context before event resolution
 
         # Allow for a custom serializer or a concise json serialization
         self._serializer = serializer or partial(json.dumps, separators=(",", ":"), cls=Encoder)
@@ -502,11 +512,17 @@ def resolve(self, event, context) -> Dict[str, Any]:
                 "You don't need to serialize event to Event Source Data Class when using Event Handler; see issue #1152"
             )
             event = event.raw_event
+
         if self._debug:
             print(self._json_dump(event), end="")
+
+        # Populate router(s) dependencies without keeping a reference to each registered router
         BaseRouter.current_event = self._to_proxy_event(event)
         BaseRouter.lambda_context = context
-        return self._resolve().build(self.current_event, self._cors)
+
+        response = self._resolve().build(self.current_event, self._cors)
+        self.clear_context()
+        return response
 
     def __call__(self, event, context) -> Any:
         return self.resolve(event, context)
@@ -705,7 +721,7 @@ def _json_dump(self, obj: Any) -> str:
         return self._serializer(obj)
 
     def include_router(self, router: "Router", prefix: Optional[str] = None) -> None:
-        """Adds all routes defined in a router
+        """Adds all routes and context defined in a router
 
         Parameters
         ----------
@@ -718,6 +734,11 @@ def include_router(self, router: "Router", prefix: Optional[str] = None) -> None
         # Add reference to parent ApiGatewayResolver to support use cases where people subclass it to add custom logic
         router.api_resolver = self
 
+        # Merge app and router context
+        self.context.update(**router.context)
+        # use pointer to allow context clearance after event is processed e.g., resolve(evt, ctx)
+        router.context = self.context
+
         for route, func in router._routes.items():
             if prefix:
                 rule = route[0]
@@ -733,6 +754,7 @@ class Router(BaseRouter):
     def __init__(self):
         self._routes: Dict[tuple, Callable] = {}
         self.api_resolver: Optional[BaseRouter] = None
+        self.context = {}  # early init as customers might add context before event resolution
 
     def route(
         self,

aws_lambda_powertools/event_handler/appsync.py

@@ -12,6 +12,7 @@
 class BaseRouter:
     current_event: AppSyncResolverEventT  # type: ignore[valid-type]
     lambda_context: LambdaContext
+    context: dict
 
     def __init__(self):
         self._resolvers: dict = {}
@@ -34,6 +35,14 @@ def register_resolver(func):
 
         return register_resolver
 
+    def append_context(self, **additional_context):
+        """Append key=value data as routing context"""
+        self.context.update(**additional_context)
+
+    def clear_context(self):
+        """Resets routing context"""
+        self.context.clear()
+
 
 class AppSyncResolver(BaseRouter):
     """
@@ -68,6 +77,7 @@ def common_field() -> str:
 
     def __init__(self):
         super().__init__()
+        self.context = {}  # early init as customers might add context before event resolution
 
     def resolve(
         self, event: dict, context: LambdaContext, data_model: Type[AppSyncResolverEvent] = AppSyncResolverEvent
@@ -144,8 +154,12 @@ def lambda_handler(event, context):
         # Maintenance: revisit generics/overload to fix [attr-defined] in mypy usage
         BaseRouter.current_event = data_model(event)
         BaseRouter.lambda_context = context
+
         resolver = self._get_resolver(BaseRouter.current_event.type_name, BaseRouter.current_event.field_name)
-        return resolver(**BaseRouter.current_event.arguments)
+        response = resolver(**BaseRouter.current_event.arguments)
+        self.clear_context()
+
+        return response
 
     def _get_resolver(self, type_name: str, field_name: str) -> Callable:
         """Get resolver for field_name
@@ -182,9 +196,15 @@ def include_router(self, router: "Router") -> None:
         router : Router
             A router containing a dict of field resolvers
         """
+        # Merge app and router context
+        self.context.update(**router.context)
+        # use pointer to allow context clearance after event is processed e.g., resolve(evt, ctx)
+        router.context = self.context
+
         self._resolvers.update(router._resolvers)
 
 
 class Router(BaseRouter):
     def __init__(self):
         super().__init__()
+        self.context = {}  # early init as customers might add context before event resolution

docs/core/event_handler/api_gateway.md

@@ -449,6 +449,28 @@ When necessary, you can set a prefix when including a router object. This means
     --8<-- "examples/event_handler_rest/src/split_route_prefix_module.py"
     ```
 
+#### Sharing contextual data
+
+You can use `append_context` when you want to share data between your App and Router instances. Any data you share will be available via the `context` dictionary available in your App or Router context.
+
+???+ info
+    For safety, we always clear any data available in the `context` dictionary after each invocation.
+
+???+ tip
+    This can also be useful for middlewares injecting contextual information before a request is processed.
+
+=== "split_route_append_context.py"
+
+	```python hl_lines="18"
+    --8<-- "examples/event_handler_rest/src/split_route_append_context.py"
+	```
+
+=== "split_route_append_context_module.py"
+
+	```python hl_lines="16"
+    --8<-- "examples/event_handler_rest/src/split_route_append_context_module.py"
+	```
+
 #### Sample layout
 
 This is a sample project layout for a monolithic function with routes split in different files (`/todos`, `/health`).

docs/core/event_handler/appsync.md

@@ -226,6 +226,28 @@ Let's assume you have `split_operation.py` as your Lambda function entrypoint an
     --8<-- "examples/event_handler_graphql/src/split_operation.py"
     ```
 
+#### Sharing contextual data
+
+You can use `append_context` when you want to share data between your App and Router instances. Any data you share will be available via the `context` dictionary available in your App or Router context.
+
+???+ info
+    For safety, we always clear any data available in the `context` dictionary after each invocation.
+
+???+ tip
+    This can also be useful for middlewares injecting contextual information before a request is processed.
+
+=== "split_route_append_context.py"
+
+	```python hl_lines="17"
+    --8<-- "examples/event_handler_graphql/src/split_operation_append_context.py"
+	```
+
+=== "split_route_append_context_module.py"
+
+	```python hl_lines="29"
+    --8<-- "examples/event_handler_graphql/src/split_operation_append_context_module.py"
+	```
+
 ## Testing your code
 
 You can test your resolvers by passing a mocked or actual AppSync Lambda event that you're expecting.

examples/event_handler_graphql/src/split_operation_append_context.py

@@ -0,0 +1,18 @@
+import split_operation_append_context_module
+
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler import AppSyncResolver
+from aws_lambda_powertools.logging import correlation_paths
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+tracer = Tracer()
+logger = Logger()
+app = AppSyncResolver()
+app.include_router(split_operation_append_context_module.router)
+
+
+@logger.inject_lambda_context(correlation_id_path=correlation_paths.APPSYNC_RESOLVER)
+@tracer.capture_lambda_handler
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    app.append_context(is_admin=True)  # arbitrary number of key=value data
+    return app.resolve(event, context)

examples/event_handler_graphql/src/split_operation_append_context_module.py

@@ -0,0 +1,30 @@
+import sys
+
+if sys.version_info >= (3, 8):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+from typing import List
+
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler.appsync import Router
+
+tracer = Tracer()
+logger = Logger()
+router = Router()
+
+
+class Location(TypedDict, total=False):
+    id: str  # noqa AA03 VNE003, required due to GraphQL Schema
+    name: str
+    description: str
+    address: str
+
+
+@router.resolver(field_name="listLocations")
+@router.resolver(field_name="locations")
+@tracer.capture_method
+def get_locations(name: str, description: str = "") -> List[Location]:  # match GraphQL Query arguments
+    is_admin: bool = router.context.get("is_admin", False)
+    return [{"name": name, "description": description}] if is_admin else []

examples/event_handler_rest/src/split_route_append_context.py

@@ -0,0 +1,19 @@
+import split_route_append_context_module
+
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler import APIGatewayRestResolver
+from aws_lambda_powertools.logging import correlation_paths
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+tracer = Tracer()
+logger = Logger()
+app = APIGatewayRestResolver()
+app.include_router(split_route_append_context_module.router)
+
+
+# You can continue to use other utilities just as before
+@logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
+@tracer.capture_lambda_handler
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    app.append_context(is_admin=True)  # arbitrary number of key=value data
+    return app.resolve(event, context)

examples/event_handler_rest/src/split_route_append_context_module.py

@@ -0,0 +1,25 @@
+import requests
+from requests import Response
+
+from aws_lambda_powertools import Tracer
+from aws_lambda_powertools.event_handler.api_gateway import Router
+
+tracer = Tracer()
+router = Router()
+
+endpoint = "https://jsonplaceholder.typicode.com/todos"
+
+
+@router.get("/todos")
+@tracer.capture_method
+def get_todos():
+    is_admin: bool = router.context.get("is_admin", False)
+    todos = {}
+
+    if is_admin:
+        todos: Response = requests.get(endpoint)
+        todos.raise_for_status()
+        todos = todos.json()[:10]
+
+    # for brevity, we'll limit to the first 10 only
+    return {"todos": todos}

tests/functional/event_handler/test_api_gateway.py

@@ -1296,3 +1296,66 @@ def test_response_with_status_code_only():
     assert ret.status_code == 204
     assert ret.body is None
     assert ret.headers == {}
+
+
+def test_append_context():
+    app = APIGatewayRestResolver()
+    app.append_context(is_admin=True)
+    assert app.context.get("is_admin") is True
+
+
+def test_router_append_context():
+    router = Router()
+    router.append_context(is_admin=True)
+    assert router.context.get("is_admin") is True
+
+
+def test_route_context_is_cleared_after_resolve():
+    # GIVEN a Http API V1 proxy type event
+    app = APIGatewayRestResolver()
+    app.append_context(is_admin=True)
+
+    @app.get("/my/path")
+    def my_path():
+        return {"is_admin": app.context["is_admin"]}
+
+    # WHEN event resolution kicks in
+    app.resolve(LOAD_GW_EVENT, {})
+
+    # THEN context should be empty
+    assert app.context == {}
+
+
+def test_router_has_access_to_app_context(json_dump):
+    # GIVEN a Router with registered routes
+    app = ApiGatewayResolver()
+    router = Router()
+    ctx = {"is_admin": True}
+
+    @router.get("/my/path")
+    def my_path():
+        return {"is_admin": router.context["is_admin"]}
+
+    app.include_router(router)
+
+    # WHEN context is added and event resolution kicks in
+    app.append_context(**ctx)
+    ret = app.resolve(LOAD_GW_EVENT, {})
+
+    # THEN response include initial context
+    assert ret["body"] == json_dump(ctx)
+    assert router.context == {}
+
+
+def test_include_router_merges_context():
+    # GIVEN
+    app = APIGatewayRestResolver()
+    router = Router()
+
+    # WHEN
+    app.append_context(is_admin=True)
+    router.append_context(product_access=True)
+
+    app.include_router(router)
+
+    assert app.context == router.context