feat(event_handler): add API Gateway WebSocket resolver

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

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",

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",
+]

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)

aws_lambda_powertools/event_handler/api_gateway_websocket/base.py

@@ -0,0 +1,59 @@
+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
+
+    @abstractmethod
+    def use(self, middlewares: list[Callable]) -> None:
+        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