diff --git a/etebase_fastapi/main.py b/etebase_fastapi/main.py index 8e8469c..c07d975 100644 --- a/etebase_fastapi/main.py +++ b/etebase_fastapi/main.py @@ -5,6 +5,7 @@ from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.trustedhost import TrustedHostMiddleware +from .middleware import DjangoDbConnectionCleanupMiddleware from .exceptions import CustomHttpException from .msgpack import MsgpackResponse from .routers.authentication import authentication_router @@ -42,6 +43,7 @@ def create_application(prefix="", middlewares=[]): app.include_router(test_reset_view_router, prefix=f"{BASE_PATH}/test/authentication") + app.add_middleware(DjangoDbConnectionCleanupMiddleware) app.add_middleware( CORSMiddleware, allow_origin_regex="https?://.*", diff --git a/etebase_fastapi/middleware.py b/etebase_fastapi/middleware.py new file mode 100644 index 0000000..06e347b --- /dev/null +++ b/etebase_fastapi/middleware.py @@ -0,0 +1,15 @@ +from starlette.types import ASGIApp, Receive, Scope, Send +from django.db import close_old_connections, reset_queries + + +class DjangoDbConnectionCleanupMiddleware: + def __init__(self, app: ASGIApp): + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send): + reset_queries() + close_old_connections() + try: + await self.app(scope, receive, send) + finally: + close_old_connections()