asgi middleware

Author: krylosov-aa

context_async_sqlalchemy/__init__.py

@@ -23,9 +23,15 @@
     close_all_sessions,
 )
 from .run_in_new_context import run_in_new_ctx
+
+from .asgi_utils import (
+    ASGIHTTPDBSessionMiddleware,
+)
+
 from .starlette_utils import (
     add_starlette_http_db_session_middleware,
     starlette_http_db_session_middleware,
+    StarletteHTTPDBSessionMiddleware,
 )
 
 from .fastapi_utils import (
@@ -53,8 +59,10 @@
     "commit_all_sessions",
     "rollback_all_sessions",
     "close_all_sessions",
+    "ASGIHTTPDBSessionMiddleware",
     "add_starlette_http_db_session_middleware",
     "starlette_http_db_session_middleware",
+    "StarletteHTTPDBSessionMiddleware",
     "fastapi_http_db_session_middleware",
     "add_fastapi_http_db_session_middleware",
 ]

context_async_sqlalchemy/asgi_utils/__init__.py

@@ -0,0 +1,5 @@
+from .middleware import ASGIHTTPDBSessionMiddleware
+
+__all__ = [
+    "ASGIHTTPDBSessionMiddleware",
+]

context_async_sqlalchemy/asgi_utils/middleware.py

@@ -0,0 +1,65 @@
+from starlette.types import ASGIApp, Receive, Scope, Send
+from http import HTTPStatus
+
+from context_async_sqlalchemy import (
+    init_db_session_ctx,
+    is_context_initiated,
+    reset_db_session_ctx,
+    auto_commit_by_status_code,
+    rollback_all_sessions,
+)
+
+
+class ASGIHTTPDBSessionMiddleware:
+    """Database session lifecycle management."""
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        """
+        Database session lifecycle management.
+        The session itself is created on demand in db_session().
+
+        Transaction auto-commit is implemented if there is no exception and
+            the response status is < 400. Otherwise, a rollback is performed.
+
+        But you can commit or rollback manually in the handler.
+        """
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        # Tests have different session management rules
+        # so if the context variable is already set, we do nothing
+        if is_context_initiated():
+            await self.app(scope, receive, send)
+            return
+
+        # We set the context here, meaning all child coroutines
+        # will receive the same context.
+        # And even if a child coroutine requests the session first,
+        # the container itself is shared, and this coroutine will
+        # add the session to container = shared context.
+        token = init_db_session_ctx()
+
+        status_code = HTTPStatus.INTERNAL_SERVER_ERROR
+
+        async def send_wrapper(message):
+            nonlocal status_code
+            if message["type"] == "http.response.start":
+                status_code = message["status"]
+            await send(message)
+
+        try:
+            await self.app(scope, receive, send_wrapper)
+            # using the status code, we decide to commit or rollback
+            # all sessions
+            await auto_commit_by_status_code(status_code)
+        except Exception:
+            # If an exception occurs, we roll all sessions back
+            await rollback_all_sessions()
+            raise
+        finally:
+            # Close all sessions and clear the context
+            await reset_db_session_ctx(token)

context_async_sqlalchemy/starlette_utils/__init__.py

@@ -1,9 +1,11 @@
 from .http_middleware import (
     add_starlette_http_db_session_middleware,
     starlette_http_db_session_middleware,
+    StarletteHTTPDBSessionMiddleware,
 )
 
 __all__ = [
     "add_starlette_http_db_session_middleware",
     "starlette_http_db_session_middleware",
+    "StarletteHTTPDBSessionMiddleware",
 ]

context_async_sqlalchemy/starlette_utils/http_middleware.py

@@ -22,6 +22,17 @@ def add_starlette_http_db_session_middleware(app: Starlette) -> None:
     )
 
 
+class StarletteHTTPDBSessionMiddleware(BaseHTTPMiddleware):
+    """Database session lifecycle management."""
+
+    async def dispatch(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        return await starlette_http_db_session_middleware(
+            request, call_next
+        )
+
+
 async def starlette_http_db_session_middleware(
     request: Request, call_next: RequestResponseEndpoint
 ) -> Response:

examples/fastapi_example/setup_app.py

@@ -26,7 +26,8 @@ def setup_app() -> FastAPI:
     """
     A convenient entry point for app configuration.
     Convenient for testing.
-    You don't have to follow my example (though I recommend it).
+
+    You don't have to follow my example here.
     """
     app = FastAPI(
         lifespan=lifespan,

examples/starlette_example/setup_app.py

@@ -6,7 +6,7 @@
 from starlette.applications import Starlette
 from starlette.routing import Route
 
-
+from context_async_sqlalchemy import ASGIHTTPDBSessionMiddleware
 from context_async_sqlalchemy.starlette_utils import (
     add_starlette_http_db_session_middleware,
 )
@@ -29,7 +29,8 @@ def setup_app() -> Starlette:
     """
     A convenient entry point for app configuration.
     Convenient for testing.
-    You don't have to follow my example.
+
+    You don't have to follow my example here.
     """
     app = Starlette(
         debug=True,