✨ feat(sqla): dependency cache and validate

Author: ProgramRipper

nonebot_plugin_orm/param.py

@@ -9,13 +9,13 @@
 
 from pydantic.fields import FieldInfo
 from nonebot.dependencies import Param
-from nonebot.params import DependParam
+from nonebot.params import Depends, DependParam
 from sqlalchemy import Row, Result, ScalarResult, select
 from sqlalchemy.sql.selectable import ExecutableReturnsRows
 from sqlalchemy.ext.asyncio import AsyncResult, AsyncScalarResult
 
 from .model import Model
-from .utils import Option, compile_dependency, generic_issubclass
+from .utils import Option, Dependency, generic_issubclass
 
 if sys.version_info >= (3, 10):
     from typing import Annotated, get_args, get_origin
@@ -119,12 +119,6 @@
 @dataclass
 class SQLDependsInner:
     dependency: ExecutableReturnsRows
-
-    if sys.version_info >= (3, 10):
-        from dataclasses import KW_ONLY
-
-        _: KW_ONLY
-
     use_cache: bool = True
     validate: bool | FieldInfo = False
 
@@ -135,7 +129,7 @@ def SQLDepends(
     use_cache: bool = True,
     validate: bool | FieldInfo = False,
 ) -> Any:
-    return SQLDependsInner(dependency, use_cache=use_cache, validate=validate)
+    return SQLDependsInner(dependency, use_cache, validate)
 
 
 class ORMParam(DependParam):
@@ -164,26 +158,35 @@ def _check_param(
             models = (models,)
 
         if depends_inner is not None:
-            dependency = compile_dependency(depends_inner.dependency, option)
+            statement = depends_inner.dependency
         elif all(map(isclass, models)) and all(
             map(issubclass, cast(Tuple[type, ...], models), repeat(Model))
         ):
             models = cast(Tuple[Type[Model], ...], models)
-            dependency = compile_dependency(
-                select(*models).where(
-                    *(
-                        getattr(model, name) == param.default
-                        for model in models
-                        for name, param in model.__signature__.parameters.items()
-                    )
-                ),
-                option,
+            # NOTE: statement is generated (see below)
+            statement = select(*models).where(
+                *(
+                    getattr(model, name) == param.default
+                    for model in models
+                    for name, param in model.__signature__.parameters.items()
+                )
             )
         else:
             return
 
-        return super()._check_param(param.replace(default=dependency), allow_types)
+        return super()._check_param(
+            param.replace(
+                default=Depends(
+                    Dependency(statement, option),
+                    use_cache=(
+                        depends_inner.use_cache if depends_inner else False
+                    ),  # NOTE: default use_cache=False as it is impossible to reuse a generated statement (see above)
+                    validate=depends_inner.validate if depends_inner else False,
+                )
+            ),
+            allow_types,
+        )
 
     @classmethod
-    def _check_parameterless(cls, *_) -> Param | None:
+    def _check_parameterless(cls, *_) -> None:
         return

nonebot_plugin_orm/utils.py

@@ -11,9 +11,9 @@
 from operator import methodcaller
 from typing_extensions import Annotated
 from dataclasses import field, dataclass
-from typing import Any, TypeVar, Coroutine
 from inspect import Parameter, Signature, isclass
 from collections.abc import Callable, Iterable, Generator
+from typing import TYPE_CHECKING, Any, TypeVar, Coroutine
 from importlib.metadata import Distribution, PackageNotFoundError, distribution
 
 import click
@@ -36,6 +36,10 @@
     from typing_extensions import ParamSpec, get_args, get_origin
 
 
+if TYPE_CHECKING:
+    from . import async_scoped_session
+
+
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 
@@ -73,60 +77,67 @@ def write(self, buffer: str):
         while frame and frame.f_code.co_name != "print_stdout":
             frame = frame.f_back
             depth += 1
-        depth += 1
 
         for line in buffer.rstrip().splitlines():
-            logger.opt(depth=depth).log(self._level, line.rstrip())
+            logger.opt(depth=depth + 1).log(self._level, line.rstrip())
 
     def flush(self):
         pass
 
 
-@dataclass
+@dataclass(unsafe_hash=True)
 class Option:
     stream: bool = True
     scalars: bool = False
     result: methodcaller | None = None
-    calls: list[methodcaller] = field(default_factory=list)
+    calls: tuple[methodcaller] = field(default_factory=tuple)
 
 
-def compile_dependency(statement: ExecutableReturnsRows, option: Option) -> Any:
-    from . import async_scoped_session
+@dataclass
+class Dependency:
+    __signature__: Signature = field(init=False)
+
+    statement: ExecutableReturnsRows
+    option: Option
+
+    def __post_init__(self) -> None:
+        from . import async_scoped_session
+
+        self.__signature__ = Signature(
+            [
+                Parameter(
+                    "_session", Parameter.KEYWORD_ONLY, annotation=async_scoped_session
+                ),
+                *(
+                    Parameter(name, Parameter.KEYWORD_ONLY, default=depends)
+                    for name, depends in self.statement.compile().params.items()
+                    if isinstance(depends, DependsInner)
+                ),
+            ]
+        )
 
-    async def __dependency(*, __session: async_scoped_session, **params: Any):
-        if option.stream:
-            result = await __session.stream(statement, params)
+    async def __call__(self, *, _session: async_scoped_session, **params: Any) -> Any:
+        if self.option.stream:
+            result = await _session.stream(self.statement, params)
         else:
-            result = await __session.execute(statement, params)
+            result = await _session.execute(self.statement, params)
 
-        for call in option.calls:
+        for call in self.option.calls:
             result = call(result)
 
-        if option.scalars:
+        if self.option.scalars:
             result = result.scalars()
 
-        if call := option.result:
+        if call := self.option.result:
             result = call(result)
 
-            if option.stream:
+            if self.option.stream:
                 result = await result
 
         return result
 
-    __dependency.__signature__ = Signature(
-        [
-            Parameter(
-                "__session", Parameter.KEYWORD_ONLY, annotation=async_scoped_session
-            ),
-            *(
-                Parameter(name, Parameter.KEYWORD_ONLY, default=depends)
-                for name, depends in statement.compile().params.items()
-                if isinstance(depends, DependsInner)
-            ),
-        ]
-    )
-
-    return Depends(__dependency)
+    def __hash__(self) -> int:
+        return hash((self.statement, self.option))
 
 
 def generic_issubclass(scls: Any, cls: Any) -> Any: