entari-plugin-database 0.1.0__tar.gz

entari_plugin_database-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 ARCLET
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

entari_plugin_database-0.1.0/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: entari-plugin-database
+Version: 0.1.0
+Summary: Entari plugin for SQLAlchemy ORM
+Author-Email: RF-Tar-Railt <rf_tar_railt@qq.com>
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: sqlalchemy>=2.0.42
+Requires-Dist: aiosqlite>=0.21.0
+Requires-Dist: graia-amnesia>=0.10.1
+Requires-Dist: tarina<0.8.0,>=0.7.1
+Requires-Dist: arclet-entari>=0.15.0
+Description-Content-Type: text/markdown
+
+# entari-plugin-database
+Entari plugin for SQLAlchemy ORM

entari_plugin_database-0.1.0/README.md

@@ -0,0 +1,2 @@
+# entari-plugin-database
+Entari plugin for SQLAlchemy ORM

entari_plugin_database-0.1.0/pyproject.toml

@@ -0,0 +1,33 @@
+[project]
+name = "entari-plugin-database"
+version = "0.1.0"
+description = "Entari plugin for SQLAlchemy ORM"
+authors = [
+    { name = "RF-Tar-Railt", email = "rf_tar_railt@qq.com" },
+]
+dependencies = [
+    "sqlalchemy>=2.0.42",
+    "aiosqlite>=0.21.0",
+    "graia-amnesia>=0.10.1",
+    "tarina<0.8.0,>=0.7.1",
+    "arclet-entari>=0.15.0",
+]
+requires-python = ">=3.9"
+readme = "README.md"
+
+[project.license]
+text = "MIT"
+
+[build-system]
+requires = [
+    "pdm-backend",
+]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true
+
+[tool.pdm.dev-dependencies]
+dev = [
+    "arclet-entari[full]>=0.15.0",
+]

entari_plugin_database-0.1.0/src/entari_plugin_database/__init__.py

@@ -0,0 +1,154 @@
+from dataclasses import field
+from typing import Optional, Any, Literal
+from typing import Union
+
+from sqlalchemy.ext.asyncio import create_async_engine
+from arclet.letoderea.scope import global_providers, global_propagators
+from arclet.entari import BasicConfModel, plugin, logger
+from arclet.entari.config import config_model_validate
+from arclet.entari.event.config import ConfigReload
+from graia.amnesia.builtins.sqla import SqlalchemyService
+from graia.amnesia.builtins.sqla.model import register_callback, remove_callback
+from graia.amnesia.builtins.sqla.model import Base as Base
+from graia.amnesia.builtins.sqla.types import EngineOptions
+from sqlalchemy.engine.url import URL
+from sqlalchemy.ext import asyncio as sa_async
+from sqlalchemy.orm import Mapped as Mapped
+from sqlalchemy.orm import mapped_column as mapped_column
+
+from .param import db_supplier, sess_provider, orm_factory
+from .param import SQLDepends as SQLDepends
+
+
+class UrlInfo(BasicConfModel):
+    type: str = "sqlite"
+    """数据库类型，默认为 sqlite"""
+    name: str = "data.db"
+    """数据库名称/文件路径"""
+    driver: str = "aiosqlite"
+    """数据库驱动，默认为 aiosqlite；其他类型的数据库驱动参考 SQLAlchemy 文档"""
+    host: Optional[str] = None
+    """数据库主机地址。如果是 SQLite 数据库，此项可不填。"""
+    port: Optional[int] = None
+    """数据库端口号。如果是 SQLite 数据库，此项可不填。"""
+    username: Optional[str] = None
+    """数据库用户名。如果是 SQLite 数据库，此项可不填。"""
+    password: Optional[str] = None
+    """数据库密码。如果是 SQLite 数据库，此项可不填。"""
+    query: dict[str, Union[list[str], str]] = field(default_factory=dict)
+    """数据库连接参数，默认为空字典。可以传入如 `{"timeout": "30"}` 的参数。"""
+
+    @property
+    def url(self) -> URL:
+        if self.type == "sqlite":
+            return URL.create(f"{self.type}+{self.driver}", database=self.name, query=self.query)
+        return URL.create(
+            f"{self.type}+{self.driver}", self.username, self.password, self.host, self.port, self.name, self.query
+        )
+
+
+class Config(UrlInfo):
+    options: EngineOptions = field(default_factory=lambda: {"echo": None, "pool_pre_ping": True})
+    """数据库连接选项，默认为 `{"echo": None, "pool_pre_ping": True}`"""
+    session_options: Union[dict[str, Any], None] = field(default=None)
+    """数据库会话选项，默认为 None。可以传入如 `{"expire_on_commit": False}` 的字典。"""
+    binds: dict[str, UrlInfo]  = field(default_factory=dict)
+    """数据库绑定配置，默认为 None。可以传入如 `{"bind1": UrlInfo(...), "bind2": UrlInfo(...)}` 的字典。"""
+    create_table_at: Literal['preparing', 'prepared', 'blocking'] = "preparing"
+    """在指定阶段创建数据库表，默认为 'preparing'。可选值为 'preparing', 'prepared', 'blocking'。"""
+
+
+plugin.declare_static()
+plugin.metadata(
+    "Database 服务",
+    [{"name": "RF-Tar-Railt", "email": "rf_tar_railt@qq.com"}],
+    "0.1.0",
+    description="基于 SQLAlchemy 的数据库服务插件",
+    urls={
+        "homepage": "https://github.com/ArcletProject/entari-plugin-database",
+    },
+    config=Config,
+)
+plugin.collect_disposes(
+    lambda: global_propagators.remove(db_supplier),
+    lambda: global_providers.remove(sess_provider),
+    lambda: global_providers.remove(orm_factory),
+)
+
+log = logger.log.wrapper("[Database]")
+_config = plugin.get_config(Config)
+
+try:
+    plugin.add_service(
+        service := SqlalchemyService(
+            _config.url,
+            _config.options,
+            _config.session_options,
+            {key: value.url for key, value in _config.binds.items()},
+            _config.create_table_at
+        )
+    )
+except Exception as e:
+    raise RuntimeError("Failed to initialize SqlalchemyService. Please check your database configuration.") from e
+
+
+@plugin.listen(ConfigReload)
+async def reload_config(event: ConfigReload, serv: SqlalchemyService):
+    if event.scope != "plugin":
+        return None
+    if event.key not in ("database", "entari_plugin_database"):
+        return None
+    new_conf = config_model_validate(Config, event.value)
+    for engine in serv.engines.values():
+        await engine.dispose(close=True)
+    engine_options = {"echo": "debug", "pool_pre_ping": True}
+    serv.engines = {"": create_async_engine(new_conf.url, **(new_conf.options or engine_options))}
+    for key, bind in (new_conf.binds or {}).items():
+        serv.engines[key] = create_async_engine(bind.url, **(new_conf.options or engine_options))
+    serv.create_table_at = new_conf.create_table_at
+    serv.session_options = new_conf.session_options or {"expire_on_commit": False}
+
+    binds = await serv.initialize()
+    log.success("Database initialized!")
+    for key, models in binds.items():
+        async with serv.engines[key].begin() as conn:
+            await conn.run_sync(
+                serv.base_class.metadata.create_all, tables=[m.__table__ for m in models], checkfirst=True
+            )
+    log.success("Database tables created!")
+    return True
+
+
+def _setup_tablename(cls: type[Base], kwargs: dict):
+    if "tablename" in kwargs:
+        return
+    for attr in ("__tablename__", "__table__"):
+        if getattr(cls, attr, None):
+            return
+
+    cls.__tablename__ = cls.__name__.lower()
+
+    if plg := plugin.get_plugin(3):
+        cls.__tablename__ = f"{plg.id.replace('-', '_')}_{cls.__tablename__}"
+
+
+register_callback(_setup_tablename)
+plugin.collect_disposes(lambda: remove_callback(_setup_tablename))
+
+
+BaseOrm = Base
+AsyncSession = sa_async.AsyncSession
+get_session = service.get_session
+
+
+__all__ = [
+    "AsyncSession",
+    "Base",
+    "BaseOrm",
+    "Mapped",
+    "mapped_column",
+    "service",
+    "SQLDepends",
+    "get_session",
+    "SqlalchemyService"
+]

entari_plugin_database-0.1.0/src/entari_plugin_database/param.py

@@ -0,0 +1,244 @@
+from dataclasses import dataclass, field
+from inspect import Signature, Parameter
+from operator import methodcaller
+
+from collections.abc import Iterator, Sequence, AsyncIterator
+from typing import Optional
+
+from sqlalchemy import Row, Result, ScalarResult, select
+from sqlalchemy.ext.asyncio import AsyncResult, AsyncScalarResult
+from sqlalchemy.sql.selectable import ExecutableReturnsRows
+from tarina import generic_issubclass
+from tarina.generic import origin_is_union, isclass
+from typing_extensions import Any, cast, get_args, get_origin
+
+from creart import it
+from launart import Launart
+from graia.amnesia.builtins.sqla import SqlalchemyService, Base
+
+from arclet.letoderea import Propagator, Contexts, STACK, Provider, ProviderFactory, Param, Depend, Subscriber
+from arclet.letoderea.ref import Deref, generate
+from arclet.letoderea.scope import global_propagators, global_providers
+from sqlalchemy.ext import asyncio as sa_async
+
+
+class DatabasePropagator(Propagator):
+    async def supply(self, ctx: Contexts, serv: Optional[SqlalchemyService] = None):
+        if serv is None:
+            return
+        session = serv.get_session()
+        stack = ctx[STACK]
+        session = await stack.enter_async_context(session)
+        return {"$db_session": session}
+
+    def compose(self):
+        yield self.supply, True, 20
+
+
+class SessionProvider(Provider[sa_async.AsyncSession]):
+    priority = 10
+
+    async def __call__(self, context: Contexts):
+        if "$db_session" in context:
+            return context["$db_session"]
+        try:
+            db = it(Launart).get_component(SqlalchemyService)
+            stack = context[STACK]
+            sess = await stack.enter_async_context(db.get_session())
+            context["$db_session"] = sess
+            return sess
+        except ValueError:
+            return
+
+
+@dataclass(unsafe_hash=True)
+class Option:
+    stream: bool = True
+    scalars: bool = False
+    calls: tuple[methodcaller, ...] = field(default_factory=tuple)
+    result: methodcaller | None = None
+
+
+PATTERNS = {
+    AsyncIterator[Sequence[Row[tuple[Any, ...]]]]: Option(
+        True,
+        False,
+        (methodcaller("partitions"),),
+    ),
+    AsyncIterator[Sequence[tuple[Any, ...]]]: Option(
+        True,
+        False,
+        (methodcaller("partitions"),),
+    ),
+    AsyncIterator[Sequence[Any]]: Option(
+        True,
+        True,
+        (methodcaller("partitions"),),
+    ),
+    Iterator[Sequence[Row[tuple[Any, ...]]]]: Option(
+        False,
+        False,
+        (methodcaller("partitions"),),
+    ),
+    Iterator[Sequence[tuple[Any, ...]]]: Option(
+        False,
+        False,
+        (methodcaller("partitions"),),
+    ),
+    Iterator[Sequence[Any]]: Option(
+        False,
+        True,
+        (methodcaller("partitions"),),
+    ),
+    AsyncResult[tuple[Any, ...]]: Option(
+        True,
+        False,
+    ),
+    AsyncScalarResult[Any]: Option(
+        True,
+        True,
+    ),
+    Result[tuple[Any, ...]]: Option(
+        False,
+        False,
+    ),
+    ScalarResult[Any]: Option(
+        False,
+        True,
+    ),
+    AsyncIterator[Row[tuple[Any, ...]]]: Option(
+        True,
+        False,
+    ),
+    Iterator[Row[tuple[Any, ...]]]: Option(
+        False,
+        False,
+    ),
+    Sequence[Row[tuple[Any, ...]]]: Option(
+        True,
+        False,
+        (),
+        methodcaller("all"),
+    ),
+    Sequence[tuple[Any, ...]]: Option(
+        True,
+        False,
+        (),
+        methodcaller("all"),
+    ),
+    Sequence[Any]: Option(
+        True,
+        True,
+        (),
+        methodcaller("all"),
+    ),
+    tuple[Any, ...]: Option(
+        True,
+        False,
+        (),
+        methodcaller("one_or_none"),
+    ),
+    Any: Option(
+        True,
+        True,
+        (),
+        methodcaller("one_or_none"),
+    ),
+}
+
+
+class SQLDepend(Depend):
+    def __init__(self, statement: ExecutableReturnsRows, option: Option = Option(), cache: bool = False):
+        super().__init__(lambda : None, cache)
+        self.statement = statement
+        self.option = option
+
+        async def target(session: sa_async.AsyncSession,  **params):
+            if self.option.stream:
+                result = await session.stream(self.statement, params)
+            else:
+                result = await session.execute(self.statement, params)
+            if self.option.scalars:
+                result = result.scalars()
+            for call in self.option.calls:
+                result = call(result)
+            if call := self.option.result:
+                result = call(result)
+                if self.option.stream:
+                    result = await result
+            return result
+
+        parameters = [Parameter("session", Parameter.KEYWORD_ONLY, annotation=sa_async.AsyncSession)]
+        for name, depends in self.statement.compile().params.items():
+            if isinstance(depends, Depend):
+                parameters.append(Parameter(name, Parameter.KEYWORD_ONLY, default=depends))
+            elif isinstance(depends, Deref):
+                parameters.append(Parameter(name, Parameter.KEYWORD_ONLY, default=Depend(generate(depends))))
+        target.__signature__ = Signature(parameters)  # type: ignore
+        self.target = target
+
+
+def SQLDepends(statement: ExecutableReturnsRows, option: Option = Option(), cache: bool = False) -> Any:
+    return SQLDepend(statement, option, cache)
+
+
+class ORMProviderFactory(ProviderFactory):
+    priority = 10
+
+    class _ModelProvider(Provider[Any]):
+        def __init__(self, statement: ExecutableReturnsRows, option: Option):
+            super().__init__()
+            self.statement = statement
+            self.option = option
+
+        async def __call__(self, context: Contexts):
+            if "$db_session" not in context:
+                return
+            sess: sa_async.AsyncSession = context["$db_session"]
+            if self.option.stream:
+                result = await sess.stream(self.statement)
+            else:
+                result = await sess.execute(self.statement)
+            if self.option.scalars:
+                result = result.scalars()
+            for call in self.option.calls:
+                result = call(result)
+            if call := self.option.result:
+                result = call(result)
+                if self.option.stream:
+                    result = await result
+            return result
+
+    def validate(self, param: Param):
+        if isinstance(param.default, SQLDepend):
+            return
+        for pattern, option in PATTERNS.items():
+            if models := cast("list[Any]", generic_issubclass(pattern, param.annotation, list_=True)):
+                break
+        else:
+            models, option = [], Option()
+
+        for index, model in enumerate(models):
+            if origin_is_union(get_origin(model)):
+                models[index] = next(
+                    (
+                        arg
+                        for arg in get_args(model)
+                        if isclass(arg) and issubclass(arg, Base)
+                    ),
+                    None,
+                )
+
+            if not (isclass(models[index]) and issubclass(models[index], Base)):
+                models = []
+                break
+        if not models:
+            return
+
+        statement = select(*models)
+        return self._ModelProvider(statement, option)
+
+
+global_propagators.append(db_supplier := DatabasePropagator())
+global_providers.append(sess_provider := SessionProvider())
+global_providers.append(orm_factory := ORMProviderFactory())