Adding exception handler + resolvers

Author: leandrodamascena

aws_lambda_powertools/event_handler/api_gateway.py

@@ -17,6 +17,7 @@
 from typing_extensions import override
 
 from aws_lambda_powertools.event_handler import content_types
+from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager
 from aws_lambda_powertools.event_handler.exceptions import NotFoundError, ServiceError
 from aws_lambda_powertools.event_handler.openapi.config import OpenAPIConfig
 from aws_lambda_powertools.event_handler.openapi.constants import (
@@ -1543,6 +1544,7 @@ def __init__(
         self.processed_stack_frames = []
         self._response_builder_class = ResponseBuilder[BaseProxyEvent]
         self.openapi_config = OpenAPIConfig()  # starting an empty dataclass
+        self.exception_handler_manager = ExceptionHandlerManager()
         self._has_response_validation_error = response_validation_error_http_code is not None
         self._response_validation_error_http_code = self._validate_response_validation_error_http_code(
             response_validation_error_http_code,
@@ -2409,7 +2411,7 @@ def not_found_handler():
                 return Response(status_code=204, content_type=None, headers=_headers, body="")
 
             # Customer registered 404 route? Call it.
-            custom_not_found_handler = self._lookup_exception_handler(NotFoundError)
+            custom_not_found_handler = self.exception_handler_manager.lookup_exception_handler(NotFoundError)
             if custom_not_found_handler:
                 return custom_not_found_handler(NotFoundError())
 
@@ -2482,26 +2484,10 @@ def not_found(self, func: Callable | None = None):
         return self.exception_handler(NotFoundError)(func)
 
     def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
-        def register_exception_handler(func: Callable):
-            if isinstance(exc_class, list):  # pragma: no cover
-                for exp in exc_class:
-                    self._exception_handlers[exp] = func
-            else:
-                self._exception_handlers[exc_class] = func
-            return func
-
-        return register_exception_handler
-
-    def _lookup_exception_handler(self, exp_type: type) -> Callable | None:
-        # Use "Method Resolution Order" to allow for matching against a base class
-        # of an exception
-        for cls in exp_type.__mro__:
-            if cls in self._exception_handlers:
-                return self._exception_handlers[cls]
-        return None
+        return self.exception_handler_manager.exception_handler(exc_class=exc_class)
 
     def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuilder | None:
-        handler = self._lookup_exception_handler(type(exp))
+        handler = self.exception_handler_manager.lookup_exception_handler(type(exp))
         if handler:
             try:
                 return self._response_builder_class(response=handler(exp), serializer=self._serializer, route=route)
@@ -2595,7 +2581,7 @@ def include_router(self, router: Router, prefix: str | None = None) -> None:
         self._router_middlewares = self._router_middlewares + router._router_middlewares
 
         logger.debug("Appending Router exception_handler into App exception_handler.")
-        self._exception_handlers.update(router._exception_handlers)
+        self.exception_handler_manager.update_exception_handlers(router._exception_handlers)
 
         # use pointer to allow context clearance after event is processed e.g., resolve(evt, ctx)
         router.context = self.context

aws_lambda_powertools/event_handler/appsync.py

@@ -5,6 +5,7 @@
 import warnings
 from typing import TYPE_CHECKING, Any, Callable
 
+from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager
 from aws_lambda_powertools.event_handler.graphql_appsync.exceptions import InvalidBatchResponse, ResolverNotFoundError
 from aws_lambda_powertools.event_handler.graphql_appsync.router import Router
 from aws_lambda_powertools.utilities.data_classes import AppSyncResolverEvent
@@ -52,6 +53,7 @@ def __init__(self):
         Initialize a new instance of the AppSyncResolver.
         """
         super().__init__()
+        self.exception_handler_manager = ExceptionHandlerManager()
         self.context = {}  # early init as customers might add context before event resolution
         self._exception_handlers: dict[type, Callable] = {}
 
@@ -151,7 +153,7 @@ def lambda_handler(event, context):
                 Router.current_event = data_model(event)
                 response = self._call_single_resolver(event=event, data_model=data_model)
         except Exception as exp:
-            response_builder = self._lookup_exception_handler(type(exp))
+            response_builder = self.exception_handler_manager.lookup_exception_handler(type(exp))
             if response_builder:
                 return response_builder(exp)
             raise
@@ -493,31 +495,4 @@ def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
             A decorator function that registers the exception handler.
         """
 
-        def register_exception_handler(func: Callable):
-            if isinstance(exc_class, list):  # pragma: no cover
-                for exp in exc_class:
-                    self._exception_handlers[exp] = func
-            else:
-                self._exception_handlers[exc_class] = func
-            return func
-
-        return register_exception_handler
-
-    def _lookup_exception_handler(self, exp_type: type) -> Callable | None:
-        """
-        Looks up the registered exception handler for the given exception type or its base classes.
-
-        Parameters
-        ----------
-        exp_type (type):
-            The exception type to look up the handler for.
-
-        Returns
-        -------
-        Callable | None:
-            The registered exception handler function if found, otherwise None.
-        """
-        for cls in exp_type.__mro__:
-            if cls in self._exception_handlers:
-                return self._exception_handlers[cls]
-        return None
+        return self.exception_handler_manager.exception_handler(exc_class=exc_class)

aws_lambda_powertools/event_handler/appsync_events.py

@@ -1,56 +1,137 @@
 from __future__ import annotations
 
+import asyncio
+import logging
 from typing import TYPE_CHECKING, Any, Callable
 
-from aws_lambda_powertools.event_handler.events_appsync._registry import ResolverEventsRegistry
+from aws_lambda_powertools.event_handler.events_appsync.router import Router
+from aws_lambda_powertools.event_handler.exception_handling import ExceptionHandlerManager
+from aws_lambda_powertools.utilities.data_classes.appsync_resolver_events_event import AppSyncResolverEventsEvent
 
 if TYPE_CHECKING:
-    from aws_lambda_powertools.utilities.data_classes.appsync_resolver_events_event import AppSyncResolverEventsEvent
     from aws_lambda_powertools.utilities.typing.lambda_context import LambdaContext
 
-
-class AppSyncEventsResolver:
+logger = logging.getLogger(__name__)
+class AppSyncEventsResolver(Router):
     """
     AppSync Events API Resolver
     """
 
     def __init__(self):
+        super().__init__()
         self.context = {}  # early init as customers might add context before event resolution
-        self._publish_registry = ResolverEventsRegistry()
-        self._async_publish_registry = ResolverEventsRegistry()
-        self._subscribe_registry = ResolverEventsRegistry()
-        self._async_subscribe_registry = ResolverEventsRegistry()
+        self.exception_handler_manager = ExceptionHandlerManager()
+        self._exception_handlers: dict[type, Callable] = {}
+
+    def __call__(
+        self,
+        event: dict,
+        context: LambdaContext,
+        data_model: type[AppSyncResolverEventsEvent] = AppSyncResolverEventsEvent,
+    ) -> Any:
+        """Implicit lambda handler which internally calls `resolve`"""
+        return self.resolve(event, context, data_model)
 
     def resolve(
         self,
         event: AppSyncResolverEventsEvent,
         context: LambdaContext,
+        data_model: type[AppSyncResolverEventsEvent] = AppSyncResolverEventsEvent,
     ) -> Any:
         """Resolves the response based on the provide event and decorator operation and namespaces"""
-        print(self._publish_registry.__dict__)
 
-    def publish(
+        self.lambda_context = context
+        Router.lambda_context = context
+
+        Router.current_event = data_model(event)
+        self.current_event = data_model(event)
+
+        try:
+            if self.current_event.info.operation == "PUBLISH":
+                response = self._call_publish_events(payload=self.current_event.events, data_model=data_model)
+            else:
+                response = self._call_subscribe_events(event=event, data_model=data_model)
+        except Exception as exp:
+            response_builder = self.exception_handler_manager.lookup_exception_handler(type(exp))
+            if response_builder:
+                return response_builder(exp)
+            raise
+
+        self.clear_context()
+
+        return response
+
+    def _call_subscribe_events(self, payload: list[dict[str, Any]], data_model: type[AppSyncResolverEventsEvent]) -> Any:
+        # PLACEHOLDER
+        pass
+
+    def _call_publish_events(self, payload: list[dict[str, Any]], data_model: type[AppSyncResolverEventsEvent]) -> Any:
+        """Call single event resolver
+
+        Parameters
+        ----------
+        event : dict
+            Event
+        data_model : type[AppSyncResolverEvent]
+            Data_model to decode AppSync event, by default it is of AppSyncResolverEvent type or subclass of it
+        """
+
+        result = []
+        logger.debug("Processing direct resolver event")
+
+        #self.current_event = data_model(event)
+        resolver = self._publish_registry.find_resolver(self.current_event.info.channel_path)
+        if not resolver:
+            print(f"No resolver found for '{self.current_event.info.channel_path}'")
+        print(resolver)
+
+        if not resolver["aggregate"]:
+            return resolver["func"](payload=self.current_event.events)
+        else:
+            for i in self.current_event.events:
+                result.append(resolver["func"](payload=i))
+
+            return result
+
+    def on_publish(
         self,
         path: str = "/default/*",
         aggregate: bool = True,
     ) -> Callable:
         return self._publish_registry.register(path=path, aggregate=aggregate)
 
-    def async_publish(
+    def async_on_publish(
         self,
         path: str = "/default/*",
         aggregate: bool = True,
     ) -> Callable:
         return self._async_publish_registry.register(path=path, aggregate=aggregate)
 
-    def subscribe(
+    def on_subscribe(
         self,
         path: str = "/default/*",
     ) -> Callable:
         return self._subscribe_registry.register(path=path)
 
-    def async_subscribe(
+    def async_on_subscribe(
         self,
         path: str = "/default/*",
     ) -> Callable:
         return self._async_subscribe_registry.register(path=path)
+
+    def exception_handler(self, exc_class: type[Exception] | list[type[Exception]]):
+        """
+        A decorator function that registers a handler for one or more exception types.
+
+        Parameters
+        ----------
+        exc_class (type[Exception] | list[type[Exception]])
+            A single exception type or a list of exception types.
+
+        Returns
+        -------
+        Callable:
+            A decorator function that registers the exception handler.
+        """
+
+        return self.exception_handler_manager.exception_handler(exc_class=exc_class)

aws_lambda_powertools/event_handler/events_appsync/base.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Callable
+
+
+class BaseRouter(ABC):
+    """Abstract base class for Router (resolvers)"""
+
+    @abstractmethod
+    def on_publish(
+        self,
+        path: str = "/default/*",
+        aggregate: bool = True,
+    ) -> Callable:
+        raise NotImplementedError
+
+    @abstractmethod
+    def async_on_publish(
+        self,
+        path: str = "/default/*",
+        aggregate: bool = True,
+    ) -> Callable:
+        raise NotImplementedError
+
+    @abstractmethod
+    def on_subscribe(
+        self,
+        path: str = "/default/*",
+    ) -> Callable:
+        raise NotImplementedError
+
+    @abstractmethod
+    def async_on_subscribe(
+        self,
+        path: str = "/default/*",
+    ) -> Callable:
+        raise NotImplementedError
+    @abstractmethod
+    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

aws_lambda_powertools/event_handler/events_appsync/router.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+from aws_lambda_powertools.event_handler.events_appsync._registry import ResolverEventsRegistry
+from aws_lambda_powertools.event_handler.events_appsync.base import BaseRouter
+from aws_lambda_powertools.utilities.data_classes.appsync_resolver_events_event import AppSyncResolverEventsEvent
+
+if TYPE_CHECKING:
+    from aws_lambda_powertools.utilities.data_classes.appsync_resolver_event import AppSyncResolverEvent
+    from aws_lambda_powertools.utilities.typing.lambda_context import LambdaContext
+
+
+class Router(BaseRouter):
+
+    context: dict
+    current_event: AppSyncResolverEventsEvent | None = None
+    lambda_context: LambdaContext | None = None
+
+    def __init__(self):
+        self.context = {}  # early init as customers might add context before event resolution
+        self._publish_registry = ResolverEventsRegistry()
+        self._async_publish_registry = ResolverEventsRegistry()
+        self._subscribe_registry = ResolverEventsRegistry()
+        self._async_subscribe_registry = ResolverEventsRegistry()
+
+    def on_publish(
+        self,
+        path: str = "/default/*",
+        aggregate: bool = True,
+    ) -> Callable:
+        return self._publish_registry.register(path=path, aggregate=aggregate)
+
+    def async_on_publish(
+        self,
+        path: str = "/default/*",
+        aggregate: bool = True,
+    ) -> Callable:
+        return self._async_publish_registry.register(path=path, aggregate=aggregate)
+
+    def on_subscribe(
+        self,
+        path: str = "/default/*",
+    ) -> Callable:
+        return self._subscribe_registry.register(path=path)
+
+    def async_on_subscribe(
+        self,
+        path: str = "/default/*",
+    ) -> Callable:
+        return self._async_subscribe_registry.register(path=path)
+
+    def append_context(self, **additional_context):
+        """Append key=value data as routing context"""
+        self.context.update(**additional_context)
+
+    def clear_context(self):
+        """Resets routing context"""
+        self.context.clear()