⚡ perf(sqla)!: sync function and async dependency of session

Author: ProgramRipper

nonebot_plugin_orm/__init__.py

@@ -23,7 +23,7 @@
 from . import migrate
 from .param import ORMParam
 from .config import Config, plugin_config
-from .utils import LoguruHandler, StreamToLogger, get_subclasses
+from .utils import LoguruHandler, StreamToLogger, coroutine, get_subclasses
 
 if sys.version_info >= (3, 9):
     from typing import Annotated
@@ -118,18 +118,21 @@ def get_session(**local_kw: Any) -> sa_async.AsyncSession:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
 
-AsyncSession = Annotated[sa_async.AsyncSession, Depends(get_session)]
+# NOTE: NoneBot DI will run sync function in thread pool executor,
+# which is poor performance for this simple function, so we wrap it as a coroutine function.
+AsyncSession = Annotated[sa_async.AsyncSession, Depends(coroutine(get_session))]
 
 
-async def get_scoped_session() -> sa_async.async_scoped_session[sa_async.AsyncSession]:
+def get_scoped_session() -> sa_async.async_scoped_session[sa_async.AsyncSession]:
     try:
         return _scoped_sessions
     except NameError:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
 
 async_scoped_session = Annotated[
-    sa_async.async_scoped_session[sa_async.AsyncSession], Depends(get_scoped_session)
+    sa_async.async_scoped_session[sa_async.AsyncSession],
+    Depends(coroutine(get_scoped_session)),
 ]
 
 

nonebot_plugin_orm/utils.py

@@ -251,6 +251,14 @@ def get_subclasses(cls: type[_T]) -> Generator[type[_T], None, None]:
         yield from get_subclasses(subclass)
 
 
+def coroutine(func: Callable[_P, _T]) -> Callable[_P, Coroutine[Any, Any, _T]]:
+    @wraps(func)
+    async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 if sys.version_info >= (3, 10):
     from inspect import get_annotations as get_annotations  # nopycln: import
 else: