feat(event-handler): Add exception_handler support

aws_lambda_powertools/event_handler/api_gateway.py

@@ -10,7 +10,7 @@
 from enum import Enum
 from functools import partial
 from http import HTTPStatus
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
 from aws_lambda_powertools.event_handler import content_types
 from aws_lambda_powertools.event_handler.exceptions import ServiceError
@@ -435,6 +435,7 @@ def __init__(
         self._proxy_type = proxy_type
         self._routes: List[Route] = []
         self._route_keys: List[str] = []
+        self._exception_handlers: Dict[Type, Callable] = {}
         self._cors = cors
         self._cors_enabled: bool = cors is not None
         self._cors_methods: Set[str] = {"OPTIONS"}
@@ -618,7 +619,11 @@ def _call_route(self, route: Route, args: Dict[str, str]) -> ResponseBuilder:
                 ),
                 route,
             )
-        except Exception:
+        except Exception as exc:
+            handler = self._lookup_exception_handler(exc)
+            if handler:
+                return ResponseBuilder(handler(exc))
+
             if self._debug:
                 # If the user has turned on debug mode,
                 # we'll let the original exception propagate so
@@ -676,6 +681,18 @@ def include_router(self, router: "Router", prefix: Optional[str] = None) -> None
 
             self.route(*route)(func)
 
+    def exception_handler(self, exception):
+        def register_exception_handler(func: Callable):
+            self._exception_handlers[exception] = func
+
+        return register_exception_handler
+
+    def _lookup_exception_handler(self, exc: Exception) -> Optional[Callable]:
+        for cls in type(exc).__mro__:
+            if cls in self._exception_handlers:
+                return self._exception_handlers[cls]
+        return None
+
 
 class Router(BaseRouter):
     """Router helper class to allow splitting ApiGatewayResolver into multiple files"""

tests/functional/event_handler/test_api_gateway.py

@@ -163,7 +163,7 @@ def patch_func():
     def handler(event, context):
         return app.resolve(event, context)
 
-    # Also check check the route configurations
+    # Also check the route configurations
     routes = app._routes
     assert len(routes) == 5
     for route in routes:
@@ -1076,3 +1076,30 @@ def foo():
 
     assert result["statusCode"] == 200
     assert result["headers"]["Content-Type"] == content_types.APPLICATION_JSON
+
+
+def test_exception_handler():
+    # GIVEN a resolver with an exception handler defined for ValueError
+    app = ApiGatewayResolver(proxy_type=ProxyEventType.APIGatewayProxyEvent)
+
+    @app.exception_handler(ValueError)
+    def handle_value_error(ex: ValueError):
+        print(f"request path is '{app.current_event.path}'")
+        return Response(
+            status_code=418,
+            content_type=content_types.TEXT_HTML,
+            body=str(ex),
+        )
+
+    @app.get("/my/path")
+    def get_lambda() -> Response:
+        raise ValueError("Foo!")
+
+    # WHEN calling the event handler
+    # AND a ValueError is raised
+    result = app(LOAD_GW_EVENT, {})
+
+    # THEN call the exception_handler
+    assert result["statusCode"] == 418
+    assert result["headers"]["Content-Type"] == content_types.TEXT_HTML
+    assert result["body"] == "Foo!"