✨ feat(sqla): Model support dependency injection

Author: ProgramRipper

nonebot_plugin_orm/__init__.py

@@ -1,28 +1,32 @@
 from __future__ import annotations
 
+import sys
 import logging
-from typing import Any
+from asyncio import gather
+from operator import methodcaller
 from functools import wraps, partial
+from typing import Any, AsyncGenerator
 
+from nonebot.params import Depends
 from nonebot.log import LoguruHandler
 from nonebot.plugin import PluginMetadata
+import sqlalchemy.ext.asyncio as sa_asyncio
 from sqlalchemy.util import greenlet_spawn
 from nonebot.matcher import current_matcher
 from nonebot import logger, require, get_driver
 from sqlalchemy import URL, Table, MetaData, make_url
-from sqlalchemy.ext.asyncio import (
-    AsyncEngine,
-    AsyncSession,
-    async_sessionmaker,
-    create_async_engine,
-    async_scoped_session,
-)
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
 from . import migrate
 from .model import Model
 from .utils import StreamToLogger
 from .config import Config, plugin_config
 
+if sys.version_info >= (3, 10):
+    from typing import Annotated
+else:
+    from typing_extensions import Annotated
+
 __all__ = (
     # __init__
     "init_orm",
@@ -75,7 +79,7 @@
 _binds: dict[type[Model], AsyncEngine]
 _engines: dict[str, AsyncEngine]
 _metadatas: dict[str, MetaData]
-_session_factory: async_sessionmaker[AsyncSession]
+_session_factory: sa_asyncio.async_sessionmaker[AsyncSession]
 
 _driver = get_driver()
 
@@ -86,7 +90,7 @@ async def init_orm():
 
     _init_engines()
     _init_table()
-    _session_factory = async_sessionmaker(
+    _session_factory = sa_asyncio.async_sessionmaker(
         _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
     )
 
@@ -100,21 +104,34 @@ async def init_orm():
 
 
 @wraps(lambda: None)  # NOTE: for dependency injection
-def get_session(**local_kw: Any) -> AsyncSession:
+def get_session(**local_kw: Any) -> sa_asyncio.AsyncSession:
     try:
         return _session_factory(**local_kw)
     except NameError:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
 
-def get_scoped_session() -> async_scoped_session[AsyncSession]:
+AsyncSession = Annotated[sa_asyncio.AsyncSession, Depends(get_session)]
+
+
+async def get_scoped_session() -> (
+    AsyncGenerator[sa_asyncio.async_scoped_session[AsyncSession], None]
+):
     try:
-        return async_scoped_session(
+        scoped_session = async_scoped_session(
             _session_factory, scopefunc=partial(current_matcher.get, None)
         )
+        yield scoped_session
     except NameError:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
+    await gather(*map(methodcaller("close"), scoped_session.registry.registry.values()))
+
+
+async_scoped_session = Annotated[
+    sa_asyncio.async_scoped_session[AsyncSession], Depends(get_scoped_session)
+]
+
 
 def _create_engine(engine: str | URL | dict[str, Any] | AsyncEngine) -> AsyncEngine:
     if isinstance(engine, AsyncEngine):

nonebot_plugin_orm/model.py

@@ -3,17 +3,20 @@
 import sys
 from inspect import Parameter, Signature
 from typing import TYPE_CHECKING, Any, ClassVar
+from typing_extensions import Self, Unpack, Annotated
 
-from sqlalchemy import Table, MetaData
+from nonebot.params import Depends
 from nonebot import get_plugin_by_module_name
+from sqlalchemy import Table, MetaData, select
+from pydantic.typing import get_args, get_origin
 from sqlalchemy.orm import Mapped, DeclarativeBase
 
 from .utils import DependsInner, get_annotations
 
 if sys.version_info >= (3, 9):
-    from typing import Annotated, get_args, get_origin
+    from typing import Annotated
 else:
-    from typing_extensions import Annotated, get_args, get_origin
+    from typing_extensions import Annotated
 
 __all__ = ("Model",)
 
@@ -27,12 +30,39 @@
 }
 
 
+class Model(DeclarativeBase):
+    metadata = MetaData(naming_convention=NAMING_CONVENTION)
+
+    if TYPE_CHECKING:
+        __args__: ClassVar[tuple[type[Self], Unpack[tuple[Any, ...]]]]
+        __origin__: type[Annotated]
+
+        __table__: ClassVar[Table]
+        __bind_key__: ClassVar[str]
+
+    def __init_subclass__(cls) -> None:
+        _setup_di(cls)
+        _setup_tablename(cls)
+
+        super().__init_subclass__()
+
+        if not hasattr(cls, "__table__"):
+            return
+
+        _setup_bind(cls)
+
+
 def _setup_di(cls: type[Model]) -> None:
     """Get signature for NoneBot's dependency injection,
     and set annotations for SQLAlchemy declarative class.
     """
+    from . import async_scoped_session
 
-    parameters: list[Parameter] = []
+    parameters: list[Parameter] = [
+        Parameter(
+            "__session__", Parameter.KEYWORD_ONLY, annotation=async_scoped_session
+        )
+    ]
 
     annotations: dict[str, Any] = {}
     for base in reversed(cls.__mro__):
@@ -41,13 +71,13 @@ def _setup_di(cls: type[Model]) -> None:
     for name, type_annotation in annotations.items():
         # Check if the attribute is both a dependent and a mapped column
         depends_inner = None
-        if isinstance(get_origin(type_annotation), Annotated):
+        if get_origin(type_annotation) is Annotated:
             (type_annotation, *extra_args) = get_args(type_annotation)
             depends_inner = next(
                 (x for x in extra_args if isinstance(x, DependsInner)), None
             )
 
-        if not isinstance(get_origin(type_annotation), Mapped):
+        if get_origin(type_annotation) is not Mapped:
             continue
 
         default = getattr(cls, name, Signature.empty)
@@ -71,7 +101,14 @@ def _setup_di(cls: type[Model]) -> None:
         if default is not Signature.empty and not isinstance(default, Mapped):
             delattr(cls, name)
 
-    cls.__signature__ = Signature(parameters)
+    async def dependency(
+        *, __session__: async_scoped_session, **kwargs: Any
+    ) -> Model | None:
+        return await __session__.scalar(select(cls).filter_by(**kwargs))
+
+    dependency.__signature__ = Signature(parameters)
+    cls.__args__ = (Model, Depends(dependency))
+    cls.__origin__ = Annotated
 
 
 def _setup_tablename(cls: type[Model]) -> None:
@@ -95,23 +132,3 @@ def _setup_bind(cls: type[Model]) -> None:
             bind_key = ""
 
     cls.__table__.info["bind_key"] = bind_key
-
-
-class Model(DeclarativeBase):
-    metadata = MetaData(naming_convention=NAMING_CONVENTION)
-
-    if TYPE_CHECKING:
-        __table__: ClassVar[Table]
-        __bind_key__: ClassVar[str]
-        __signature__: ClassVar[Signature]
-
-    def __init_subclass__(cls) -> None:
-        _setup_di(cls)
-        _setup_tablename(cls)
-
-        super().__init_subclass__()
-
-        if not hasattr(cls, "__table__"):
-            return
-
-        _setup_bind(cls)