29 lines
958 B
Python
29 lines
958 B
Python
from typing import Any, Awaitable, Callable
|
|
|
|
from aiogram import BaseMiddleware
|
|
from aiogram.types import TelegramObject
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
|
|
class DbSessionMiddleware(BaseMiddleware):
|
|
"""Opens one SQLAlchemy async session per update; commits on success."""
|
|
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
async def __call__(
|
|
self,
|
|
handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
|
|
event: TelegramObject,
|
|
data: dict[str, Any],
|
|
) -> Any:
|
|
async with self._session_factory() as session:
|
|
data["session"] = session
|
|
try:
|
|
result = await handler(event, data)
|
|
await session.commit()
|
|
return result
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|