From 35bd80d7ea20f20911ce74db970db8ccf9787d82 Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 15:51:53 +1000 Subject: [PATCH 1/7] feat(data_classes): add authorizer and callback_url to API Gateway WebSocket event MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Groundwork for the API Gateway WebSocket event handler (#1165), kept as a separate commit so the data-class changes are reviewable on their own. The resolver's auth and send-to-client patterns need both properties: the repo already supports authoring WebSocket Lambda authorizers (APIGatewayAuthorizerResponseWebSocket); this adds the consuming side — the authorizer context API Gateway injects into every invocation for a connection — plus the Management API callback URL for replying to clients. --- .../api_gateway_websocket_event.py | 14 +++++++++++++ .../test_api_gateway_websocket_event.py | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/aws_lambda_powertools/utilities/data_classes/api_gateway_websocket_event.py b/aws_lambda_powertools/utilities/data_classes/api_gateway_websocket_event.py index bb93cac7fe2..e02ebd3a04e 100644 --- a/aws_lambda_powertools/utilities/data_classes/api_gateway_websocket_event.py +++ b/aws_lambda_powertools/utilities/data_classes/api_gateway_websocket_event.py @@ -69,6 +69,11 @@ def request_time_epoch(self) -> int: def identity(self) -> APIGatewayWebSocketEventIdentity: return APIGatewayWebSocketEventIdentity(self["identity"]) + @property + def authorizer(self) -> dict[str, Any]: + """Lambda authorizer context, injected on every invocation for the connection. Values are stringified.""" + return self.get("authorizer") or {} + @property def request_id(self) -> str: return self["requestId"] @@ -85,6 +90,15 @@ def connection_id(self) -> str: def api_id(self) -> str: return self["apiId"] + @property + def callback_url(self) -> str: + """Management API endpoint (`https://{domain_name}/{stage}`) for sending messages back to this connection. + + For connections made through a custom domain name, this URL is only valid when the domain's + API mapping path matches the stage name. + """ + return f"https://{self['domainName']}/{self['stage']}" + class APIGatewayWebSocketEvent(DictWrapper): """AWS proxy integration event for WebSocket API diff --git a/tests/unit/data_classes/required_dependencies/test_api_gateway_websocket_event.py b/tests/unit/data_classes/required_dependencies/test_api_gateway_websocket_event.py index 4151429ad50..4e4acfeccf4 100644 --- a/tests/unit/data_classes/required_dependencies/test_api_gateway_websocket_event.py +++ b/tests/unit/data_classes/required_dependencies/test_api_gateway_websocket_event.py @@ -36,6 +36,8 @@ def test_connect_api_gateway_websocket_event(): assert request_context.domain_name == request_context_raw["domainName"] assert request_context.connection_id == request_context_raw["connectionId"] assert request_context.api_id == request_context_raw["apiId"] + assert request_context.authorizer == {} + assert request_context.callback_url == f"https://{request_context_raw['domainName']}/{request_context_raw['stage']}" identity = request_context.identity identity_raw = request_context_raw["identity"] @@ -71,6 +73,8 @@ def test_disconnect_api_gateway_websocket_event(): assert request_context.domain_name == request_context_raw["domainName"] assert request_context.connection_id == request_context_raw["connectionId"] assert request_context.api_id == request_context_raw["apiId"] + assert request_context.authorizer == {} + assert request_context.callback_url == f"https://{request_context_raw['domainName']}/{request_context_raw['stage']}" identity = request_context.identity identity_raw = request_context_raw["identity"] @@ -108,8 +112,24 @@ def test_message_api_gateway_websocket_event(): assert request_context.domain_name == request_context_raw["domainName"] assert request_context.connection_id == request_context_raw["connectionId"] assert request_context.api_id == request_context_raw["apiId"] + assert request_context.authorizer == {} + assert request_context.callback_url == f"https://{request_context_raw['domainName']}/{request_context_raw['stage']}" identity = request_context.identity identity_raw = request_context_raw["identity"] assert identity.source_ip == identity_raw["sourceIp"] assert identity.user_agent is None + + +def test_lambda_authorizer_api_gateway_websocket_event(): + raw_event = load_event("apiGatewayWebSocketApiMessage.json") + raw_event["requestContext"]["authorizer"] = { + "principalId": "user123", + "tenantId": "tenant-a", + "numKey": "1", + "boolKey": "true", + } + parsed_event = APIGatewayWebSocketEvent(raw_event) + + assert parsed_event.request_context.authorizer == raw_event["requestContext"]["authorizer"] + assert parsed_event.request_context.authorizer["principalId"] == "user123" From af7dd46495d737b19c156d99752e5d78e3631791 Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 17:32:50 +1000 Subject: [PATCH 2/7] feat(event_handler): add API Gateway WebSocket resolver Standalone resolver for API Gateway WebSocket APIs (#1165), built on APIGatewayWebSocketEvent and mirroring the AppSync Events package structure: route()/on_connect/on_disconnect/on_default registration with exact route-key matching, response normalization to the statusCode/body shape API Gateway expects, and exception handling via the shared ExceptionHandlerManager with a catch-all that returns a bare 500 so exception details never reach connected clients. Includes functional tests (full coverage), docs page with getting started examples and SAM template, and mkdocs nav. --- .../event_handler/__init__.py | 2 + .../api_gateway_websocket/__init__.py | 7 + .../api_gateway_websocket/_registry.py | 87 ++++ .../api_gateway_websocket/base.py | 55 +++ .../api_gateway_websocket/router.py | 236 +++++++++ .../api_gateway_websocket/types.py | 22 + .../api_gateway_websocket/websocket.py | 216 +++++++++ .../event_handler/api_gateway_websocket.md | 132 +++++ .../getting_started_with_websocket_api.yaml | 100 ++++ .../src/getting_started_with_connect.py | 20 + .../src/getting_started_with_custom_route.py | 19 + .../src/getting_started_with_disconnect.py | 21 + .../src/getting_started_with_testing.py | 36 ++ .../getting_started_with_testing_event.json | 49 ++ .../src/working_with_exception_handling.py | 24 + mkdocs.yml | 2 + .../api_gateway_websocket/__init__.py | 0 .../test_api_gateway_websocket_resolver.py | 455 ++++++++++++++++++ 18 files changed, 1483 insertions(+) create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/__init__.py create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/_registry.py create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/base.py create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/router.py create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/types.py create mode 100644 aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py create mode 100644 docs/core/event_handler/api_gateway_websocket.md create mode 100644 examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml create mode 100644 examples/event_handler_api_gateway_websocket/src/getting_started_with_connect.py create mode 100644 examples/event_handler_api_gateway_websocket/src/getting_started_with_custom_route.py create mode 100644 examples/event_handler_api_gateway_websocket/src/getting_started_with_disconnect.py create mode 100644 examples/event_handler_api_gateway_websocket/src/getting_started_with_testing.py create mode 100644 examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_exception_handling.py create mode 100644 tests/functional/event_handler/required_dependencies/api_gateway_websocket/__init__.py create mode 100644 tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py diff --git a/aws_lambda_powertools/event_handler/__init__.py b/aws_lambda_powertools/event_handler/__init__.py index 2e4c008ec63..105a5443a56 100644 --- a/aws_lambda_powertools/event_handler/__init__.py +++ b/aws_lambda_powertools/event_handler/__init__.py @@ -10,6 +10,7 @@ CORSConfig, Response, ) +from aws_lambda_powertools.event_handler.api_gateway_websocket.websocket import APIGatewayWebSocketResolver from aws_lambda_powertools.event_handler.appsync import AppSyncResolver from aws_lambda_powertools.event_handler.bedrock_agent import BedrockAgentResolver, BedrockResponse from aws_lambda_powertools.event_handler.bedrock_agent_function import ( @@ -30,6 +31,7 @@ "AppSyncEventsResolver", "APIGatewayRestResolver", "APIGatewayHttpResolver", + "APIGatewayWebSocketResolver", "ALBResolver", "ApiGatewayResolver", "BedrockAgentResolver", diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/__init__.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/__init__.py new file mode 100644 index 00000000000..0abcdd31523 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/__init__.py @@ -0,0 +1,7 @@ +from aws_lambda_powertools.event_handler.api_gateway_websocket.router import Router +from aws_lambda_powertools.event_handler.api_gateway_websocket.websocket import APIGatewayWebSocketResolver + +__all__ = [ + "APIGatewayWebSocketResolver", + "Router", +] diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/_registry.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/_registry.py new file mode 100644 index 00000000000..5ee619f69c2 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/_registry.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import logging +import warnings +from typing import TYPE_CHECKING + +from aws_lambda_powertools.event_handler.api_gateway_websocket.types import WebSocketRoute +from aws_lambda_powertools.warnings import PowertoolsUserWarning + +if TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) + + +class RouteRegistry: + def __init__(self): + self.routes: dict[str, WebSocketRoute] = {} + + def register( + self, + route_key: str, + middlewares: list[Callable] | None = None, + ) -> Callable: + """Registers a route handler for a route key + + Parameters + ---------- + route_key : str + Route key produced by the API's route selection expression, e.g. `$connect` or a custom key + middlewares : list[Callable] | None + Middlewares to run around the handler for this route + + Return + ---------- + Callable + A decorator that registers the handler + """ + + def _register(func: Callable) -> Callable: + if not route_key: + warnings.warn( + f"The route key registered for `{getattr(func, '__name__', func)}` is empty and will be skipped.", + stacklevel=2, + category=PowertoolsUserWarning, + ) + return func + + if route_key in self.routes: + warnings.warn( + f"A route handler is already registered for route key `{route_key}`. " + "The last registration will be used.", + stacklevel=2, + category=PowertoolsUserWarning, + ) + + logger.debug(f"Adding route handler `{func.__name__}` for route key `{route_key}`") + self.routes[route_key] = WebSocketRoute(func=func, middlewares=middlewares or []) + return func + + return _register + + def find_route(self, route_key: str) -> WebSocketRoute | None: + """Find a route handler by exact route key match + + Parameters + ---------- + route_key : str + Route key to look up + + Return + ---------- + WebSocketRoute | None + The registered route, or None when no handler is registered for the route key + """ + logger.debug(f"Looking for route handler for route key `{route_key}`") + return self.routes.get(route_key) + + def merge(self, other_registry: RouteRegistry) -> None: + """Update current registry with routes from an incoming registry + + Parameters + ---------- + other_registry : RouteRegistry + Registry to merge from + """ + self.routes.update(**other_registry.routes) diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py new file mode 100644 index 00000000000..d8b45cf2b47 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + +ROUTE_KEY_CONNECT = "$connect" +ROUTE_KEY_DISCONNECT = "$disconnect" +ROUTE_KEY_DEFAULT = "$default" + + +class BaseRouter(ABC): + """Abstract base class for WebSocket routers (resolvers)""" + + @abstractmethod + def route( + self, + route_key: str, + middlewares: list[Callable] | None = None, + ) -> Callable: + raise NotImplementedError + + @abstractmethod + def on_connect( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + raise NotImplementedError + + @abstractmethod + def on_disconnect( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + raise NotImplementedError + + @abstractmethod + def on_default( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + raise NotImplementedError + + def append_context(self, **additional_context) -> None: + """ + Appends context information available under any route. + + Parameters + ----------- + **additional_context: dict + Additional context key-value pairs to append. + """ + raise NotImplementedError diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py new file mode 100644 index 00000000000..ee5306e5f10 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aws_lambda_powertools.event_handler.api_gateway_websocket._registry import RouteRegistry +from aws_lambda_powertools.event_handler.api_gateway_websocket.base import ( + ROUTE_KEY_CONNECT, + ROUTE_KEY_DEFAULT, + ROUTE_KEY_DISCONNECT, + BaseRouter, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from aws_lambda_powertools.utilities.data_classes.api_gateway_websocket_event import APIGatewayWebSocketEvent + from aws_lambda_powertools.utilities.typing.lambda_context import LambdaContext + + +class Router(BaseRouter): + """ + Router for API Gateway WebSocket API event handling. + + Registers route handlers for WebSocket route keys so they can be split across + files and later included in an `APIGatewayWebSocketResolver` via `include_router`. + + Parameters + ---------- + context : dict + Dictionary to store context information accessible across route handlers + current_event : APIGatewayWebSocketEvent + Current event being processed + lambda_context : LambdaContext + Lambda context from the AWS Lambda function + + Examples + -------- + Create a router and define route handlers: + + >>> from aws_lambda_powertools.event_handler.api_gateway_websocket import Router + >>> + >>> orders_router = Router() + >>> + >>> @orders_router.route("orderUpdate") + >>> def order_update(): + >>> order = orders_router.current_event.json_body + >>> return {"orderId": order["id"], "status": "updated"} + """ + + context: dict + current_event: APIGatewayWebSocketEvent + lambda_context: LambdaContext + + def __init__(self): + self.context = {} # early init as customers might add context before event resolution + self._route_registry = RouteRegistry() + self._exception_handlers: dict[type[Exception], Callable] = {} + + def route( + self, + route_key: str, + middlewares: list[Callable] | None = None, + ) -> Callable: + """ + Register a handler for a WebSocket route key. + + Route keys are matched exactly against the key selected by the API's + route selection expression — there is no pattern matching. + + Parameters + ---------- + route_key : str + The route key to register the handler for, e.g. `$connect`, `$disconnect`, + `$default`, or a custom route key such as `orderUpdate` + middlewares : list[Callable], optional + Middlewares to run around the handler for this route + + Returns + ------- + Callable + Decorator function that registers the route handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.route("orderUpdate") + >>> def order_update(): + >>> return {"received": app.current_event.json_body["id"]} + """ + return self._route_registry.register(route_key=route_key, middlewares=middlewares) + + def on_connect( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + """ + Register a handler for the `$connect` route key. + + The returned status code decides whether API Gateway accepts the connection: + 2xx accepts the WebSocket upgrade, anything else rejects it. + + Parameters + ---------- + middlewares : list[Callable], optional + Middlewares to run around the handler for this route + + Returns + ------- + Callable + Decorator function that registers the connect handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.on_connect() + >>> def connect(): + >>> connection_id = app.current_event.request_context.connection_id + >>> return None # 200: accept the connection + """ + return self.route(route_key=ROUTE_KEY_CONNECT, middlewares=middlewares) + + def on_disconnect( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + """ + Register a handler for the `$disconnect` route key. + + This handler is invoked on a best-effort basis when the connection closes, + including for connections whose `$connect` was rejected. + + Parameters + ---------- + middlewares : list[Callable], optional + Middlewares to run around the handler for this route + + Returns + ------- + Callable + Decorator function that registers the disconnect handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.on_disconnect() + >>> def disconnect(): + >>> reason = app.current_event.request_context.disconnect_reason + """ + return self.route(route_key=ROUTE_KEY_DISCONNECT, middlewares=middlewares) + + def on_default( + self, + middlewares: list[Callable] | None = None, + ) -> Callable: + """ + Register a handler for the `$default` route key. + + Invoked when the route selection expression does not match any custom route, + or when the API has no custom routes. + + Parameters + ---------- + middlewares : list[Callable], optional + Middlewares to run around the handler for this route + + Returns + ------- + Callable + Decorator function that registers the default handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.on_default() + >>> def default(): + >>> return {"echo": app.current_event.body} + """ + return self.route(route_key=ROUTE_KEY_DEFAULT, middlewares=middlewares) + + def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) -> Callable: + """ + Register a handler for one or more exception types. + + The handler receives the exception and its return value goes through the + same response normalization as route handler return values. + + Parameters + ---------- + exc_class : type[Exception] | list[type[Exception]] + A single exception type or a list of exception types + + Returns + ------- + Callable + Decorator function that registers the exception handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.exception_handler(ValueError) + >>> def handle_invalid_value(exc: ValueError): + >>> return {"error": str(exc)}, 400 + """ + + def register_exception_handler(func: Callable) -> Callable: + if isinstance(exc_class, list): + for exp in exc_class: + self._exception_handlers[exp] = func + else: + self._exception_handlers[exc_class] = func + return func + + return register_exception_handler + + def append_context(self, **additional_context) -> None: + """Append key=value data as routing context""" + self.context.update(**additional_context) + + def clear_context(self) -> None: + """Resets routing context""" + self.context.clear() diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/types.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/types.py new file mode 100644 index 00000000000..36cfad9c6a3 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/types.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypedDict + +if TYPE_CHECKING: + from collections.abc import Callable + + +class WebSocketRoute(TypedDict): + """ + Type definition for a registered WebSocket route + + Parameters + ---------- + func: Callable[..., Any] + Route handler function + middlewares: list[Callable[..., Any]] + Middlewares to run around the handler for this route + """ + + func: Callable[..., Any] + middlewares: list[Callable[..., Any]] diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py new file mode 100644 index 00000000000..0496caeed12 --- /dev/null +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import json +import logging +import warnings +from http import HTTPStatus +from typing import TYPE_CHECKING, Any + +from aws_lambda_powertools.event_handler.api_gateway_websocket.router import Router +from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager +from aws_lambda_powertools.shared.json_encoder import Encoder +from aws_lambda_powertools.utilities.data_classes.api_gateway_websocket_event import APIGatewayWebSocketEvent +from aws_lambda_powertools.warnings import PowertoolsUserWarning + +if TYPE_CHECKING: + from collections.abc import Callable + + from aws_lambda_powertools.utilities.typing.lambda_context import LambdaContext + +logger = logging.getLogger(__name__) + + +class APIGatewayWebSocketResolver(Router): + """ + API Gateway WebSocket API resolver. + + Dispatches WebSocket events to handlers registered by route key (`$connect`, + `$disconnect`, `$default`, or custom route keys selected by the API's route + selection expression) and normalizes handler return values into the + `{"statusCode": ..., "body": ...}` shape API Gateway expects. + + Handlers take no arguments and access the request through `app.current_event` + (an `APIGatewayWebSocketEvent`), `app.lambda_context`, and `app.context`. + + Attributes + ---------- + context: dict + Dictionary to store context information accessible across route handlers + current_event: APIGatewayWebSocketEvent + Current event being processed + lambda_context: LambdaContext + Lambda context from the AWS Lambda function + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.on_connect() + >>> def connect(): + >>> return None # 200: accept the connection + >>> + >>> @app.on_disconnect() + >>> def disconnect(): + >>> connection_id = app.current_event.request_context.connection_id + >>> + >>> @app.route("orderUpdate") + >>> def order_update(): + >>> return {"orderId": app.current_event.json_body["id"]} + >>> + >>> @app.on_default() + >>> def default(): + >>> return {"echo": app.current_event.body} + >>> + >>> def lambda_handler(event, context): + >>> return app.resolve(event, context) + """ + + def __init__(self): + super().__init__() + self.exception_handler_manager = ExceptionHandlerManager() + + def __call__( + self, + event: dict | APIGatewayWebSocketEvent, + context: LambdaContext, + ) -> dict[str, Any]: + """Implicit lambda handler which internally calls `resolve`.""" + return self.resolve(event, context) + + def resolve( + self, + event: dict | APIGatewayWebSocketEvent, + context: LambdaContext, + ) -> dict[str, Any]: + """ + Resolve a WebSocket event to the handler registered for its route key. + + Parameters + ---------- + event: dict | APIGatewayWebSocketEvent + The API Gateway WebSocket event to process + context: LambdaContext + The Lambda context + + Returns + ------- + dict[str, Any] + The normalized Lambda response, always containing `statusCode` and + optionally `body`. On `$connect`, a non-2xx status code rejects the + connection; on other routes the body is delivered back to the client + only when the route has a route response configured. + + Examples + -------- + >>> app = APIGatewayWebSocketResolver() + >>> + >>> def lambda_handler(event, context): + >>> return app.resolve(event, context) + """ + try: + self._setup_context(event, context) + return self._resolve_route() + except Exception as exc: + return self._handle_exception(exc) + finally: + self.clear_context() + + def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) -> Callable: + """ + Register a handler for one or more exception types. + + The handler receives the exception and its return value goes through the + same response normalization as route handler return values. Exceptions + with no registered handler are logged and become a bare + `{"statusCode": 500}` so no exception details reach the client. + + Parameters + ---------- + exc_class : type[Exception] | list[type[Exception]] + A single exception type or a list of exception types + + Returns + ------- + Callable + Decorator function that registers the exception handler + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> @app.exception_handler(ValueError) + >>> def handle_invalid_value(exc: ValueError): + >>> return {"error": str(exc)}, 400 + """ + return self.exception_handler_manager.exception_handler(exc_class=exc_class) + + def _resolve_route(self) -> dict[str, Any]: + """Dispatch the current event to its route handler and normalize the response.""" + route_key = self.current_event.request_context.route_key + route = self._route_registry.find_route(route_key) + if route is None: + warnings.warn( + f"No route handler registered for route key `{route_key}`.", + stacklevel=2, + category=PowertoolsUserWarning, + ) + return {"statusCode": HTTPStatus.BAD_REQUEST.value} + + logger.debug(f"Dispatching route key `{route_key}` to `{route['func'].__name__}`") + return self._to_response(route["func"]()) + + def _handle_exception(self, exc: Exception) -> dict[str, Any]: + """Resolve an exception to a response via registered handlers, or a bare 500. + + An unhandled exception must never propagate to the Lambda runtime: the runtime's + default error response (exception type, message, and stack trace) is delivered + to the connected client. + """ + handler = self.exception_handler_manager.lookup_exception_handler(type(exc)) + if handler: + try: + return self._to_response(handler(exc)) + except Exception: + logger.exception("Exception handler raised an exception") + return {"statusCode": HTTPStatus.INTERNAL_SERVER_ERROR.value} + + logger.exception(f"Unhandled exception while resolving route key `{self._safe_route_key()}`") + return {"statusCode": HTTPStatus.INTERNAL_SERVER_ERROR.value} + + def _safe_route_key(self) -> str: + """Best-effort route key for logging, tolerating malformed events.""" + try: + return self.current_event.request_context.route_key + except Exception: + return "" + + def _to_response(self, result: Any) -> dict[str, Any]: + """Normalize a handler return value into the Lambda response shape. + + A 2-tuple sets the status code explicitly; any other value maps to 200. + A returned dict is always body content — it is never inspected for a + `statusCode` key. + """ + if isinstance(result, tuple) and len(result) == 2: + body, status_code = result + return self._format_response(body, int(status_code)) + return self._format_response(result, HTTPStatus.OK.value) + + def _format_response(self, body: Any, status_code: int) -> dict[str, Any]: + if body is None: + return {"statusCode": status_code} + if isinstance(body, str): + return {"statusCode": status_code, "body": body} + return {"statusCode": status_code, "body": json.dumps(body, separators=(",", ":"), cls=Encoder)} + + def _setup_context(self, event: dict | APIGatewayWebSocketEvent, context: LambdaContext) -> None: + """Set up the current event and context, shared with included routers via class attributes.""" + self.lambda_context = context + Router.lambda_context = context + + Router.current_event = event if isinstance(event, APIGatewayWebSocketEvent) else APIGatewayWebSocketEvent(event) + self.current_event = Router.current_event diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md new file mode 100644 index 00000000000..7541656e0a0 --- /dev/null +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -0,0 +1,132 @@ +--- +title: API Gateway WebSocket APIs +description: Core utility +status: new +--- + +Event Handler for Amazon API Gateway WebSocket APIs. + +## Key Features + +* Route WebSocket events by route key with dedicated `$connect`, `$disconnect`, and `$default` decorators +* Automatic response normalization to the `statusCode`/`body` shape API Gateway expects +* Accept or reject connections from your `$connect` handler with a status code +* Exception handling with a safety net that prevents stack traces from reaching connected clients +* No dependencies beyond the standard library — routing works with the core Powertools for AWS package + +## Terminology + +**[WebSocket API](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api.html){target="_blank"}**. An API Gateway API type that maintains a persistent, bidirectional connection between clients and your backend. API Gateway manages the connections; your Lambda function receives one event per client message or lifecycle change. + +**Route key**. The value API Gateway uses to select which route (and backend integration) handles a message. Lifecycle events use the predefined `$connect`, `$disconnect`, and `$default` route keys; your API can define custom route keys for application messages. + +**[Route selection expression](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-develop-routes.html){target="_blank"}**. An API-level expression, commonly `$request.body.action`, that extracts the route key from each incoming message. Messages that don't match any custom route fall through to `$default`. + +**Connection ID**. A unique identifier API Gateway assigns to each connection. After `$connect`, it is the only identity a message carries unless a Lambda authorizer is configured — treat it as the session key and look up anything else you need against it. + +**[Route response](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-route-response.html){target="_blank"}**. Opt-in configuration that makes API Gateway deliver your handler's returned body back to the calling client. Without a route response on the route, returned bodies are discarded. + +**[Lambda authorizer](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-lambda-auth.html){target="_blank"}**. An optional authorization step you configure on your API's `$connect` route. API Gateway invokes it before your handler; on success, it stores the authorizer's output with the connection and injects it into every subsequent invocation. + +## Getting started + +### Required resources + +You need an API Gateway WebSocket API with its routes integrated with your Lambda function. The template below wires `$connect`, `$disconnect`, `$default`, and a custom `orderUpdate` route to a single function. + +???+ warning "Returned bodies need a route response" + A handler's returned body reaches the calling client only on routes with a **route response** configured. Without it, API Gateway silently discards the body. The template configures route responses on `orderUpdate` and `$default` so the replies in the examples below actually reach the client. + +=== "getting_started_with_websocket_api.yaml" + + ```yaml hl_lines="27 50-80" + --8<-- "examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml" + ``` + +### Route decorators + +Register handlers with `@app.on_connect()`, `@app.on_disconnect()`, `@app.on_default()`, or `@app.route("yourRouteKey")` for custom route keys. Handlers take no arguments; they access the request through `app.current_event` and return a value that becomes the Lambda response. + +Route keys are matched exactly — API Gateway resolves the route key through the route selection expression before invoking your function, so there is nothing to pattern-match. + +=== "getting_started_with_connect.py" + + ```python hl_lines="6 9 12 16" + --8<-- "examples/event_handler_api_gateway_websocket/src/getting_started_with_connect.py" + ``` + + 1. Headers and cookies are only available on `$connect` — later messages carry only the connection ID. + 2. Returning a non-2xx status code rejects the WebSocket upgrade. + 3. `None` normalizes to `{"statusCode": 200}`, accepting the connection. + +=== "getting_started_with_disconnect.py" + + ```python hl_lines="6 9" + --8<-- "examples/event_handler_api_gateway_websocket/src/getting_started_with_disconnect.py" + ``` + + 1. `$disconnect` is delivered on a best-effort basis, including for connections whose `$connect` was rejected — don't assume your connect handler accepted this connection. + +=== "getting_started_with_custom_route.py" + + ```python hl_lines="4 7 13" + --8<-- "examples/event_handler_api_gateway_websocket/src/getting_started_with_custom_route.py" + ``` + + 1. With the route selection expression `$request.body.action`, this handles messages like `{"action": "orderUpdate", "orderId": 123}`. + 2. Messages whose route key doesn't match any route fall through to `$default`. + +### Response format + +The resolver normalizes your handler's return value into the response shape API Gateway expects. Return a plain value, or a 2-tuple to set the status code: + +| Handler returns | Lambda response | +| --------------- | --------------- | +| `None` | `{"statusCode": 200}` | +| `str` | `{"statusCode": 200, "body": }` | +| `dict` / `list` | `{"statusCode": 200, "body": json.dumps(value)}` | +| 2-tuple `(value, status_code)` | status code from the tuple; body from `value` per the rules above | + +A returned dict is always body content — it is never inspected for a `statusCode` key. + +???+ note "`None` maps to 200, not an empty response" + Unlike the REST resolver, `None` returns `{"statusCode": 200}` because `$connect` requires a status code, and accepting the connection is the only sensible default. + +The status code has different effects depending on the route: + +* **`$connect`**: the status code decides the connection — 2xx accepts the WebSocket upgrade, anything else rejects it. +* **All other routes**: the body is delivered back to the client only when the route has a route response configured, and a non-2xx status code does **not** drop the connection. + +## Advanced + +### Exception handling + +Register handlers for specific exception types with `@app.exception_handler`; it also accepts a list of types, and lookup respects inheritance. The handler receives the exception, and its return value goes through the same [response normalization](#response-format). + +=== "working_with_exception_handling.py" + + ```python hl_lines="9 12 19" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_exception_handling.py" + ``` + + 1. Also accepts a list of exception types, and matches subclasses of the registered type. + 2. The handler's return value goes through the same normalization as route handler returns. + +???+ note "Unhandled exceptions never reach the client" + An exception with no registered handler is logged and becomes a bare `{"statusCode": 500}`. If it propagated to the Lambda runtime instead, the runtime's error response — exception type, message, and stack trace — would be delivered to the connected client. + +## Testing your code + +Test your handlers by passing a WebSocket event payload to `app.resolve()` — no mocks of API Gateway needed. + +=== "getting_started_with_testing.py" + + ```python + --8<-- "examples/event_handler_api_gateway_websocket/src/getting_started_with_testing.py" + ``` + +=== "getting_started_with_testing_event.json" + + ```json + --8<-- "examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json" + ``` diff --git a/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml b/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml new file mode 100644 index 00000000000..ef5d7b22e69 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml @@ -0,0 +1,100 @@ +AWSTemplateFormatVersion: '2010-09-09' +Transform: AWS::Serverless-2016-10-31 + +Globals: + Function: + Timeout: 5 + MemorySize: 256 + Runtime: python3.13 + Tracing: Active + Environment: + Variables: + POWERTOOLS_LOG_LEVEL: INFO + POWERTOOLS_SERVICE_NAME: websocket + +Resources: + WebSocketHandlerFunction: + Type: AWS::Serverless::Function + Properties: + Handler: app.lambda_handler + CodeUri: websocket_handler + + WebSocketApi: + Type: AWS::ApiGatewayV2::Api + Properties: + Name: PowertoolsWebSocketApi + ProtocolType: WEBSOCKET + RouteSelectionExpression: $request.body.action + + WebSocketIntegration: + Type: AWS::ApiGatewayV2::Integration + Properties: + ApiId: !Ref WebSocketApi + IntegrationType: AWS_PROXY + IntegrationUri: !Sub arn:${AWS::Partition}:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${WebSocketHandlerFunction.Arn}/invocations + + ConnectRoute: + Type: AWS::ApiGatewayV2::Route + Properties: + ApiId: !Ref WebSocketApi + RouteKey: $connect + Target: !Sub integrations/${WebSocketIntegration} + + DisconnectRoute: + Type: AWS::ApiGatewayV2::Route + Properties: + ApiId: !Ref WebSocketApi + RouteKey: $disconnect + Target: !Sub integrations/${WebSocketIntegration} + + # RouteResponseSelectionExpression + RouteResponse deliver the handler's + # returned body back to the client; without them the body is discarded. + OrderUpdateRoute: + Type: AWS::ApiGatewayV2::Route + Properties: + ApiId: !Ref WebSocketApi + RouteKey: orderUpdate + RouteResponseSelectionExpression: $default + Target: !Sub integrations/${WebSocketIntegration} + + OrderUpdateRouteResponse: + Type: AWS::ApiGatewayV2::RouteResponse + Properties: + ApiId: !Ref WebSocketApi + RouteId: !Ref OrderUpdateRoute + RouteResponseKey: $default + + DefaultRoute: + Type: AWS::ApiGatewayV2::Route + Properties: + ApiId: !Ref WebSocketApi + RouteKey: $default + RouteResponseSelectionExpression: $default + Target: !Sub integrations/${WebSocketIntegration} + + DefaultRouteResponse: + Type: AWS::ApiGatewayV2::RouteResponse + Properties: + ApiId: !Ref WebSocketApi + RouteId: !Ref DefaultRoute + RouteResponseKey: $default + + WebSocketStage: + Type: AWS::ApiGatewayV2::Stage + Properties: + ApiId: !Ref WebSocketApi + StageName: prod + AutoDeploy: true + + WebSocketInvokePermission: + Type: AWS::Lambda::Permission + Properties: + FunctionName: !Ref WebSocketHandlerFunction + Action: lambda:InvokeFunction + Principal: apigateway.amazonaws.com + SourceArn: !Sub arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/* + +Outputs: + WebSocketURL: + Description: WebSocket endpoint clients connect to + Value: !Sub wss://${WebSocketApi}.execute-api.${AWS::Region}.amazonaws.com/${WebSocketStage} diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_connect.py b/examples/event_handler_api_gateway_websocket/src/getting_started_with_connect.py new file mode 100644 index 00000000000..58b4e3946b6 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_connect.py @@ -0,0 +1,20 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() +app = APIGatewayWebSocketResolver() + + +@app.on_connect() +def connect(): + if "Authorization" not in app.current_event.headers: # (1)! + return None, 401 # (2)! + + connection_id = app.current_event.request_context.connection_id + logger.info("Connection accepted", connection_id=connection_id) + return None # (3)! + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_custom_route.py b/examples/event_handler_api_gateway_websocket/src/getting_started_with_custom_route.py new file mode 100644 index 00000000000..b6800f7ea18 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_custom_route.py @@ -0,0 +1,19 @@ +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() + + +@app.route("orderUpdate") # (1)! +def order_update(): + order = app.current_event.json_body + return {"orderId": order["orderId"], "status": "received"} + + +@app.on_default() # (2)! +def default(): + return {"error": "unknown action"}, 400 + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_disconnect.py b/examples/event_handler_api_gateway_websocket/src/getting_started_with_disconnect.py new file mode 100644 index 00000000000..9e2fe8d98e3 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_disconnect.py @@ -0,0 +1,21 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() +app = APIGatewayWebSocketResolver() + + +@app.on_disconnect() +def disconnect(): # (1)! + request_context = app.current_event.request_context + logger.info( + "Connection closed", + connection_id=request_context.connection_id, + status_code=request_context.disconnect_status_code, + reason=request_context.disconnect_reason, + ) + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing.py b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing.py new file mode 100644 index 00000000000..5d9ff84c4d4 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing.py @@ -0,0 +1,36 @@ +import json +from pathlib import Path + +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + + +class LambdaContext: + def __init__(self): + self.function_name = "test-func" + self.memory_limit_in_mb = 128 + self.invoked_function_arn = "arn:aws:lambda:eu-west-1:809313241234:function:test-func" + self.aws_request_id = "52fdfc07-2182-154f-163f-5f0f9a621d72" + + def get_remaining_time_in_millis(self) -> int: + return 1000 + + +def test_connect_is_accepted(): + # GIVEN a sample $connect event + with Path.open(Path("getting_started_with_testing_event.json"), "r") as f: + event = json.load(f) + + lambda_context = LambdaContext() + + # GIVEN a resolver with a $connect handler + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def connect(): + return None + + # WHEN we resolve the event + result = app.resolve(event, lambda_context) + + # THEN the connection is accepted + assert result == {"statusCode": 200} diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json new file mode 100644 index 00000000000..188d5869326 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json @@ -0,0 +1,49 @@ +{ + "headers": { + "Host": "fjnq7njcv2.execute-api.us-east-1.amazonaws.com", + "Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits", + "Sec-WebSocket-Key": "+W5xw47OHh3OTFsWKjGu9Q==", + "Sec-WebSocket-Version": "13", + "X-Amzn-Trace-Id": "Root=1-6731ebfc-08e1e656421db73c5d2eef31", + "X-Forwarded-For": "166.90.225.1", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Host": ["fjnq7njcv2.execute-api.us-east-1.amazonaws.com"], + "Sec-WebSocket-Extensions": ["permessage-deflate; client_max_window_bits"], + "Sec-WebSocket-Key": ["+W5xw47OHh3OTFsWKjGu9Q=="], + "Sec-WebSocket-Version": ["13"], + "X-Amzn-Trace-Id": ["Root=1-6731ebfc-08e1e656421db73c5d2eef31"], + "X-Forwarded-For": ["166.90.225.1"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "queryStringParameters": { + "userId": "user123", + "token": "abc.def.ghi" + }, + "multiValueQueryStringParameters": { + "userId": ["123"], + "token": ["abc.def.ghi"], + "filter": ["new", "unread"] + }, + "requestContext": { + "routeKey": "$connect", + "eventType": "CONNECT", + "extendedRequestId": "BFHPhFe3IAMF95g=", + "requestTime": "11/Nov/2024:11:35:24 +0000", + "messageDirection": "IN", + "stage": "prod", + "connectedAt": 1731324924553, + "requestTimeEpoch": 1731324924561, + "identity": { + "sourceIp": "166.90.225.1" + }, + "requestId": "BFHPhFe3IAMF95g=", + "domainName": "asasasas.execute-api.us-east-1.amazonaws.com", + "connectionId": "BFHPhfCWIAMCKlQ=", + "apiId": "asasasas" + }, + "isBase64Encoded": false +} \ No newline at end of file diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_exception_handling.py b/examples/event_handler_api_gateway_websocket/src/working_with_exception_handling.py new file mode 100644 index 00000000000..bbb4baaa9a6 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_exception_handling.py @@ -0,0 +1,24 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() +app = APIGatewayWebSocketResolver() + + +@app.exception_handler(ValueError) # (1)! +def handle_invalid_message(exc: ValueError): + logger.error(f"Invalid message: {exc}") + return {"error": str(exc)}, 400 # (2)! + + +@app.route("orderUpdate") +def order_update(): + order = app.current_event.json_body + if "orderId" not in order: + raise ValueError("missing orderId") + return {"orderId": order["orderId"], "status": "received"} + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/mkdocs.yml b/mkdocs.yml index 265560a55d5..00297992dfe 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -24,6 +24,7 @@ nav: - core/event_handler/openapi.md - core/event_handler/appsync.md - core/event_handler/appsync_events.md + - core/event_handler/api_gateway_websocket.md - core/event_handler/bedrock_agents.md - utilities/parameters.md - utilities/batch.md @@ -245,6 +246,7 @@ plugins: - core/event_handler/api_gateway.md - core/event_handler/appsync.md - core/event_handler/appsync_events.md + - core/event_handler/api_gateway_websocket.md - core/event_handler/bedrock_agents.md Utilities: - utilities/parameters.md diff --git a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/__init__.py b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py new file mode 100644 index 00000000000..ad217028e87 --- /dev/null +++ b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py @@ -0,0 +1,455 @@ +import json +from copy import deepcopy + +import pytest + +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.event_handler.api_gateway_websocket import Router +from aws_lambda_powertools.event_handler.api_gateway_websocket._registry import RouteRegistry +from aws_lambda_powertools.utilities.data_classes import APIGatewayWebSocketEvent +from aws_lambda_powertools.warnings import PowertoolsUserWarning +from tests.functional.utils import load_event + + +class LambdaContext: + def __init__(self): + self.function_name = "test-func" + self.memory_limit_in_mb = 128 + self.invoked_function_arn = "arn:aws:lambda:eu-west-1:809313241234:function:test-func" + self.aws_request_id = "52fdfc07-2182-154f-163f-5f0f9a621d72" + + def get_remaining_time_in_millis(self) -> int: + return 1000 + + +@pytest.fixture(scope="module") +def lambda_context() -> LambdaContext: + return LambdaContext() + + +@pytest.fixture(scope="module") +def connect_event(): + return load_event("apiGatewayWebSocketApiConnect.json") + + +@pytest.fixture(scope="module") +def disconnect_event(): + return load_event("apiGatewayWebSocketApiDisconnect.json") + + +@pytest.fixture(scope="module") +def message_event(): + return load_event("apiGatewayWebSocketApiMessage.json") + + +def test_connect_route_dispatch(connect_event, lambda_context): + # GIVEN a resolver with a $connect handler returning None + app = APIGatewayWebSocketResolver() + invocations = [] + + @app.on_connect() + def connect(): + invocations.append(app.current_event.request_context.connection_id) + + # WHEN a $connect event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the handler runs and None normalizes to a bare 200 (accept the connection) + assert invocations == [connect_event["requestContext"]["connectionId"]] + assert result == {"statusCode": 200} + + +def test_disconnect_route_dispatch(disconnect_event, lambda_context): + # GIVEN a resolver with a $disconnect handler reading the disconnect details + app = APIGatewayWebSocketResolver() + captured = {} + + @app.on_disconnect() + def disconnect(): + captured["status_code"] = app.current_event.request_context.disconnect_status_code + captured["reason"] = app.current_event.request_context.disconnect_reason + + # WHEN a $disconnect event is resolved + result = app.resolve(disconnect_event, lambda_context) + + # THEN the handler sees the close details from the event + assert result == {"statusCode": 200} + assert captured["status_code"] == disconnect_event["requestContext"]["disconnectStatusCode"] + assert captured["reason"] == disconnect_event["requestContext"]["disconnectReason"] + + +def test_default_route_reads_json_body(message_event, lambda_context): + # GIVEN a $default handler and a message event with a JSON body + event = deepcopy(message_event) + event["requestContext"]["routeKey"] = "$default" + app = APIGatewayWebSocketResolver() + + @app.on_default() + def default(): + return {"received": app.current_event.json_body["message"]} + + # WHEN the event is resolved + result = app.resolve(event, lambda_context) + + # THEN the handler reads the parsed body and the dict is JSON-serialized + assert result == {"statusCode": 200, "body": json.dumps({"received": "Hello from client"}, separators=(",", ":"))} + + +def test_custom_route_key_dispatch(message_event, lambda_context): + # GIVEN a handler registered for a custom route key + event = deepcopy(message_event) + event["requestContext"]["routeKey"] = "orderUpdate" + app = APIGatewayWebSocketResolver() + + @app.route("orderUpdate") + def order_update(): + return "order received" + + # WHEN an event with that route key is resolved + result = app.resolve(event, lambda_context) + + # THEN the custom route handler is dispatched and str returns pass through as body + assert result == {"statusCode": 200, "body": "order received"} + + +@pytest.mark.parametrize( + ("handler_return", "expected"), + [ + (None, {"statusCode": 200}), + ("plain text", {"statusCode": 200, "body": "plain text"}), + ({"statusCode": 999}, {"statusCode": 200, "body": '{"statusCode":999}'}), + (["a", "b"], {"statusCode": 200, "body": '["a","b"]'}), + ((None, 401), {"statusCode": 401}), + (({"reason": "full"}, 503), {"statusCode": 503, "body": '{"reason":"full"}'}), + (("denied", 403), {"statusCode": 403, "body": "denied"}), + ], + ids=["none", "str", "dict-with-statusCode-key-is-body", "list", "tuple-none", "tuple-dict", "tuple-str"], +) +def test_return_value_normalization(connect_event, lambda_context, handler_return, expected): + # GIVEN a handler returning each supported shape + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def connect(): + return handler_return + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the response matches the normalization table exactly + assert result == expected + + +def test_unknown_route_key_warns_and_returns_400(message_event, lambda_context): + # GIVEN a resolver with no handler for the event's route key + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def connect(): + return None + + # WHEN the event is resolved + with pytest.warns(PowertoolsUserWarning, match="No route handler registered for route key `chat`"): + result = app.resolve(message_event, lambda_context) + + # THEN a bare 400 is returned + assert result == {"statusCode": 400} + + +def test_duplicate_route_registration_warns_last_wins(connect_event, lambda_context): + # GIVEN two handlers registered for the same route key + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def first(): + return "first" + + with pytest.warns(PowertoolsUserWarning, match="already registered for route key `\\$connect`"): + + @app.on_connect() + def second(): + return "second" + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the last registration wins + assert result == {"statusCode": 200, "body": "second"} + + +def test_empty_route_key_registration_warns_and_skips(connect_event, lambda_context): + # GIVEN a handler registered with an empty route key + app = APIGatewayWebSocketResolver() + + with pytest.warns(PowertoolsUserWarning, match="route key registered for `handler` is empty"): + + @app.route("") + def handler(): + return None + + # THEN the function is returned undecorated and nothing is registered + assert handler() is None + assert app._route_registry.routes == {} + + +def test_current_event_and_lambda_context(connect_event, lambda_context): + # GIVEN a resolver and a handler inspecting the event and context + app = APIGatewayWebSocketResolver() + captured = {} + + @app.on_connect() + def connect(): + captured["event"] = app.current_event + captured["context"] = app.lambda_context + + # WHEN resolving a raw dict event + app.resolve(connect_event, lambda_context) + + # THEN the event is wrapped in the data class and the context is set + assert isinstance(captured["event"], APIGatewayWebSocketEvent) + assert captured["event"].request_context.route_key == "$connect" + assert captured["context"] is lambda_context + + +def test_already_wrapped_event_is_not_rewrapped(connect_event, lambda_context): + # GIVEN an event already wrapped in the data class + wrapped = APIGatewayWebSocketEvent(connect_event) + app = APIGatewayWebSocketResolver() + captured = {} + + @app.on_connect() + def connect(): + captured["event"] = app.current_event + + # WHEN it is resolved + result = app.resolve(wrapped, lambda_context) + + # THEN the same instance is used and dispatch works + assert captured["event"] is wrapped + assert result == {"statusCode": 200} + + +def test_unhandled_exception_returns_bare_500(connect_event, lambda_context): + # GIVEN a handler raising an exception with no registered handler + app = APIGatewayWebSocketResolver() + secret_message = "sensitive detail about the failure" + + @app.on_connect() + def connect(): + raise RuntimeError(secret_message) + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN a bare 500 is returned with no exception details anywhere in the response + assert result == {"statusCode": 500} + serialized = json.dumps(result) + assert secret_message not in serialized + assert "RuntimeError" not in serialized + assert "Traceback" not in serialized + + +def test_registered_exception_handler_response_is_normalized(connect_event, lambda_context): + # GIVEN an exception handler registered for ValueError + app = APIGatewayWebSocketResolver() + + @app.exception_handler(ValueError) + def handle_value_error(exc: ValueError): + return {"error": str(exc)}, 400 + + @app.on_connect() + def connect(): + raise ValueError("bad input") + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the handler's return value goes through the same normalization + assert result == {"statusCode": 400, "body": '{"error":"bad input"}'} + + +def test_exception_handler_lookup_respects_inheritance(connect_event, lambda_context): + # GIVEN a handler registered for a base exception class + class OrderError(Exception): ... + + class OrderNotFoundError(OrderError): ... + + app = APIGatewayWebSocketResolver() + + @app.exception_handler(OrderError) + def handle_order_error(exc: OrderError): + return None, 404 + + @app.on_connect() + def connect(): + raise OrderNotFoundError("order 123 not found") + + # WHEN a subclass of the registered exception is raised + result = app.resolve(connect_event, lambda_context) + + # THEN the base class handler is used + assert result == {"statusCode": 404} + + +def test_disconnect_dispatches_even_when_connect_was_rejected(connect_event, disconnect_event, lambda_context): + # GIVEN a resolver whose $connect handler rejects the connection + app = APIGatewayWebSocketResolver() + dispatched = [] + + @app.on_connect() + def connect(): + dispatched.append("$connect") + return None, 401 + + @app.on_disconnect() + def disconnect(): + dispatched.append("$disconnect") + + # WHEN $connect is rejected and API Gateway later delivers $disconnect anyway + connect_result = app.resolve(connect_event, lambda_context) + disconnect_result = app.resolve(disconnect_event, lambda_context) + + # THEN the resolver keeps no per-connection state and dispatches both + assert connect_result == {"statusCode": 401} + assert disconnect_result == {"statusCode": 200} + assert dispatched == ["$connect", "$disconnect"] + + +def test_resolver_is_callable_as_lambda_handler(connect_event, lambda_context): + # GIVEN a resolver with a $connect handler + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def connect(): + return None + + # WHEN the resolver instance itself is used as the Lambda handler + result = app(connect_event, lambda_context) + + # THEN it resolves the event + assert result == {"statusCode": 200} + + +def test_append_context_visible_in_handler_and_cleared_after_resolve(connect_event, lambda_context): + # GIVEN context appended before resolution + app = APIGatewayWebSocketResolver() + app.append_context(tenant_id="tenant-a") + captured = {} + + @app.on_connect() + def connect(): + captured["tenant_id"] = app.context.get("tenant_id") + + # WHEN the event is resolved + app.resolve(connect_event, lambda_context) + + # THEN the handler sees the context and it is cleared afterwards + assert captured["tenant_id"] == "tenant-a" + assert app.context == {} + + +def test_malformed_event_returns_bare_500(lambda_context): + # GIVEN an event missing requestContext entirely + app = APIGatewayWebSocketResolver() + + @app.on_connect() + def connect(): + return None + + # WHEN the event is resolved + result = app.resolve({}, lambda_context) + + # THEN the resolver returns a bare 500 with no exception details + assert result == {"statusCode": 500} + + +def test_exception_handler_raising_returns_bare_500(connect_event, lambda_context): + # GIVEN an exception handler that itself raises + app = APIGatewayWebSocketResolver() + secret_message = "sensitive detail about the failure" + + @app.exception_handler(ValueError) + def handle_value_error(exc: ValueError): + raise RuntimeError(secret_message) + + @app.on_connect() + def connect(): + raise ValueError("bad input") + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN a bare 500 is returned with no exception details + assert result == {"statusCode": 500} + assert secret_message not in json.dumps(result) + + +def test_router_exception_handler_registration(): + # GIVEN a Router with exception handlers registered for a single and a list of types + router = Router() + + @router.exception_handler(ValueError) + def handle_value_error(exc: ValueError): ... + + @router.exception_handler([KeyError, TypeError]) + def handle_lookup_errors(exc: Exception): ... + + # THEN the handlers are stored for merging via include_router + assert router._exception_handlers[ValueError] is handle_value_error + assert router._exception_handlers[KeyError] is handle_lookup_errors + assert router._exception_handlers[TypeError] is handle_lookup_errors + + +def test_route_registry_merge(): + # GIVEN two registries with registered routes + first = RouteRegistry() + second = RouteRegistry() + + @first.register("$connect") + def connect(): ... + + @second.register("orderUpdate") + def order_update(): ... + + # WHEN one registry is merged into the other + first.merge(second) + + # THEN both routes are found in the merged registry + assert first.find_route("$connect")["func"] is connect + assert first.find_route("orderUpdate")["func"] is order_update + + +def test_duplicate_disconnect_delivery_dispatches_both(disconnect_event, lambda_context): + # GIVEN a resolver with a $disconnect handler ($disconnect delivery is at-least-once in production) + app = APIGatewayWebSocketResolver() + dispatched = [] + + @app.on_disconnect() + def disconnect(): + dispatched.append(app.current_event.request_context.connection_id) + + # WHEN the same $disconnect event is delivered twice + first = app.resolve(disconnect_event, lambda_context) + second = app.resolve(disconnect_event, lambda_context) + + # THEN both deliveries dispatch to the handler and succeed + assert first == {"statusCode": 200} + assert second == {"statusCode": 200} + assert dispatched == [disconnect_event["requestContext"]["connectionId"]] * 2 + + +def test_system_route_key_registrable_via_route_primitive(disconnect_event, lambda_context): + # GIVEN a $disconnect handler registered via the route() primitive instead of on_disconnect() + app = APIGatewayWebSocketResolver() + dispatched = [] + + @app.route("$disconnect") + def disconnect(): + dispatched.append(app.current_event.request_context.route_key) + + # WHEN a $disconnect event is resolved + result = app.resolve(disconnect_event, lambda_context) + + # THEN it dispatches identically to the sugar decorator + assert result == {"statusCode": 200} + assert dispatched == ["$disconnect"] From 144b62e413ecd277caf6e672634202a899d65f01 Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 18:28:49 +1000 Subject: [PATCH 3/7] feat(event_handler): add middleware support to API Gateway WebSocket resolver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reuses the shared middleware framework (MiddlewareFrame, BaseMiddlewareHandler) with a WebSocket terminal adapter: global middlewares registered via use() run first, then route-level middlewares=[...], then the handler. Return values — including short-circuits that skip next_middleware — go through the same response normalization as handler returns, and middleware exceptions follow the existing exception-handling path. Global middlewares also run around the 400 for unmatched route keys, matching the REST resolver's not-found behaviour. Includes middleware tests (ordering, short-circuit, isolation across routes, exception paths, class-based middleware) and docs: Middleware plus Authentication patterns (Lambda authorizer first, then authenticate/inject_user middleware with a bring-your-own connection store). --- .../api_gateway_websocket/base.py | 4 + .../api_gateway_websocket/router.py | 28 +++ .../api_gateway_websocket/websocket.py | 40 +++- .../event_handler/api_gateway_websocket.md | 50 ++++ .../src/working_with_lambda_authorizer.py | 20 ++ .../src/working_with_middleware.py | 31 +++ .../working_with_middleware_authentication.py | 46 ++++ .../test_api_gateway_websocket_resolver.py | 223 ++++++++++++++++++ 8 files changed, 439 insertions(+), 3 deletions(-) create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_middleware.py create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py index d8b45cf2b47..35596bb5b42 100644 --- a/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/base.py @@ -43,6 +43,10 @@ def on_default( ) -> Callable: raise NotImplementedError + @abstractmethod + def use(self, middlewares: list[Callable]) -> None: + raise NotImplementedError + def append_context(self, **additional_context) -> None: """ Appends context information available under any route. diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py index ee5306e5f10..31854a520c7 100644 --- a/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py @@ -55,6 +55,7 @@ def __init__(self): self.context = {} # early init as customers might add context before event resolution self._route_registry = RouteRegistry() self._exception_handlers: dict[type[Exception], Callable] = {} + self._router_middlewares: list[Callable] = [] def route( self, @@ -189,6 +190,33 @@ def on_default( """ return self.route(route_key=ROUTE_KEY_DEFAULT, middlewares=middlewares) + def use(self, middlewares: list[Callable]) -> None: + """ + Add one or more global middlewares that run before route-specific middlewares. + + Middlewares run in insertion order: global middlewares first, then route-level + `middlewares=[...]`, then the route handler. + + Parameters + ---------- + middlewares : list[Callable] + List of middlewares to run on every event, including unmatched route keys + + Examples + -------- + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> from aws_lambda_powertools.event_handler.middlewares import NextMiddleware + >>> + >>> app = APIGatewayWebSocketResolver() + >>> + >>> def log_route(app: APIGatewayWebSocketResolver, next_middleware: NextMiddleware): + >>> print(f"dispatching {app.current_event.request_context.route_key}") + >>> return next_middleware(app) + >>> + >>> app.use(middlewares=[log_route]) + """ + self._router_middlewares.extend(middlewares) + def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) -> Callable: """ Register a handler for one or more exception types. diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py index 0496caeed12..b4428894dc7 100644 --- a/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py @@ -6,7 +6,9 @@ from http import HTTPStatus from typing import TYPE_CHECKING, Any +from aws_lambda_powertools.event_handler.api_gateway import MiddlewareFrame from aws_lambda_powertools.event_handler.api_gateway_websocket.router import Router +from aws_lambda_powertools.event_handler.api_gateway_websocket.types import WebSocketRoute from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager from aws_lambda_powertools.shared.json_encoder import Encoder from aws_lambda_powertools.utilities.data_classes.api_gateway_websocket_event import APIGatewayWebSocketEvent @@ -20,6 +22,20 @@ logger = logging.getLogger(__name__) +def _registered_route_adapter(app: APIGatewayWebSocketResolver, next_middleware: Callable[..., Any]) -> Any: + """Terminal middleware frame: calls the registered route handler (no arguments) and returns its raw value. + + Response normalization happens once, after the whole middleware chain returns, so middlewares + may short-circuit with any of the same shapes route handlers can return. + """ + return next_middleware() + + +def _route_not_found() -> tuple[None, int]: + """Stand-in handler for unmatched route keys, so global middlewares still run around the 400.""" + return None, HTTPStatus.BAD_REQUEST.value + + class APIGatewayWebSocketResolver(Router): """ API Gateway WebSocket API resolver. @@ -70,6 +86,7 @@ class APIGatewayWebSocketResolver(Router): def __init__(self): super().__init__() self.exception_handler_manager = ExceptionHandlerManager() + self.processed_stack_frames: list[str] = [] # used by MiddlewareFrame for debug traces def __call__( self, @@ -110,6 +127,7 @@ def resolve( >>> return app.resolve(event, context) """ try: + self.processed_stack_frames.clear() self._setup_context(event, context) return self._resolve_route() except Exception as exc: @@ -149,7 +167,7 @@ def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) return self.exception_handler_manager.exception_handler(exc_class=exc_class) def _resolve_route(self) -> dict[str, Any]: - """Dispatch the current event to its route handler and normalize the response.""" + """Dispatch the current event through the middleware chain to its route handler and normalize the response.""" route_key = self.current_event.request_context.route_key route = self._route_registry.find_route(route_key) if route is None: @@ -158,10 +176,26 @@ def _resolve_route(self) -> dict[str, Any]: stacklevel=2, category=PowertoolsUserWarning, ) - return {"statusCode": HTTPStatus.BAD_REQUEST.value} + # global middlewares still run around the 400, as REST does for not-found routes + route = WebSocketRoute(func=_route_not_found, middlewares=[]) logger.debug(f"Dispatching route key `{route_key}` to `{route['func'].__name__}`") - return self._to_response(route["func"]()) + middleware_stack = self._build_middleware_stack(route) + return self._to_response(middleware_stack(self)) + + def _build_middleware_stack(self, route: WebSocketRoute) -> Callable: + """Wrap the route handler in middleware frames: global (`use`) first, then route-level, then + the terminal adapter that calls the handler. Frames are wrapped in reverse so middlewares + run in the order they were added.""" + stack: Callable = route["func"] + all_middlewares = [*self._router_middlewares, *route["middlewares"], _registered_route_adapter] + for middleware in reversed(all_middlewares): + stack = MiddlewareFrame(current_middleware=middleware, next_middleware=stack) + return stack + + def _push_processed_stack_frame(self, frame: str) -> None: + """Record a processed middleware frame; MiddlewareFrame calls this for debug traces.""" + self.processed_stack_frames.append(frame) def _handle_exception(self, exc: Exception) -> dict[str, Any]: """Resolve an exception to a response via registered handlers, or a bare 500. diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md index 7541656e0a0..a4d8c02fddf 100644 --- a/docs/core/event_handler/api_gateway_websocket.md +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -99,6 +99,56 @@ The status code has different effects depending on the route: ## Advanced +### Middleware + +The resolver reuses the same middleware framework as the REST resolver. A middleware is a callable receiving the resolver instance and the next handler in the chain; it runs code before and/or after the route handler, and its return value goes through the same [response normalization](#response-format) as handler returns. + +Execution order is: global middlewares (`app.use`), then route-level `middlewares=[...]`, then the route handler. Return without calling `next_middleware(app)` to short-circuit the chain; exceptions raised in middlewares follow the same [exception handling](#exception-handling) flow as handlers. + +=== "working_with_middleware.py" + + ```python hl_lines="9 11 17 21 24" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_middleware.py" + ``` + + 1. Middlewares receive the resolver instance and the next handler in the chain. + 2. Call `next_middleware(app)` to continue the chain; its return value is the route response. + 3. Short-circuit: the handler never runs and this value becomes the response. + 4. Global middlewares run on every event — including unmatched route keys — before route-level ones. + 5. Route-level middlewares run for this route only. + +### Authentication patterns + +`$connect` is the only route where credentials exist: headers and cookies are present on the WebSocket handshake only, and later messages carry nothing but the connection ID. There are two ways to bridge that gap. + +#### Lambda authorizer + +Prefer a Lambda authorizer on `$connect` when you can run a separate authorizer function. API Gateway persists the authorizer's output with the connection and injects it into `request_context.authorizer` on **every** invocation for that connection — a managed connection-to-identity store, no code needed in your handlers beyond reading it. + +=== "working_with_lambda_authorizer.py" + + ```python hl_lines="9 14-15" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py" + ``` + + 1. Injected by API Gateway on every route for this connection — message routes, `$default`, and `$disconnect` included. + 2. Authorizer context values arrive stringified (a numeric `1` arrives as `"1"`). + +#### Authenticating with middleware + +Without a Lambda authorizer, authenticate on `$connect` with a middleware and persist identity against the connection ID in your own store; a second middleware resolves it on message routes and shares it via `app.append_context`. + +=== "working_with_middleware_authentication.py" + + ```python hl_lines="6 13-14 16 25 29 39" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py" + ``` + + 1. Bring your own connection store. The in-memory dictionary keeps this example short — real invocations for one connection can hit different Lambda environments, so use an external store such as DynamoDB with a TTL. + 2. Headers only exist on `$connect`. + 3. Returning without calling `next_middleware` short-circuits the chain — the handler never runs and the connection is rejected. + 4. Handlers and later middlewares read it from `app.context`. + ### Exception handling Register handlers for specific exception types with `@app.exception_handler`; it also accepts a list of types, and lookup respects inheritance. The handler receives the exception, and its return value goes through the same [response normalization](#response-format). diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py b/examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py new file mode 100644 index 00000000000..a1c9177810d --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_lambda_authorizer.py @@ -0,0 +1,20 @@ +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() + + +@app.route("orderUpdate") +def order_update(): + identity = app.current_event.request_context.authorizer # (1)! + order = app.current_event.json_body + return { + "orderId": order["orderId"], + "status": "received", + "processedFor": identity["principalId"], + "tenant": identity.get("tenantId"), # (2)! + } + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_middleware.py b/examples/event_handler_api_gateway_websocket/src/working_with_middleware.py new file mode 100644 index 00000000000..22758526055 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_middleware.py @@ -0,0 +1,31 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() +app = APIGatewayWebSocketResolver() + + +def log_route(app, next_middleware): # (1)! + logger.info("Dispatching", route_key=app.current_event.request_context.route_key) + return next_middleware(app) # (2)! + + +def enforce_order_limit(app, next_middleware): + order = app.current_event.json_body + if order.get("quantity", 0) > 100: + return {"error": "quantity over limit"}, 400 # (3)! + return next_middleware(app) + + +app.use(middlewares=[log_route]) # (4)! + + +@app.route("orderUpdate", middlewares=[enforce_order_limit]) # (5)! +def order_update(): + order = app.current_event.json_body + return {"orderId": order["orderId"], "status": "received"} + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py b/examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py new file mode 100644 index 00000000000..9c3e187d0bb --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_middleware_authentication.py @@ -0,0 +1,46 @@ +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() + +sessions: dict[str, str] = {} # (1)! + + +def validate_token(token: str | None) -> str | None: + return "alice" if token == "Bearer valid-token" else None + + +def authenticate(app, next_middleware): + user = validate_token(app.current_event.headers.get("Authorization")) # (2)! + if user is None: + return None, 401 # (3)! + sessions[app.current_event.request_context.connection_id] = user + return next_middleware(app) + + +def inject_user(app, next_middleware): + user = sessions.get(app.current_event.request_context.connection_id) + if user is None: + return {"error": "unauthenticated"}, 401 + app.append_context(user=user) # (4)! + return next_middleware(app) + + +@app.on_connect(middlewares=[authenticate]) +def connect(): + return None + + +@app.on_disconnect() +def disconnect(): + sessions.pop(app.current_event.request_context.connection_id, None) + + +@app.route("orderUpdate", middlewares=[inject_user]) +def order_update(): + order = app.current_event.json_body + return {"orderId": order["orderId"], "status": "received", "processedFor": app.context["user"]} + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py index ad217028e87..1097826a50c 100644 --- a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py +++ b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py @@ -453,3 +453,226 @@ def disconnect(): # THEN it dispatches identically to the sugar decorator assert result == {"statusCode": 200} assert dispatched == ["$disconnect"] + + +def test_route_middleware_runs_before_and_after_handler(connect_event, lambda_context): + # GIVEN a route-level middleware capturing execution order around the handler + app = APIGatewayWebSocketResolver() + order = [] + + def middleware(app_, next_middleware): + order.append("before") + result = next_middleware(app_) + order.append("after") + return result + + @app.on_connect(middlewares=[middleware]) + def connect(): + order.append("handler") + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the middleware wraps the handler and the response is normalized + assert order == ["before", "handler", "after"] + assert result == {"statusCode": 200} + + +def test_global_middleware_runs_before_route_middleware(connect_event, lambda_context): + # GIVEN a global middleware (use) and a route-level middleware + app = APIGatewayWebSocketResolver() + order = [] + + def global_middleware(app_, next_middleware): + order.append("global_before") + result = next_middleware(app_) + order.append("global_after") + return result + + def route_middleware(app_, next_middleware): + order.append("route_before") + result = next_middleware(app_) + order.append("route_after") + return result + + app.use(middlewares=[global_middleware]) + + @app.on_connect(middlewares=[route_middleware]) + def connect(): + order.append("handler") + + # WHEN the event is resolved + app.resolve(connect_event, lambda_context) + + # THEN the full order is global -> route -> handler -> route -> global + assert order == ["global_before", "route_before", "handler", "route_after", "global_after"] + + +def test_middleware_short_circuit_skips_handler(connect_event, lambda_context): + # GIVEN a $connect middleware rejecting without calling next_middleware + app = APIGatewayWebSocketResolver() + handler_invoked = [] + + def authenticate(app_, next_middleware): + if "Authorization" not in app_.current_event.headers: + return None, 401 + return next_middleware(app_) + + @app.on_connect(middlewares=[authenticate]) + def connect(): + handler_invoked.append(True) + + # WHEN an event without an Authorization header is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the middleware's return value is normalized and the handler never runs + assert result == {"statusCode": 401} + assert handler_invoked == [] + + +def test_middleware_append_context_visible_in_handler(connect_event, lambda_context): + # GIVEN a middleware sharing data through append_context + app = APIGatewayWebSocketResolver() + captured = {} + + def inject_user(app_, next_middleware): + app_.append_context(user="alice") + return next_middleware(app_) + + @app.on_connect(middlewares=[inject_user]) + def connect(): + captured["user"] = app.context.get("user") + + # WHEN the event is resolved + app.resolve(connect_event, lambda_context) + + # THEN the handler sees the context value and context is cleared after resolve + assert captured["user"] == "alice" + assert app.context == {} + + +def test_middleware_exception_uses_exception_handling_path(connect_event, lambda_context): + # GIVEN a middleware that raises and a registered exception handler + app = APIGatewayWebSocketResolver() + + @app.exception_handler(PermissionError) + def handle_permission_error(exc: PermissionError): + return {"error": str(exc)}, 403 + + def broken_middleware(app_, next_middleware): + raise PermissionError("not allowed") + + @app.on_connect(middlewares=[broken_middleware]) + def connect(): + return None + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the middleware exception goes through the same exception-handling path + assert result == {"statusCode": 403, "body": '{"error":"not allowed"}'} + + +def test_middleware_unhandled_exception_returns_bare_500(connect_event, lambda_context): + # GIVEN a middleware raising an exception with no registered handler + app = APIGatewayWebSocketResolver() + + def broken_middleware(app_, next_middleware): + raise RuntimeError("middleware secret detail") + + @app.on_connect(middlewares=[broken_middleware]) + def connect(): + return None + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN a bare 500 is returned with no exception details + assert result == {"statusCode": 500} + assert "middleware secret detail" not in json.dumps(result) + + +def test_class_based_middleware_from_shared_framework(connect_event, lambda_context): + # GIVEN a middleware built on the shared BaseMiddlewareHandler framework + from aws_lambda_powertools.event_handler.middlewares import BaseMiddlewareHandler + + order = [] + + class TrackingMiddleware(BaseMiddlewareHandler): + def handler(self, app_, next_middleware): + order.append("class_before") + result = next_middleware(app_) + order.append("class_after") + return result + + app = APIGatewayWebSocketResolver() + + @app.on_connect(middlewares=[TrackingMiddleware()]) + def connect(): + order.append("handler") + + # WHEN the event is resolved + result = app.resolve(connect_event, lambda_context) + + # THEN the class-based middleware runs around the handler + assert order == ["class_before", "handler", "class_after"] + assert result == {"statusCode": 200} + + +def test_global_middleware_runs_for_unknown_route_key(message_event, lambda_context): + # GIVEN a global middleware and no handler registered for the event's route key + app = APIGatewayWebSocketResolver() + observed = [] + + def observe(app_, next_middleware): + observed.append(app_.current_event.request_context.route_key) + return next_middleware(app_) + + app.use(middlewares=[observe]) + + @app.on_connect() + def connect(): ... + + # WHEN the event is resolved + with pytest.warns(PowertoolsUserWarning, match="No route handler registered"): + result = app.resolve(message_event, lambda_context) + + # THEN the global middleware observes the miss and the 400 is still returned + assert result == {"statusCode": 400} + assert observed == ["chat"] + + +def test_route_middleware_isolation_across_routes(connect_event, message_event, lambda_context): + # GIVEN a global middleware, route A with a route-level middleware, and route B without + event_a = deepcopy(message_event) + event_a["requestContext"]["routeKey"] = "routeA" + event_b = deepcopy(message_event) + event_b["requestContext"]["routeKey"] = "routeB" + + app = APIGatewayWebSocketResolver() + order = [] + + def global_middleware(app_, next_middleware): + order.append(f"global:{app_.current_event.request_context.route_key}") + return next_middleware(app_) + + def route_a_middleware(app_, next_middleware): + order.append(f"route_a_mw:{app_.current_event.request_context.route_key}") + return next_middleware(app_) + + app.use(middlewares=[global_middleware]) + + @app.route("routeA", middlewares=[route_a_middleware]) + def route_a(): + order.append("handler_a") + + @app.route("routeB") + def route_b(): + order.append("handler_b") + + # WHEN both routes are resolved + app.resolve(event_a, lambda_context) + app.resolve(event_b, lambda_context) + + # THEN the global middleware runs on both routes; the route-level middleware runs only on route A + assert order == ["global:routeA", "route_a_mw:routeA", "handler_a", "global:routeB", "handler_b"] From 7a58a7b52940b40768a5e540acfdf36051140751 Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 18:56:08 +1000 Subject: [PATCH 4/7] feat(event_handler): add Router support to API Gateway WebSocket resolver include_router merges a Router's routes, global middlewares, exception handlers, and context into the resolver, so route keys can be split across files. Router handlers access the request through the router instance (router.current_event), mirroring the AppSync Events and REST resolver conventions. On route key collisions the router-registered handler wins, matching AppSync Events. Includes tests for dispatch, context pointer sharing, middleware and exception-handler merging and ordering, route-key precedence, plus docs with a two-file example. --- .../api_gateway_websocket/websocket.py | 37 +++++ .../event_handler/api_gateway_websocket.md | 20 +++ .../src/working_with_router.py | 16 ++ .../src/working_with_router_orders.py | 15 ++ .../test_api_gateway_websocket_resolver.py | 149 ++++++++++++++++++ 5 files changed, 237 insertions(+) create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_router.py create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_router_orders.py diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py index b4428894dc7..299c14dc590 100644 --- a/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/websocket.py @@ -166,6 +166,43 @@ def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]) """ return self.exception_handler_manager.exception_handler(exc_class=exc_class) + def include_router(self, router: Router) -> None: + """ + Add all routes, middlewares, exception handlers, and context defined in a router. + + Parameters + ---------- + router : Router + A router containing routes to include, typically defined in a separate module + + Examples + -------- + >>> # orders.py + >>> from aws_lambda_powertools.event_handler.api_gateway_websocket import Router + >>> + >>> router = Router() + >>> + >>> @router.route("orderUpdate") + >>> def order_update(): + >>> return {"orderId": router.current_event.json_body["orderId"]} + >>> + >>> # app.py + >>> import orders + >>> from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver + >>> + >>> app = APIGatewayWebSocketResolver() + >>> app.include_router(orders.router) + """ + logger.debug("Merging router and app 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 + + logger.debug("Merging router routes, middlewares, and exception handlers") + self._router_middlewares.extend(router._router_middlewares) + self._route_registry.merge(router._route_registry) + self.exception_handler_manager.update_exception_handlers(router._exception_handlers) + def _resolve_route(self) -> dict[str, Any]: """Dispatch the current event through the middleware chain to its route handler and normalize the response.""" route_key = self.current_event.request_context.route_key diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md index a4d8c02fddf..d1a4b31e387 100644 --- a/docs/core/event_handler/api_gateway_websocket.md +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -149,6 +149,26 @@ Without a Lambda authorizer, authenticate on `$connect` with a middleware and pe 3. Returning without calling `next_middleware` short-circuits the chain — the handler never runs and the connection is rejected. 4. Handlers and later middlewares read it from `app.context`. +### Split routes with Router + +As your API grows, group related route keys in separate files with `Router`, then include them in the resolver. `include_router` merges the router's routes, global middlewares (`use`), exception handlers, and context into the app. Inside a router file, access the request through the router instance — `router.current_event`, `router.lambda_context`, and `router.context`. + +=== "working_with_router_orders.py" + + ```python hl_lines="3 8" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_router_orders.py" + ``` + + 1. Use the router instance to access the current event inside router files. + +=== "working_with_router.py" + + ```python hl_lines="1 7" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_router.py" + ``` + + 1. Registers every route, middleware, and exception handler defined on the router. + ### Exception handling Register handlers for specific exception types with `@app.exception_handler`; it also accepts a list of types, and lookup respects inheritance. The handler receives the exception, and its return value goes through the same [response normalization](#response-format). diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_router.py b/examples/event_handler_api_gateway_websocket/src/working_with_router.py new file mode 100644 index 00000000000..73a37449cc6 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_router.py @@ -0,0 +1,16 @@ +import working_with_router_orders + +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() +app.include_router(working_with_router_orders.router) # (1)! + + +@app.on_connect() +def connect(): + return None + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_router_orders.py b/examples/event_handler_api_gateway_websocket/src/working_with_router_orders.py new file mode 100644 index 00000000000..b0951005fbd --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_router_orders.py @@ -0,0 +1,15 @@ +from aws_lambda_powertools.event_handler.api_gateway_websocket import Router + +router = Router() + + +@router.route("orderUpdate") +def order_update(): + order = router.current_event.json_body # (1)! + return {"orderId": order["orderId"], "status": "received"} + + +@router.route("orderCancel") +def order_cancel(): + order = router.current_event.json_body + return {"orderId": order["orderId"], "status": "cancelled"} diff --git a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py index 1097826a50c..69d3447776e 100644 --- a/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py +++ b/tests/functional/event_handler/required_dependencies/api_gateway_websocket/test_api_gateway_websocket_resolver.py @@ -676,3 +676,152 @@ def route_b(): # THEN the global middleware runs on both routes; the route-level middleware runs only on route A assert order == ["global:routeA", "route_a_mw:routeA", "handler_a", "global:routeB", "handler_b"] + + +def test_included_router_routes_are_dispatched(message_event, lambda_context): + # GIVEN routes registered on a standalone Router + event = deepcopy(message_event) + event["requestContext"]["routeKey"] = "orderUpdate" + + router = Router() + + @router.route("orderUpdate") + def order_update(): + return "router handled it" + + # WHEN the router is included and the event resolved + app = APIGatewayWebSocketResolver() + app.include_router(router) + result = app.resolve(event, lambda_context) + + # THEN the router-registered handler is dispatched + assert result == {"statusCode": 200, "body": "router handled it"} + + +def test_router_current_event_and_lambda_context_in_handler(message_event, lambda_context): + # GIVEN a router handler accessing the event and context through the router instance + event = deepcopy(message_event) + event["requestContext"]["routeKey"] = "orderUpdate" + + router = Router() + captured = {} + + @router.route("orderUpdate") + def order_update(): + captured["event"] = router.current_event + captured["context"] = router.lambda_context + + app = APIGatewayWebSocketResolver() + app.include_router(router) + + # WHEN the event is resolved through the app + app.resolve(event, lambda_context) + + # THEN the router sees the same event and context as the app (class-attribute mechanism) + assert isinstance(captured["event"], APIGatewayWebSocketEvent) + assert captured["event"].request_context.route_key == "orderUpdate" + assert captured["context"] is lambda_context + + +def test_include_router_merges_and_shares_context(connect_event, lambda_context): + # GIVEN context appended to a router before inclusion + router = Router() + router.append_context(source="router") + captured = {} + + @router.route("$connect") + def connect(): + captured["source"] = router.context.get("source") + captured["app_added"] = router.context.get("app_added") + + app = APIGatewayWebSocketResolver() + app.include_router(router) + + # THEN the router context is merged into the app and the pointer is shared + assert app.context["source"] == "router" + assert router.context is app.context + + # WHEN the app appends more context and the event is resolved + app.append_context(app_added=True) + app.resolve(connect_event, lambda_context) + + # THEN the handler saw both values and clear_context cleared the shared dict + assert captured == {"source": "router", "app_added": True} + assert app.context == {} + assert router.context == {} + + +def test_router_level_use_middlewares_are_merged(connect_event, lambda_context): + # GIVEN global middlewares on both the app and a router, and a route-level middleware on the router route + app = APIGatewayWebSocketResolver() + router = Router() + order = [] + + def app_middleware(app_, next_middleware): + order.append("app") + return next_middleware(app_) + + def router_middleware(app_, next_middleware): + order.append("router") + return next_middleware(app_) + + def route_middleware(app_, next_middleware): + order.append("route_mw") + return next_middleware(app_) + + app.use(middlewares=[app_middleware]) + router.use(middlewares=[router_middleware]) + + @router.route("$connect", middlewares=[route_middleware]) + def connect(): + order.append("handler") + + # WHEN the router is included and the event resolved + app.include_router(router) + result = app.resolve(connect_event, lambda_context) + + # THEN globals run in merge order (app then router), then the route-level middleware, then the handler + assert result == {"statusCode": 200} + assert order == ["app", "router", "route_mw", "handler"] + + +def test_router_exception_handlers_are_merged(connect_event, lambda_context): + # GIVEN an exception handler registered on a router + router = Router() + + @router.exception_handler(ValueError) + def handle_value_error(exc: ValueError): + return {"error": str(exc)}, 400 + + @router.route("$connect") + def connect(): + raise ValueError("router boom") + + # WHEN the router is included and the event resolved + app = APIGatewayWebSocketResolver() + app.include_router(router) + result = app.resolve(connect_event, lambda_context) + + # THEN the router-registered exception handler is used + assert result == {"statusCode": 400, "body": '{"error":"router boom"}'} + + +def test_included_router_route_wins_over_app_route(connect_event, lambda_context): + # GIVEN the app and a router both registering the same route key + app = APIGatewayWebSocketResolver() + router = Router() + + @app.on_connect() + def app_connect(): + return "from app" + + @router.route("$connect") + def router_connect(): + return "from router" + + # WHEN the router is included and the event resolved + app.include_router(router) + result = app.resolve(connect_event, lambda_context) + + # THEN the router-registered handler takes precedence (last-wins merge semantics) + assert result == {"statusCode": 200, "body": "from router"} From 8319886c300fda1d470c9cd5308263d9e783acda Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 19:46:40 +1000 Subject: [PATCH 5/7] docs(event_handler): finish API Gateway WebSocket resolver docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the overview and workflow mermaid diagrams, accessing the event and Lambda context, and the sending-messages recipe: a connection store capturing connection_id and callback_url at $connect, an in-resolver broadcast, and a separate function pushing a result to the requesting client later — with GoneException handling, ManageConnections IAM, and a custom domain caveat. --- .../event_handler/api_gateway_websocket.md | 129 ++++++++++++++++++ .../accessing_websocket_event_and_context.py | 19 +++ .../getting_started_with_testing_event.json | 2 +- .../src/my_connection_store.py | 21 +++ .../src/working_with_post_to_connection.py | 46 +++++++ ...ng_with_post_to_connection_report_ready.py | 25 ++++ 6 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py create mode 100644 examples/event_handler_api_gateway_websocket/src/my_connection_store.py create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py create mode 100644 examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md index d1a4b31e387..49796870827 100644 --- a/docs/core/event_handler/api_gateway_websocket.md +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -6,6 +6,32 @@ status: new Event Handler for Amazon API Gateway WebSocket APIs. +```mermaid +stateDiagram-v2 + direction LR + EventSource: WebSocket client + Gateway: Amazon API Gateway + LambdaInit: Lambda invocation + EventHandler: Event Handler + EventHandlerResolver: Route event based on route key + YourLogic: Run your registered route handler + EventHandlerResolverBuilder: Normalize response (statusCode/body) + LambdaResponse: Lambda response + + EventSource --> Gateway: $connect, messages, $disconnect + Gateway --> LambdaInit: Route selection + + LambdaInit --> EventHandler + + state EventHandler { + [*] --> EventHandlerResolver: app.resolve(event, context) + EventHandlerResolver --> YourLogic + YourLogic --> EventHandlerResolverBuilder + } + + EventHandler --> LambdaResponse +``` + ## Key Features * Route WebSocket events by route key with dedicated `$connect`, `$disconnect`, and `$default` decorators @@ -185,6 +211,109 @@ Register handlers for specific exception types with `@app.exception_handler`; it ???+ note "Unhandled exceptions never reach the client" An exception with no registered handler is logged and becomes a bare `{"statusCode": 500}`. If it propagated to the Lambda runtime instead, the runtime's error response — exception type, message, and stack trace — would be delivered to the connected client. +### Accessing the event and Lambda context + +Inside handlers and middlewares, `app.current_event` is the current `APIGatewayWebSocketEvent` ([Event Source Data Classes](../../utilities/data_classes.md){target="_blank"} utility) and `app.lambda_context` is the Lambda context. + +Use `app.append_context` / `app.context` to share data between middlewares and handlers within a single invocation — the context is cleared after each `resolve`. + +=== "accessing_websocket_event_and_context.py" + + ```python hl_lines="9 13 14" + --8<-- "examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py" + ``` + + 1. Route, connection, and identity details live in `request_context`. + 2. `json_body` parses the message body; `body` and `decoded_body` are also available. + 3. The standard Lambda context object. + +### Sending messages to connected clients + +Returned bodies only reply to the **calling** client, on routes with a route response. To push a message to a client at any time — including after the invocation that received the request has long finished — call the API Gateway Management API's `PostToConnection` with the connection ID, against the endpoint the `callback_url` property gives you. + +Capture `connection_id` and `callback_url` together at `$connect` and persist them: that record is everything **any** process needs to push to that client later. + +The example shows both sending situations. Inside the resolver, `broadcast` pushes to every stored connection from the same invocation, building the client from the current event. + +Outside it, `submitReport` acknowledges a long-running request immediately, and a **separate Lambda function** — for example the final state of a Step Functions workflow — pushes the finished result to the one client that asked, using the stored `callback_url`. Progress updates are that same call made mid-work. + +A stale connection ID raises `GoneException` — the client is no longer connected, so treat it as a signal to remove that connection from your store. The posting function also needs the `execute-api:ManageConnections` IAM permission on `arn:aws:execute-api:{region}:{account}:{api-id}/{stage}/POST/@connections/*`. + +=== "working_with_post_to_connection.py" + + ```python hl_lines="2 16 28 35 41" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py" + ``` + + 1. The only thing the WebSocket resolver and the pushing function share. + 2. Captured at `$connect`: the only moment you get for free, and all anyone needs to post to this client later. + 3. Acknowledge immediately; the result is pushed by another function when it is ready. For the ack to reach the client, configure a route response on `submitReport`. + 4. Sending from inside the resolver: the current event provides the endpoint directly. + +=== "working_with_post_to_connection_report_ready.py" + + ```python hl_lines="18 23" + --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py" + ``` + + 1. The stored `callback_url` rebuilds the Management API client in any execution environment. + 2. The client disconnected before the result was ready — remove it from the store. + +=== "my_connection_store.py" + + ```python + --8<-- "examples/event_handler_api_gateway_websocket/src/my_connection_store.py" + ``` + +???+ note "Custom domain names" + `callback_url` is built from the event's domain name. When clients connect through a custom domain whose API mapping path differs from the stage name, construct the endpoint yourself instead: `https://{api_id}.execute-api.{region}.amazonaws.com/{stage}`. + +## Event Handler workflow + +### Connection lifecycle + +Connection accepted, message exchanged, connection closed. + +
+```mermaid +sequenceDiagram + participant Client as WebSocket client + participant Gateway as Amazon API Gateway + participant Lambda as Lambda (Event Handler) + Client->>Gateway: Upgrade request (wss://) + Gateway->>Lambda: $connect event + Lambda-->>Gateway: {"statusCode": 200} + Gateway-->>Client: 101 Switching Protocols + Client->>Gateway: {"action": "orderUpdate", ...} + Gateway->>Lambda: orderUpdate event (route selection expression) + Lambda-->>Gateway: {"statusCode": 200, "body": "..."} + Gateway-->>Client: body (only with a route response configured) + Client->>Gateway: Close frame + Gateway->>Lambda: $disconnect event (best effort) + Lambda-->>Gateway: {"statusCode": 200} +``` +
+ +### Rejected connection + +Connection rejected by the `$connect` handler. + +
+```mermaid +sequenceDiagram + participant Client as WebSocket client + participant Gateway as Amazon API Gateway + participant Lambda as Lambda (Event Handler) + Client->>Gateway: Upgrade request (wss://) + Gateway->>Lambda: $connect event + Lambda-->>Gateway: {"statusCode": 401} + Gateway-->>Client: HTTP 401 (upgrade refused) + Note over Gateway,Lambda: $disconnect may still be delivered for this connection + Gateway->>Lambda: $disconnect event + Lambda-->>Gateway: {"statusCode": 200} +``` +
+ ## Testing your code Test your handlers by passing a WebSocket event payload to `app.resolve()` — no mocks of API Gateway needed. diff --git a/examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py b/examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py new file mode 100644 index 00000000000..2a9c006fe11 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/accessing_websocket_event_and_context.py @@ -0,0 +1,19 @@ +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() + + +@app.on_default() +def default(): + request_context = app.current_event.request_context # (1)! + return { + "connectionId": request_context.connection_id, + "routeKey": request_context.route_key, + "message": app.current_event.json_body, # (2)! + "remainingTimeMs": app.lambda_context.get_remaining_time_in_millis(), # (3)! + } + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json index 188d5869326..e6320cdbdd6 100644 --- a/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json +++ b/examples/event_handler_api_gateway_websocket/src/getting_started_with_testing_event.json @@ -46,4 +46,4 @@ "apiId": "asasasas" }, "isBase64Encoded": false -} \ No newline at end of file +} diff --git a/examples/event_handler_api_gateway_websocket/src/my_connection_store.py b/examples/event_handler_api_gateway_websocket/src/my_connection_store.py new file mode 100644 index 00000000000..ee3cbf573df --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/my_connection_store.py @@ -0,0 +1,21 @@ +# Illustrative connection store. In production, back these functions with a durable store +# such as a DynamoDB table: the WebSocket resolver and the functions that push results run +# in different Lambda execution environments and cannot share process memory. + +_connections: dict[str, str] = {} + + +def save(connection_id: str, callback_url: str) -> None: + _connections[connection_id] = callback_url + + +def get_callback_url(connection_id: str) -> str: + return _connections[connection_id] + + +def all_connection_ids() -> list[str]: + return list(_connections) + + +def delete(connection_id: str) -> None: + _connections.pop(connection_id, None) diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py new file mode 100644 index 00000000000..b7b920c93f6 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py @@ -0,0 +1,46 @@ +import boto3 +import my_connection_store # (1)! + +from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver +from aws_lambda_powertools.utilities.typing import LambdaContext + +app = APIGatewayWebSocketResolver() + + +def start_report_generation(report_id: str, connection_id: str) -> None: ... # e.g. start a Step Functions execution + + +@app.on_connect() +def connect(): + request_context = app.current_event.request_context + my_connection_store.save(request_context.connection_id, request_context.callback_url) # (2)! + + +@app.on_disconnect() +def disconnect(): + my_connection_store.delete(app.current_event.request_context.connection_id) + + +@app.route("submitReport") +def submit_report(): + request = app.current_event.json_body + start_report_generation(request["reportId"], app.current_event.request_context.connection_id) + return {"status": "accepted"} # (3)! + + +@app.route("broadcast") +def broadcast(): + client = boto3.client( + "apigatewaymanagementapi", + endpoint_url=app.current_event.request_context.callback_url, # (4)! + ) + message: str = app.current_event.json_body["message"] + for connection_id in my_connection_store.all_connection_ids(): + try: + client.post_to_connection(ConnectionId=connection_id, Data=message.encode()) + except client.exceptions.GoneException: + my_connection_store.delete(connection_id) + + +def lambda_handler(event: dict, context: LambdaContext) -> dict: + return app.resolve(event, context) diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py new file mode 100644 index 00000000000..7c1afda3fc8 --- /dev/null +++ b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py @@ -0,0 +1,25 @@ +# A separate Lambda function, in a different execution environment from the WebSocket +# resolver — invoked when the work completes, e.g. as the final state of a Step Functions +# workflow. The connection store is the only thing the two functions share. + +import json + +import boto3 +import my_connection_store + +from aws_lambda_powertools import Logger +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() + + +def lambda_handler(event: dict, context: LambdaContext) -> None: + connection_id = event["connectionId"] + callback_url = my_connection_store.get_callback_url(connection_id) # (1)! + + client = boto3.client("apigatewaymanagementapi", endpoint_url=callback_url) + try: + client.post_to_connection(ConnectionId=connection_id, Data=json.dumps(event["report"]).encode()) + except client.exceptions.GoneException: # (2)! + logger.info("Client disconnected before the report was ready", connection_id=connection_id) + my_connection_store.delete(connection_id) From 073e08edab3440f3c8ac179d21772b4248475087 Mon Sep 17 00:00:00 2001 From: Rodos Date: Mon, 13 Jul 2026 23:08:43 +1000 Subject: [PATCH 6/7] docs(event_handler): polish WebSocket resolver docs from final review --- README.md | 1 + .../event_handler/api_gateway_websocket/router.py | 2 +- docs/core/event_handler/api_gateway_websocket.md | 8 +++++--- docs/index.md | 2 +- .../src/my_connection_store.py | 3 ++- mkdocs.yml | 4 ++-- 6 files changed, 12 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 34d3cac965c..fa93dfa9e3f 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ Core utilities such as Tracing, Logging, Metrics, and Event Handler are availabl * **[Event handler: AppSync](https://docs.powertools.aws.dev/lambda/python/latest/core/event_handler/appsync/)** - AppSync event handler for Lambda Direct Resolver and Amplify GraphQL Transformer function * **[Event handler: AppSync Events](https://docs.powertools.aws.dev/lambda/python/latest/core/event_handler/appsync_events/)** - AppSync Events handler for real-time WebSocket APIs with pub/sub pattern * **[Event handler: API Gateway, ALB, Lambda Function URL, VPC Lattice](https://docs.powertools.aws.dev/lambda/python/latest/core/event_handler/api_gateway/)** - REST/HTTP API event handler for Lambda functions invoked via Amazon API Gateway, ALB, Lambda Function URL, and VPC Lattice +* **[Event handler: API Gateway WebSocket](https://docs.powertools.aws.dev/lambda/python/latest/core/event_handler/api_gateway_websocket/)** - WebSocket API event handler for Lambda functions invoked via Amazon API Gateway WebSocket APIs * **[Event handler: Agents for Amazon Bedrock](https://docs.powertools.aws.dev/lambda/python/latest/core/event_handler/bedrock_agents/)** - Create Agents for Amazon Bedrock, automatically generating OpenAPI schemas * **[Middleware factory](https://docs.powertools.aws.dev/lambda/python/latest/utilities/middleware_factory/)** - Decorator factory to create your own middleware to run logic before, and after each Lambda invocation * **[Parameters](https://docs.powertools.aws.dev/lambda/python/latest/utilities/parameters/)** - Retrieve and cache parameter values from Parameter Store, Secrets Manager, AppConfig, or DynamoDB diff --git a/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py index 31854a520c7..345fb0f85db 100644 --- a/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py +++ b/aws_lambda_powertools/event_handler/api_gateway_websocket/router.py @@ -24,7 +24,7 @@ class Router(BaseRouter): Registers route handlers for WebSocket route keys so they can be split across files and later included in an `APIGatewayWebSocketResolver` via `include_router`. - Parameters + Attributes ---------- context : dict Dictionary to store context information accessible across route handlers diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md index 49796870827..8b0ffb77d1c 100644 --- a/docs/core/event_handler/api_gateway_websocket.md +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -121,7 +121,9 @@ A returned dict is always body content — it is never inspected for a `statusCo The status code has different effects depending on the route: * **`$connect`**: the status code decides the connection — 2xx accepts the WebSocket upgrade, anything else rejects it. -* **All other routes**: the body is delivered back to the client only when the route has a route response configured, and a non-2xx status code does **not** drop the connection. +* **All other routes**: the body is delivered back to the client only on routes with a [route response](#terminology), and a non-2xx status code does **not** drop the connection. + +An event whose route key has no registered handler emits a `PowertoolsUserWarning` and returns `{"statusCode": 400}` — on `$connect`, that rejects the connection. ## Advanced @@ -145,7 +147,7 @@ Execution order is: global middlewares (`app.use`), then route-level `middleware ### Authentication patterns -`$connect` is the only route where credentials exist: headers and cookies are present on the WebSocket handshake only, and later messages carry nothing but the connection ID. There are two ways to bridge that gap. +`$connect` is the only route where credentials exist: headers and cookies are present on the WebSocket handshake only, and later messages carry nothing from the client but the connection ID. There are two ways to bridge that gap. #### Lambda authorizer @@ -229,7 +231,7 @@ Use `app.append_context` / `app.context` to share data between middlewares and h ### Sending messages to connected clients -Returned bodies only reply to the **calling** client, on routes with a route response. To push a message to a client at any time — including after the invocation that received the request has long finished — call the API Gateway Management API's `PostToConnection` with the connection ID, against the endpoint the `callback_url` property gives you. +Returned bodies only reply to the **calling** client. To push a message to a client at any time — including after the invocation that received the request has long finished — call the API Gateway Management API's `PostToConnection` with the connection ID, against the endpoint the `callback_url` property gives you. Capture `connection_id` and `callback_url` together at `$connect` and persist them: that record is everything **any** process needs to push to that client later. diff --git a/docs/index.md b/docs/index.md index 887b35b23fa..2610854f39f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -52,7 +52,7 @@ Powertools for AWS Lambda (Python) is a developer toolkit to implement Serverles | [Tracer](./core/tracer.md) | Decorators and utilities to trace Lambda function handlers, and both synchronous and asynchronous functions | | [Logger](./core/logger.md) | Structured logging made easier, and target to enrich structured logging with key Lambda context details | | [Metrics](./core/metrics.md) | Custom Metrics created asynchronously via CloudWatch Embedded Metric Format (EMF) | -| [Event Handler](./core/event_handler/api_gateway.md) | Event handler for API Gateway, ALB, Lambda Function URL, VPC Lattice, AppSync, and Bedrock Agents | +| [Event Handler](./core/event_handler/api_gateway.md) | Event handler for API Gateway, ALB, Lambda Function URL, VPC Lattice, API Gateway WebSocket, AppSync, and Bedrock Agents | | [Parameters](./utilities/parameters.md) | Retrieve and cache parameter values from Parameter Store, Secrets Manager, AppConfig, or DynamoDB | | [Parser](./utilities/parser.md) | Data parsing and deep validation using Pydantic | | [Batch Processing](./utilities/batch.md) | Handle partial failures for SQS, Kinesis Data Streams, and DynamoDB Streams | diff --git a/examples/event_handler_api_gateway_websocket/src/my_connection_store.py b/examples/event_handler_api_gateway_websocket/src/my_connection_store.py index ee3cbf573df..cca942adcdc 100644 --- a/examples/event_handler_api_gateway_websocket/src/my_connection_store.py +++ b/examples/event_handler_api_gateway_websocket/src/my_connection_store.py @@ -1,6 +1,7 @@ # Illustrative connection store. In production, back these functions with a durable store # such as a DynamoDB table: the WebSocket resolver and the functions that push results run -# in different Lambda execution environments and cannot share process memory. +# in different Lambda execution environments and cannot share process memory. Use a TTL — +# $disconnect delivery is best-effort, so missed disconnects would leak entries forever. _connections: dict[str, str] = {} diff --git a/mkdocs.yml b/mkdocs.yml index 00297992dfe..b438d1a57c4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -22,9 +22,9 @@ nav: - Event Handler: - core/event_handler/api_gateway.md - core/event_handler/openapi.md + - core/event_handler/api_gateway_websocket.md - core/event_handler/appsync.md - core/event_handler/appsync_events.md - - core/event_handler/api_gateway_websocket.md - core/event_handler/bedrock_agents.md - utilities/parameters.md - utilities/batch.md @@ -244,9 +244,9 @@ plugins: - core/metrics.md - core/metrics/datadog.md - core/event_handler/api_gateway.md + - core/event_handler/api_gateway_websocket.md - core/event_handler/appsync.md - core/event_handler/appsync_events.md - - core/event_handler/api_gateway_websocket.md - core/event_handler/bedrock_agents.md Utilities: - utilities/parameters.md From 72c6b670667480a88c12d3e86772d9ddebc8a848 Mon Sep 17 00:00:00 2001 From: Rodos Date: Tue, 14 Jul 2026 00:03:14 +1000 Subject: [PATCH 7/7] fix(examples): appease SonarCloud in WebSocket examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds botocore timeouts to Management API clients, moves client construction out of the handler, and enables function and access logging in the getting-started SAM template. Example code, it turns out, is held to production standards — fair enough. --- docs/core/event_handler/api_gateway_websocket.md | 6 +++--- .../sam/getting_started_with_websocket_api.yaml | 15 +++++++++++++++ .../src/working_with_post_to_connection.py | 3 +++ ...orking_with_post_to_connection_report_ready.py | 8 +++++++- 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/docs/core/event_handler/api_gateway_websocket.md b/docs/core/event_handler/api_gateway_websocket.md index 8b0ffb77d1c..84518052518 100644 --- a/docs/core/event_handler/api_gateway_websocket.md +++ b/docs/core/event_handler/api_gateway_websocket.md @@ -65,7 +65,7 @@ You need an API Gateway WebSocket API with its routes integrated with your Lambd === "getting_started_with_websocket_api.yaml" - ```yaml hl_lines="27 50-80" + ```yaml hl_lines="39 62-92" --8<-- "examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml" ``` @@ -243,7 +243,7 @@ A stale connection ID raises `GoneException` — the client is no longer connect === "working_with_post_to_connection.py" - ```python hl_lines="2 16 28 35 41" + ```python hl_lines="2 18 30 37 44" --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py" ``` @@ -254,7 +254,7 @@ A stale connection ID raises `GoneException` — the client is no longer connect === "working_with_post_to_connection_report_ready.py" - ```python hl_lines="18 23" + ```python hl_lines="24 29" --8<-- "examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py" ``` diff --git a/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml b/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml index ef5d7b22e69..789b067cadd 100644 --- a/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml +++ b/examples/event_handler_api_gateway_websocket/sam/getting_started_with_websocket_api.yaml @@ -18,6 +18,18 @@ Resources: Properties: Handler: app.lambda_handler CodeUri: websocket_handler + LoggingConfig: + LogGroup: !Ref WebSocketHandlerLogs + + WebSocketHandlerLogs: + Type: AWS::Logs::LogGroup + Properties: + RetentionInDays: 7 + + WebSocketAccessLogs: + Type: AWS::Logs::LogGroup + Properties: + RetentionInDays: 7 WebSocketApi: Type: AWS::ApiGatewayV2::Api @@ -85,6 +97,9 @@ Resources: ApiId: !Ref WebSocketApi StageName: prod AutoDeploy: true + AccessLogSettings: + DestinationArn: !GetAtt WebSocketAccessLogs.Arn + Format: '{"requestId":"$context.requestId","routeKey":"$context.routeKey","status":"$context.status","connectionId":"$context.connectionId"}' WebSocketInvokePermission: Type: AWS::Lambda::Permission diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py index b7b920c93f6..158cce56960 100644 --- a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py +++ b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection.py @@ -1,9 +1,11 @@ import boto3 import my_connection_store # (1)! +from botocore.config import Config from aws_lambda_powertools.event_handler import APIGatewayWebSocketResolver from aws_lambda_powertools.utilities.typing import LambdaContext +boto_config = Config(connect_timeout=3, read_timeout=5, retries={"total_max_attempts": 2}) app = APIGatewayWebSocketResolver() @@ -33,6 +35,7 @@ def broadcast(): client = boto3.client( "apigatewaymanagementapi", endpoint_url=app.current_event.request_context.callback_url, # (4)! + config=boto_config, ) message: str = app.current_event.json_body["message"] for connection_id in my_connection_store.all_connection_ids(): diff --git a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py index 7c1afda3fc8..52808a85f15 100644 --- a/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py +++ b/examples/event_handler_api_gateway_websocket/src/working_with_post_to_connection_report_ready.py @@ -6,18 +6,24 @@ import boto3 import my_connection_store +from botocore.config import Config from aws_lambda_powertools import Logger from aws_lambda_powertools.utilities.typing import LambdaContext logger = Logger() +boto_config = Config(connect_timeout=3, read_timeout=5, retries={"total_max_attempts": 2}) + + +def management_client_for(callback_url: str): + return boto3.client("apigatewaymanagementapi", endpoint_url=callback_url, config=boto_config) def lambda_handler(event: dict, context: LambdaContext) -> None: connection_id = event["connectionId"] callback_url = my_connection_store.get_callback_url(connection_id) # (1)! - client = boto3.client("apigatewaymanagementapi", endpoint_url=callback_url) + client = management_client_for(callback_url) try: client.post_to_connection(ConnectionId=connection_id, Data=json.dumps(event["report"]).encode()) except client.exceptions.GoneException: # (2)!