🐛 fix(sqla): close scoped session by processor

Author: ProgramRipper

nonebot_plugin_orm/__init__.py

@@ -3,22 +3,22 @@
 import sys
 import logging
 from typing import Any
-from asyncio import gather
 from argparse import Namespace
-from operator import methodcaller
-from collections.abc import AsyncGenerator
-from functools import wraps, partial, lru_cache
+from contextlib import suppress
+from functools import wraps, lru_cache
 
 import click
 from nonebot.rule import Rule
+from nonebot.adapters import Event
 import sqlalchemy.ext.asyncio as sa_async
 from nonebot.permission import Permission
-from sqlalchemy.util import greenlet_spawn
+from sqlalchemy import URL, Table, MetaData
 from nonebot.params import Depends, DefaultParam
 from nonebot.plugin import Plugin, PluginMetadata
-from nonebot.matcher import Matcher, current_matcher
-from sqlalchemy import URL, Table, MetaData, make_url
+from sqlalchemy.util import ScopedRegistry, greenlet_spawn
 from sqlalchemy.log import Identified, _qual_logger_name_for_cls
+from nonebot.message import run_postprocessor, event_postprocessor
+from nonebot.matcher import Matcher, current_event, current_matcher
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from nonebot import logger, require, get_driver, get_plugin_by_module_name
 
@@ -65,7 +65,8 @@
 _engines: dict[str, AsyncEngine]
 _metadatas: dict[str, MetaData]
 _plugins: dict[str, Plugin]
-_session_factory: sa_async.async_sessionmaker[AsyncSession]
+_session_factory: sa_async.async_sessionmaker[sa_async.AsyncSession]
+_scoped_sessions: ScopedRegistry[sa_async.async_scoped_session[sa_async.AsyncSession]]
 
 _data_dir = get_data_dir(__plugin_meta__.name)
 _driver = get_driver()
@@ -93,7 +94,7 @@ async def init_orm() -> None:
 
 
 def _init_orm():
-    global _session_factory
+    global _session_factory, _scoped_sessions
 
     _init_engines()
     _init_table()
@@ -103,6 +104,12 @@ def _init_orm():
             **plugin_config.sqlalchemy_session_options,
         }
     )
+    _scoped_sessions = ScopedRegistry(
+        lambda: sa_async.async_scoped_session(
+            _session_factory, lambda: current_matcher.get(None)
+        ),
+        lambda: id(current_event.get(None)),
+    )
 
 
 @wraps(lambda: None)  # NOTE: for dependency injection
@@ -116,25 +123,34 @@ def get_session(**local_kw: Any) -> sa_async.AsyncSession:
 AsyncSession = Annotated[sa_async.AsyncSession, Depends(get_session)]
 
 
-async def get_scoped_session() -> (
-    AsyncGenerator[sa_async.async_scoped_session[AsyncSession], None]
-):
+async def get_scoped_session() -> sa_async.async_scoped_session[sa_async.AsyncSession]:
     try:
-        scoped_session = async_scoped_session(
-            _session_factory, scopefunc=partial(current_matcher.get, None)
-        )
-        yield scoped_session
+        return _scoped_sessions()
     except NameError:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
-    await gather(*map(methodcaller("close"), scoped_session.registry.registry.values()))
-
 
 async_scoped_session = Annotated[
     sa_async.async_scoped_session[sa_async.AsyncSession], Depends(get_scoped_session)
 ]
 
 
+@event_postprocessor
+def _clear_scoped_session(event: Event) -> None:
+    with suppress(KeyError):
+        del _scoped_sessions.registry[id(event)]
+
+
+@run_postprocessor
+async def _close_scoped_session(event: Event, matcher: Matcher) -> None:
+    with suppress(KeyError):
+        session: sa_async.AsyncSession = _scoped_sessions.registry[
+            id(event)
+        ].registry.registry[matcher]
+        del _scoped_sessions.registry[id(event)].registry.registry[matcher]
+        await session.close()
+
+
 def _create_engine(engine: str | URL | dict[str, Any] | AsyncEngine) -> AsyncEngine:
     if isinstance(engine, AsyncEngine):
         return engine