Litestar integration example with ASGI middleware

examples/litestar_app.py

@@ -7,13 +7,7 @@
 
 from dishka import Provider, Scope, provide
 from dishka.integrations.base import Depends
-from dishka.integrations.litestar import (
-    inject,
-    setup_dishka,
-    startup_dishka,
-    make_dishka_container,
-    shutdown_dishka
-)
+from dishka.integrations.litestar import inject, DishkaApp
 
 
 # app core
@@ -25,7 +19,7 @@ def get(self) -> str:
 
 class FakeDbGateway(DbGateway):
     def get(self) -> str:
-        return "Hello"
+        return "Hello123"
 
 
 class Interactor:
@@ -62,13 +56,8 @@ def create_app():
         level=logging.WARNING,
         format='%(asctime)s  %(process)-7s %(module)-20s %(message)s',
     )
-    app = Litestar(
-        route_handlers=[MainController],
-        on_startup=[startup_dishka],
-        on_shutdown=[shutdown_dishka],
-        before_request=make_dishka_container
-    )
-    return setup_dishka(app, [InteractorProvider(), AdaptersProvider()])
+    app = Litestar(route_handlers=[MainController])
+    return DishkaApp([InteractorProvider(), AdaptersProvider()], app)
 
 
 if __name__ == "__main__":

src/dishka/integrations/litestar.py

@@ -1,9 +1,12 @@
 from inspect import Parameter
-from typing import Optional, Sequence, get_type_hints
+from typing import Optional, get_type_hints
 
-from litestar import Litestar, Request
+from litestar import Request
+from litestar.enums import ScopeType
+from litestar.types import ASGIApp, Receive, Scope, Send
 
-from dishka import Provider, make_async_container
+from dishka.async_container import AsyncContainer, AsyncContextWrapper
+from dishka.integrations.asgi import BaseDishkaApp
 from dishka.integrations.base import wrap_injection
 
 
@@ -32,20 +35,29 @@ def inject(func):
     )
 
 
-async def make_dishka_container(request: Request):
-    request_container = await request.app.state.dishka_container().__aenter__()
-    request.state.dishka_container = request_container
+def make_add_request_container_middleware(app: ASGIApp):
+    async def middleware(scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != ScopeType.HTTP:
+            await app(scope, receive, send)
+            return
 
+        request = Request(scope)
+        async with request.app.state.dishka_container(
+                {Request: request},
+        ) as request_container:
+            request.state.dishka_container = request_container
+            await app(scope, receive, send)
 
-async def startup_dishka(app: Litestar):
-    container = await app.state.dishka_container_wrapper.__aenter__()
-    app.state.dishka_container = container
+    return middleware
 
 
-async def shutdown_dishka(app: Litestar):
-    await app.state.dishka_container_wrapper.__aexit__(None, None, None)
+class DishkaApp(BaseDishkaApp):
+    def _init_request_middleware(
+            self, app, container_wrapper: AsyncContextWrapper,
+    ):
+        app.asgi_handler = make_add_request_container_middleware(
+            app.asgi_handler,
+        )
 
-
-def setup_dishka(app: Litestar, providers: Sequence[Provider]) -> Litestar:
-    app.state.dishka_container_wrapper = make_async_container(*providers)
-    return app
+    def _app_startup(self, app, container: AsyncContainer):
+        app.state.dishka_container = container