lavkach3
25 строк · 720.0 Байт
1from uuid import uuid4
2
3from starlette.types import ASGIApp, Receive, Scope, Send
4from taskiq import SimpleRetryMiddleware
5
6from core.db.session import set_session_context, reset_session_context, session
7
8
9class SQLAlchemyMiddleware:
10def __init__(self, app: ASGIApp) -> None:
11self.app = app
12
13async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
14session_id = str(uuid4())
15context = set_session_context(session_id=session_id)
16
17try:
18await self.app(scope, receive, send)
19except Exception as e:
20raise e
21finally:
22await session.remove()
23reset_session_context(context=context)
24
25SimpleRetryMiddleware