lavkach3

1
from uuid import uuid4
2

3
from starlette.types import ASGIApp, Receive, Scope, Send
4
from taskiq import SimpleRetryMiddleware
5

6
from core.db.session import set_session_context, reset_session_context, session
7

8

9
class SQLAlchemyMiddleware:
10
    def __init__(self, app: ASGIApp) -> None:
11
        self.app = app
12

13
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
14
        session_id = str(uuid4())
15
        context = set_session_context(session_id=session_id)
16

17
        try:
18
            await self.app(scope, receive, send)
19
        except Exception as e:
20
            raise e
21
        finally:
22
            await session.remove()
23
            reset_session_context(context=context)
24

25
SimpleRetryMiddleware