Replace BaseHTTPMiddleware with much faster ASGI equalivalent

Author: puehringer

requirements.txt

@@ -20,6 +20,7 @@ python-multipart==0.0.5
 rdkit==2022.3.5
 requests==2.28.1
 SQLAlchemy==1.4.45
+starlette-context==0.3.5
 urllib3==1.26.9
 uvicorn[standard]==0.20.0
 yamlreader==3.0.4

tdp_core/dbmanager.py

@@ -8,7 +8,7 @@
 from . import manager
 from .dbview import DBConnector
 from .middleware.close_web_sessions_middleware import CloseWebSessionsMiddleware
-from .middleware.request_context_middleware import get_request
+from .middleware.request_context_plugin import get_request
 
 _log = logging.getLogger(__name__)
 
@@ -93,11 +93,14 @@ def create_web_session(self, engine_or_id: Union[Engine, str]) -> Session:
         """
         session = self.create_session(engine_or_id)
 
+        r = get_request()
+        if not r:
+            raise Exception("No request found, did you use a create_web_sesssion outside of a request?")
         try:
-            existing_sessions = get_request().state.db_sessions
+            existing_sessions = r.state.db_sessions
         except (KeyError, AttributeError):
             existing_sessions = []
-            get_request().state.db_sessions = existing_sessions
+            r.state.db_sessions = existing_sessions
         existing_sessions.append(session)
 
         return session

tdp_core/middleware/close_web_sessions_middleware.py

@@ -1,18 +1,23 @@
-from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-from starlette.requests import Request
+from fastapi import FastAPI
 
+from .request_context_plugin import get_request
 
-class CloseWebSessionsMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
-        response = await call_next(request)
 
-        try:
-            for db_session in request.state.db_sessions:
-                try:
-                    db_session.close()
-                except Exception:
-                    pass
-        except (KeyError, AttributeError):
-            pass
+# Use basic ASGI middleware instead of BaseHTTPMiddleware as it is significantly faster: https://github.com/tiangolo/fastapi/issues/2696#issuecomment-768224643
+class CloseWebSessionsMiddleware:
+    def __init__(self, app: FastAPI):
+        self.app = app
 
-        return response
+    async def __call__(self, scope, receive, send):
+        await self.app(scope, receive, send)
+
+        r = get_request()
+        if r:
+            try:
+                for db_session in r.state.db_sessions:
+                    try:
+                        db_session.close()
+                    except Exception:
+                        pass
+            except (KeyError, AttributeError):
+                pass

tdp_core/middleware/exception_handler_middleware.py

@@ -1,20 +1,23 @@
 import logging
 
-from fastapi import HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.exception_handlers import http_exception_handler
-from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-from starlette.requests import Request
 
 from ..server.utils import detail_from_exception
+from .request_context_plugin import get_request
 
 
-class ExceptionHandlerMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
+# Use basic ASGI middleware instead of BaseHTTPMiddleware as it is significantly faster: https://github.com/tiangolo/fastapi/issues/2696#issuecomment-768224643
+class ExceptionHandlerMiddleware:
+    def __init__(self, app: FastAPI):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
         try:
-            return await call_next(request)
+            await self.app(scope, receive, send)
         except Exception as e:
             logging.exception("An error occurred in FastAPI")
             return await http_exception_handler(
-                request,
+                get_request(),  # type: ignore
                 e if isinstance(e, HTTPException) else HTTPException(status_code=500, detail=detail_from_exception(e)),
             )

tdp_core/middleware/request_context_middleware.py

@@ -0,0 +1,17 @@
+from typing import Optional
+
+from starlette.requests import HTTPConnection, Request
+from starlette_context import context
+from starlette_context.plugins.base import Plugin
+
+
+def get_request() -> Request | None:
+    return context.get("request")
+
+
+class RequestContextPlugin(Plugin):
+    # The returned value will be inserted in the context with this key
+    key = "request"
+
+    async def process_request(self, request: Request | HTTPConnection) -> Optional[Request | HTTPConnection]:
+        return request

tdp_core/middleware/request_context_plugin.py

@@ -9,7 +9,7 @@
 from fastapi.security.utils import get_authorization_scheme_param
 
 from .. import manager
-from ..middleware.request_context_middleware import get_request
+from ..middleware.request_context_plugin import get_request
 from .model import ANONYMOUS_USER, LogoutReturnValue, User
 from .store.base_store import BaseStore
 
@@ -119,17 +119,18 @@ def _delegate_stores_until_not_none(self, store_method_name: str, *args):
     @property
     def current_user(self) -> Optional[User]:
         try:
-            req = get_request()
-            # Fetch the existing user from the request if there is any
-            try:
-                user = req.state.user
-                if user:
-                    return user
-            except (KeyError, AttributeError):
-                pass
-            # If there is no user, try to load it from the request and store it in the request
-            user = req.state.user = self.load_from_request(get_request())
-            return user
+            r = get_request()
+            if r:
+                # Fetch the existing user from the request if there is any
+                try:
+                    user = r.state.user
+                    if user:
+                        return user
+                except (KeyError, AttributeError):
+                    pass
+                # If there is no user, try to load it from the request and store it in the request
+                user = r.state.user = self.load_from_request(r)
+                return user
         except HTTPException:
             return None
         except Exception:

tdp_core/security/manager.py

@@ -8,6 +8,7 @@
 from fastapi.middleware.wsgi import WSGIMiddleware
 from pydantic import create_model
 from pydantic.utils import deep_update
+from starlette_context.middleware import RawContextMiddleware
 
 from ..settings.constants import default_logging_dict
 
@@ -60,7 +61,6 @@ def create_visyn_server(
     )
 
     from ..middleware.exception_handler_middleware import ExceptionHandlerMiddleware
-    from ..middleware.request_context_middleware import RequestContextMiddleware
 
     # TODO: For some reason, a @app.exception_handler(Exception) is not called here. We use a middleware instead.
     app.add_middleware(ExceptionHandlerMiddleware)
@@ -143,8 +143,10 @@ def create_visyn_server(
     for p in plugins:
         p.plugin.init_app(app)
 
-    # Add middleware to access Request "outside"
-    app.add_middleware(RequestContextMiddleware)
+    from ..middleware.request_context_plugin import RequestContextPlugin
+
+    # Use starlette-context to store the current request globally, i.e. accessible via context['request']
+    app.add_middleware(RawContextMiddleware, plugins=(RequestContextPlugin(),))
 
     # TODO: Move up?
     app.add_api_route("/health", health)  # type: ignore