将aiogram2的LifetimeControllerMiddleware适配到aiogram3遇连接回收问题

问题背景 #

在Aiogram 2中，使用LifetimeControllerMiddleware实现数据库会话管理的中间件，代码如下：

class DatabaseMiddleware(LifetimeControllerMiddleware):
    def __init__(self, async_session_maker):
        super().__init__()
        self.async_session_maker = async_session_maker

    async def pre_process(self, obj, data, *args):
        session: AsyncSession = self.async_session_maker()
        data["session"] = session
        data["repo"] = Repo(session=session)

    async def post_process(self, obj, data, *args):
        del data["repo"]
        session: AsyncSession = data.get("session")
        await session.close()

升级到Aiogram 3后，由于LifetimeControllerMiddleware被移除，改用BaseMiddleware实现，但出现SQLAlchemy连接回收警告和错误：

class DatabaseMiddleware(BaseMiddleware):

    def __init__(self, async_session_maker):
        super().__init__()
        self.async_session_maker = async_session_maker

    async def __call__(self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: Dict[str, Any]) -> Any:
            session: AsyncSession = self.async_session_maker()
            data["session"] = session
            data["repo"] = Repo(session=session)
            return await handler(event, data)

报错信息 #

SAWarning: The garbage collector is trying to clean up non-checked-in connection <AdaptedConnection <asyncpg.connection.Connection object at 0x000001ACE7450C80>>, which will be terminated. Please ensure that SQLAlchemy pooled connections are returned to the pool explicitly, either by calling close() or by using appropriate context managers to manage their lifecycle.
schema = self._schema_type_to_method[schema['type']](schema.copy(), f)
ERROR:sqlalchemy.pool.impl.QueuePool:The garbage collector is trying to clean up non-checked-in connection <AdaptedConnection <asyncpg.connection.Connection object at 0x000001ACE613F060>>, which will be terminated. Please ensure that SQLAlchemy pooled connections are returned to the pool explicitly, either by calling close() or by using appropriate context managers to manage their lifecycle.

错误原因 #

Aiogram 2的LifetimeControllerMiddleware通过post_process方法确保会话在请求处理完成后被关闭，而当前Aiogram 3的实现仅在处理前创建会话，未在处理完成后主动关闭会话，导致SQLAlchemy连接无法正常回收到连接池，最终被垃圾回收器强制清理，触发警告和错误。

解决方案 #

在Aiogram 3的BaseMiddleware中，使用try...finally结构，确保无论请求处理成功还是抛出异常，都能关闭数据库会话并回收连接，同时清理data中的相关对象。

完整正确代码 #

from typing import Callable, Dict, Any, Awaitable
from aiogram import BaseMiddleware
from aiogram.types import TelegramObject
from sqlalchemy.ext.asyncio import AsyncSession


class DatabaseMiddleware(BaseMiddleware):
    def __init__(self, async_session_maker):
        super().__init__()
        self.async_session_maker = async_session_maker

    async def __call__(
        self,
        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
        event: TelegramObject,
        data: Dict[str, Any]
    ) -> Any:
        # 创建会话和Repo实例
        session: AsyncSession = self.async_session_maker()
        data["session"] = session
        data["repo"] = Repo(session=session)
        
        try:
            # 执行请求处理逻辑
            return await handler(event, data)
        finally:
            # 无论处理结果如何，都清理资源
            del data["repo"]
            await session.close()
            del data["session"]

补充说明 #

• 使用try...finally可以保证会话在任何情况下都会被关闭，避免连接泄漏。
• 清理data中的repo和session是良好的实践，防止内存中残留不必要的对象引用。

内容的提问来源于stack exchange，提问作者Samich80