fix: use custom integrations for context-independent use (#62)

Author: emptybutton

src/ttt/infrastructure/adapters/game_tasks.py

@@ -3,7 +3,7 @@
 
 from ttt.application.game.game.ports.game_tasks import GameTasks
 from ttt.infrastructure.taskiq.tasks.make_ai_move_in_game_task import (
-    make_ai_move_in_game_broker_task,
+    make_ai_move_in_game_task,
 )
 
 
@@ -16,4 +16,4 @@ async def make_ai_move(
         ai_id: UUID,
         /,
     ) -> None:
-        await make_ai_move_in_game_broker_task.kiq(user_id, game_id, ai_id)
+        await make_ai_move_in_game_task.kiq(user_id, game_id, ai_id)

src/ttt/infrastructure/dishka/__init__.py

@@ -0,0 +1,9 @@
+from collections.abc import Mapping
+from typing import Any, Protocol
+
+from dishka.async_container import AsyncContextWrapper
+
+
+class NextContainer(Protocol):
+    def __call__(self, context: Mapping[Any, Any] = {}) -> AsyncContextWrapper:
+        ...

src/ttt/infrastructure/dishka/next_container.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+from dishka.integrations.taskiq import CONTAINER_NAME
+from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult
+
+from ttt.infrastructure.dishka.next_container import NextContainer
+
+
+class TaskiqNextContainerMiddleware(TaskiqMiddleware):
+    def __init__(self, next_container: NextContainer) -> None:
+        super().__init__()
+        self._next_container = next_container
+
+    async def pre_execute(
+        self,
+        message: TaskiqMessage,
+    ) -> TaskiqMessage:
+        next_container = self._next_container({TaskiqMessage: message})
+
+        container = await next_container.__aenter__()  # noqa: PLC2801
+        message.labels[CONTAINER_NAME] = container
+        return message
+
+    async def on_error(
+        self,
+        message: TaskiqMessage,  # noqa: ARG002
+        result: TaskiqResult[Any],
+        exception: BaseException,  # noqa: ARG002
+    ) -> None:
+        if CONTAINER_NAME in result.labels:
+            await result.labels[CONTAINER_NAME].close()
+            del result.labels[CONTAINER_NAME]
+
+    async def post_execute(
+        self,
+        message: TaskiqMessage,  # noqa: ARG002
+        result: TaskiqResult[Any],
+    ) -> None:
+        if CONTAINER_NAME in result.labels:
+            await result.labels[CONTAINER_NAME].close()
+            del result.labels[CONTAINER_NAME]

src/ttt/infrastructure/taskiq/middlewares.py

@@ -3,5 +3,5 @@
     complete_stars_purchase_payment_task,
 )
 from ttt.infrastructure.taskiq.tasks.make_ai_move_in_game_task import (
-    make_ai_move_in_game_broker_task,
+    make_ai_move_in_game_task,
 )

src/ttt/infrastructure/taskiq/tasks/__init__.py

@@ -12,6 +12,7 @@
 
 
 @nats_tasks.task(
+    task_name="stars_purchase-stars_purchase-complete_stars_purchase_payment",
     subject="stars_purchase.stars_purchase.complete_stars_purchase_payment",
     pull_subscribe=PullSubscribe(lambda js, subject: js.pull_subscribe(
         subject,

src/ttt/infrastructure/taskiq/tasks/complete_stars_purchase_payment_task.py

@@ -9,6 +9,7 @@
 
 
 @nats_tasks.task(
+    task_name="game-game-make_ai_move_in_game",
     subject="game.game.make_ai_move_in_game",
     pull_subscribe=PullSubscribe(lambda js, subject: js.pull_subscribe(
         subject,
@@ -17,7 +18,7 @@
     )),
 )
 @inject(patch_module=True)
-async def make_ai_move_in_game_broker_task(
+async def make_ai_move_in_game_task(
     user_id: int,
     game_id: UUID,
     ai_id: UUID,

src/ttt/infrastructure/taskiq/tasks/make_ai_move_in_game_task.py

@@ -3,8 +3,6 @@
 from types import TracebackType
 from typing import Self
 
-from dishka import AsyncContainer
-from dishka.integrations.taskiq import setup_dishka
 from taskiq.receiver import Receiver
 
 
@@ -19,10 +17,7 @@ async def __aenter__(self) -> Self:
         await self._task_group.__aenter__()
         return self
 
-    async def __call__(self, container: AsyncContainer) -> None:
-        for receiver in self._receivers:
-            setup_dishka(container, receiver.broker)
-
+    async def __call__(self) -> None:
         await gather(*(
             receiver.broker.startup()
             for receiver in self._receivers

src/ttt/infrastructure/taskiq/worker.py

@@ -0,0 +1,25 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from dishka import AsyncContainer
+from dishka.async_container import AsyncContextWrapper
+
+from ttt.infrastructure.dishka.next_container import NextContainer
+
+
+@dataclass
+class NextContainerWithFilledContext(NextContainer):
+    _root_container: AsyncContainer
+    _context_type_hints: tuple[Any, ...]
+
+    def __call__(
+        self, context: Mapping[Any, Any] = {},
+    ) -> AsyncContextWrapper:
+        return self._root_container(self._filled_context(context))
+
+    def _filled_context(self, context: Mapping[Any, Any]) -> dict[Any, Any]:
+        return {
+            hint | None: context.get(hint)
+            for hint in self._context_type_hints
+        }

src/ttt/main/common/next_container.py

@@ -19,6 +19,7 @@
     FromComponent,
     Provider,
     Scope,
+    from_context,
     provide,
 )
 from dishka.integrations.aiogram import AiogramMiddlewareData
@@ -165,15 +166,11 @@ class NoMessageInEventError(Exception):
 
 
 class PresentationProvider(Provider):
-    @provide(scope=Scope.REQUEST)
-    def provide_event(self, event: TelegramObject) -> TelegramObject | None:
-        return event
-
-    @provide(scope=Scope.REQUEST)
-    def provide_aiogram_middleware_data(
-        self, data: AiogramMiddlewareData,
-    ) -> AiogramMiddlewareData | None:
-        return data
+    provide_aiogram_middleware_data = from_context(
+        AiogramMiddlewareData | None,
+        scope=Scope.REQUEST,
+    )
+    provide_event = from_context(TelegramObject | None, scope=Scope.REQUEST)
 
     @provide(scope=Scope.APP)
     def provide_strage(self, redis: Redis) -> BaseStorage:
@@ -347,13 +344,6 @@ def provide_callback_query(
             case _:
                 return None
 
-    @provide(scope=Scope.REQUEST)
-    def provide_fsm_context(
-        self,
-        middleware_data: AiogramMiddlewareData,
-    ) -> FSMContext:
-        return cast(FSMContext, middleware_data["state"])
-
     @provide(scope=Scope.REQUEST)
     def provide_stars_purchase_payment_gateway(
         self,
@@ -437,7 +427,9 @@ class ApplicationProvider(Provider):
     )
     provide_cancel_game = provide(CancelGame, scope=Scope.REQUEST)
     provide_make_move_in_game = provide(MakeMoveInGame, scope=Scope.REQUEST)
-    provide_make_ai_move_in_game = provide(MakeAiMoveInGame, scope=Scope.REQUEST)
+    provide_make_ai_move_in_game = provide(
+        MakeAiMoveInGame, scope=Scope.REQUEST,
+    )
     provide_view_game = provide(ViewGame, scope=Scope.REQUEST)
 
     provide_accept_invitation_to_game = provide(