Added ExceptionMiddleware

Author: eadwinCode

ellar/core/factory.py

@@ -7,6 +7,8 @@
 
 from ellar.constants import MODULE_METADATA, MODULE_WATERMARK
 from ellar.core import Config
+from ellar.core.exceptions.interfaces import IExceptionMiddlewareService
+from ellar.core.exceptions.service import ExceptionMiddlewareService
 from ellar.core.main import App
 from ellar.core.modules import ModuleBase
 from ellar.core.modules.ref import create_module_ref_factor
@@ -95,6 +97,9 @@ def _create_app(
 
         config = Config(app_configured=True, config_module=config_module)
         injector = EllarInjector(auto_bind=config.INJECTOR_AUTO_BIND)
+        injector.container.register(
+            IExceptionMiddlewareService, ExceptionMiddlewareService
+        )
         injector.container.register_instance(config, concrete_type=Config)
         cls._build_modules(app_module=module, injector=injector, config=config)
 

ellar/core/main.py

@@ -8,6 +8,10 @@
 from ellar.core.context import ExecutionContext, IExecutionContext
 from ellar.core.datastructures import State, URLPath
 from ellar.core.events import EventHandler, RouterEventManager
+from ellar.core.exceptions.interfaces import (
+    IExceptionHandler,
+    IExceptionMiddlewareService,
+)
 from ellar.core.guard import GuardCanActivate
 from ellar.core.middleware import (
     CORSMiddleware,
@@ -57,7 +61,6 @@ def __init__(
         self._injector: EllarInjector = injector
 
         self._global_guards = [] if global_guards is None else list(global_guards)
-        self._exception_handlers = dict(t.cast(dict, self.config.EXCEPTION_HANDLERS))
         self._user_middleware = list(t.cast(list, self.config.MIDDLEWARE))
 
         self.on_startup = RouterEventManager(
@@ -173,14 +176,11 @@ def config(self) -> Config:  # type: ignore
         return self._config
 
     def build_middleware_stack(self) -> ASGIApp:
-        error_handler = None
-        exception_handlers = {}
-
-        for key, value in self._exception_handlers.items():
-            if key in (500, Exception):
-                error_handler = value
-            else:
-                exception_handlers[key] = value
+        service_middleware = self.injector.get(
+            IExceptionMiddlewareService  # type:ignore
+        )
+        service_middleware.build_exception_handlers()
+        error_handler = service_middleware.get_500_error_handler()
 
         middleware = (
             [
@@ -212,7 +212,9 @@ def build_middleware_stack(self) -> ASGIApp:
             + self._user_middleware
             + [
                 Middleware(
-                    ExceptionMiddleware, handlers=exception_handlers, debug=self.debug
+                    ExceptionMiddleware,
+                    exception_middleware_service=service_middleware,
+                    debug=self.debug,
                 ),
             ]
         )
@@ -265,20 +267,11 @@ def _finalize_app_initialization(self) -> None:
 
     def add_exception_handler(
         self,
-        exc_class_or_status_code: t.Union[int, t.Type[Exception]],
-        handler: t.Callable,
-    ) -> None:  # pragma: no cover
-        self._exception_handlers[exc_class_or_status_code] = handler
-        self.middleware_stack = self.build_middleware_stack()
-
-    def exception_handler(
-        self, exc_class_or_status_code: t.Union[int, t.Type[Exception]]
-    ) -> t.Callable:  # pragma: nocover
-        def decorator(func: t.Callable) -> t.Callable:
-            self.add_exception_handler(exc_class_or_status_code, func)
-            return func
-
-        return decorator
+        exception_handler: IExceptionHandler,
+    ) -> None:
+        if exception_handler not in self.config.EXCEPTION_HANDLERS:
+            self.config.EXCEPTION_HANDLERS.append(exception_handler)
+            self.rebuild_middleware_stack()
 
     def rebuild_middleware_stack(self) -> None:
         self.middleware_stack = self.build_middleware_stack()

ellar/core/middleware/exceptions.py

@@ -0,0 +1,65 @@
+import typing as t
+
+from starlette.exceptions import HTTPException
+
+from ellar.constants import SCOPE_EXECUTION_CONTEXT_PROVIDER
+from ellar.types import ASGIApp, TMessage, TReceive, TScope, TSend
+
+if t.TYPE_CHECKING:  # pragma: no cover
+    from ellar.core.exceptions.service import ExceptionMiddlewareService
+
+
+class ExceptionMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        exception_middleware_service: "ExceptionMiddlewareService",
+        debug: bool = False,
+    ) -> None:
+        self.app = app
+        self.debug = debug  # TODO: We ought to handle 404 cases if debug is set.
+        self._exception_middleware_service = exception_middleware_service
+
+    async def __call__(self, scope: TScope, receive: TReceive, send: TSend) -> None:
+        if scope["type"] not in ("http", "websocket"):
+            await self.app(scope, receive, send)
+            return
+
+        response_started = False
+
+        async def sender(message: TMessage) -> None:
+            nonlocal response_started
+
+            if message["type"] == "http.response.start":
+                response_started = True
+            await send(message)
+
+        try:
+            await self.app(scope, receive, sender)
+        except Exception as exc:
+            handler = None
+
+            if isinstance(exc, HTTPException):
+                handler = self._exception_middleware_service.lookup_status_code_exception_handler(
+                    exc.status_code
+                )
+
+            if handler is None:
+                handler = self._exception_middleware_service.lookup_exception_handler(
+                    exc
+                )
+
+            if handler is None:
+                raise exc
+
+            if response_started:
+                msg = "Caught handled exception, but response already started."
+                raise RuntimeError(msg) from exc
+
+            if scope["type"] == "http":
+                response = await handler.catch(
+                    scope[SCOPE_EXECUTION_CONTEXT_PROVIDER], exc
+                )
+                await response(scope, receive, sender)
+            elif scope["type"] == "websocket":
+                await handler.catch(scope[SCOPE_EXECUTION_CONTEXT_PROVIDER], exc)