Allow for error handler registration

chalice/app.py

@@ -31,6 +31,7 @@
     from chalice.local import LambdaContext
 
 _PARAMS = re.compile(r'{\w+}')
+ErrorHandlerFuncType = Callable[[Exception], Any]
 MiddlewareFuncType = Callable[[Any, Callable[[Any], Any]], Any]
 UserHandlerFuncType = Callable[..., Any]
 HeadersType = Dict[str, Union[str, List[str]]]
@@ -730,6 +731,17 @@ def _middleware_wrapper(
             return func
         return _middleware_wrapper
 
+    def error(
+            self,
+            exception: Exception
+    ) -> Callable[[ErrorHandlerFuncType], Any]:
+        def _error_wrapper(
+                func: ErrorHandlerFuncType
+        ) -> ErrorHandlerFuncType:
+            self.register_error(exception, func)
+            return func
+        return _error_wrapper
+
     def authorizer(self, ttl_seconds: Optional[int] = None,
                    execution_role: Optional[str] = None,
                    name: Optional[str] = None,
@@ -941,6 +953,10 @@ def _register_handler(self, handler_type: str, name: str,
                           options: Optional[Dict[Any, Any]] = None) -> None:
         raise NotImplementedError("_register_handler")
 
+    def register_error(self, exception: Exception,
+                       func: ErrorHandlerFuncType) -> None:
+        raise NotImplementedError("register_error")
+
     def register_middleware(self, func: MiddlewareFuncType,
                             event_type: str = 'all') -> None:
         raise NotImplementedError("register_middleware")
@@ -957,11 +973,21 @@ def __init__(self) -> None:
         self.api: APIGateway = APIGateway()
         self.handler_map: Dict[str, Callable[..., Any]] = {}
         self.middleware_handlers: List[Tuple[MiddlewareFuncType, str]] = []
+        self.error_handlers: List[Tuple[ErrorHandlerFuncType, str]] = []
 
     def register_middleware(self, func: MiddlewareFuncType,
                             event_type: str = 'all') -> None:
         self.middleware_handlers.append((func, event_type))
 
+    def register_error(self, exception: Any,
+                       func: ErrorHandlerFuncType) -> None:
+        if not issubclass(exception, Exception):
+            raise TypeError(
+                f"{exception.__name__} is not a subclass of Exception."
+                "Error handlers can only be registered for Exception classes."
+            )
+        self.error_handlers.append((func, exception.__name__))
+
     def _do_register_handler(self, handler_type: str, name: str,
                              user_handler: UserHandlerFuncType,
                              wrapped_handler: Callable[..., Any], kwargs: Any,
@@ -1346,6 +1372,7 @@ def __call__(self, event: Any, context: Any) -> Dict[str, Any]:
         handler = RestAPIEventHandler(
             self.routes, self.api, self.log, self.debug,
             middleware_handlers=self._get_middleware_handlers('http'),
+            error_handlers=self.error_handlers
         )
         self.current_request: \
             Optional[Request] = handler.create_request_object(event, context)
@@ -1788,7 +1815,10 @@ def __call__(self, event: Dict[str, Any],
 class RestAPIEventHandler(BaseLambdaHandler):
     def __init__(self, route_table: Dict[str, Dict[str, RouteEntry]],
                  api: APIGateway, log: logging.Logger, debug: bool,
-                 middleware_handlers: Optional[List[Callable[..., Any]]] = None
+                 middleware_handlers: Optional[
+                     List[Callable[..., Any]]] = None,
+                 error_handlers: Optional[
+                     List[Tuple[ErrorHandlerFuncType, str]]] = None
                  ) -> None:
         self.routes: Dict[str, Dict[str, RouteEntry]] = route_table
         self.api: APIGateway = api
@@ -1800,6 +1830,9 @@ def __init__(self, route_table: Dict[str, Dict[str, RouteEntry]],
             middleware_handlers = []
         self._middleware_handlers: \
             List[Callable[..., Any]] = middleware_handlers
+        if error_handlers is None:
+            error_handlers = []
+        self._error_handlers = error_handlers
 
     def _global_error_handler(self, event: Any,
                               get_response: Callable[..., Any]) -> Response:
@@ -1925,13 +1958,29 @@ def _get_view_function_response(self, view_function: Callable[..., Any],
         except ChaliceViewError as e:
             # Any chalice view error should propagate.  These
             # get mapped to various HTTP status codes in API Gateway.
-            response = Response(body={'Code': e.__class__.__name__,
-                                      'Message': str(e)},
-                                status_code=e.STATUS_CODE)
-        except Exception:
-            response = self._unhandled_exception_to_response()
+            response = self._get_error_handler_response(e) or \
+                Response(body={'Code': e.__class__.__name__,
+                               'Message': str(e)},
+                         status_code=e.STATUS_CODE)
+        except Exception as e:
+            response = self._get_error_handler_response(e) or \
+                self._unhandled_exception_to_response()
         return response
 
+    def _get_error_handler_response(self, e: Exception) -> Any:
+        # Loops through the registered error handlers and returns the first
+        # `Response` result from handlers. If no handlers are matched or no
+        # matched handlers returned a `Response`, returns None to allow for
+        # chalice to handle the error.
+        raised = e.__class__.__name__
+        handlers = (func for func, exc_type in self._error_handlers if
+                    exc_type == raised)
+        for func in handlers:
+            response = func(e)
+            if isinstance(response, Response):
+                return response
+        return
+
     def _unhandled_exception_to_response(self) -> Response:
         headers: HeadersType = {}
         path = getattr(self.current_request, 'path', 'unknown')
@@ -2207,6 +2256,13 @@ def register_middleware(self, func: Callable,
             )
         )
 
+    def register_error(self, exception: Exception, func: Callable) -> None:
+        self._deferred_registrations.append(
+            lambda app, options: app.register_error(
+                exception, func
+            )
+        )
+
     def _register_handler(self, handler_type: str, name: str,
                           user_handler: UserHandlerFuncType,
                           wrapped_handler: Any, kwargs: Dict[str, Any],

docs/source/topics/errorhandler.rst

@@ -0,0 +1,92 @@
+=====================
+Custom Error Handling
+=====================
+
+While chalice middleware allow for catching of user defined errors, exceptions
+raised by a third party library can't be seen by the middleware and chalice
+will set the response without giving the middleware a chance to see the
+exception. These error handlers will only by used in the case of 'http' event,
+as the middleware for other types of events can catch other exceptions
+(see :ref:`middleware-error-handling-rest`).
+
+In the case where you want to return your own ``Response`` for those exceptions
+you can register an error handler to intercept the exception.
+
+Below is an example of Chalice error hanler:
+
+.. code-block:: python
+
+    from chalice import Chalice, Response
+    from thirdparty import func, ThirdPartyError
+
+    app = Chalice(app_name='demo-errorhandler')
+
+    @app.error(ThirdPartyError)
+    def my_error_handler(error: Exception):
+        app.log.error("ThirdPartyError was raised")
+        return Response(
+            body=e.__class__.__name__,
+            status_code=500
+        )
+
+    @app.route('/')
+    def index():
+        return func()
+
+In this example, our error handler is registered to catch ``ThirdPartyError``
+raised in a http event. In this case, if `func` were to raise a
+``ThirdPartyError``, ``my_error_handler`` will be called and our custom
+``Response`` will be returned.
+
+Writing Error Handlers
+======================
+
+Error handlers must adhere to these requirements:
+
+* Must be a callable object that accepts one parameter. It will be of
+  ``Exception`` type.
+* Must return a response. If the response is not of ``chalice.Response`` type,
+  Chalice will either try to call the next error handler registered for the
+  same error or in the event where no handlers return a valid response, Chalice
+  will return a ``chalice.ChaliceViewError``.
+
+
+Error Propagation
+-----------------
+
+Cahlice will propagatet the error through all registered handlers until a valid
+response is returned. If no handlers return a valid response, chalice will
+handle the error as if no handlers were registered.
+
+.. code-block:: python
+
+    @app.error(ThirdPartyError)
+    def my_error_handler_1(error: Exception):
+        if error.message == '1':
+            return Response(
+                body='Error 1 was raised',
+                status_code=200
+            )
+
+    @app.error(ThirdPartyError)
+    def my_error_handler_2(error: Exception):
+        if error.message == '2':
+            return Response(
+                body='Error 2 was raised',
+                status_code=400
+            )
+
+    @app.route('/1')
+    def index():
+        # The response from `my_error_handler_1` will be returned
+        raise ThirdPartyError('1')
+
+    @app.route('/2')
+    def index():
+        # The response from `my_error_handler_2` will be returned
+        raise ThirdPartyError('2')
+
+    @app.route('/3')
+    def index():
+        # A ChaliceViewError will be returned
+        raise ThirdPartyError('3')

docs/source/topics/index.rst

@@ -25,3 +25,4 @@ Topics
    experimental
    testing
    middleware
+   errorhandler

docs/source/topics/middleware.rst

@@ -96,6 +96,8 @@ If an exception is raised in a Lambda handler and no middleware catches the
 exception, the exception will be returned back to the client that invoked
 the Lambda function.
 
+.. _middleware-error-handling-rest:
+
 Rest APIs
 ~~~~~~~~~