feat: support multi engine

Author: RF-Tar-Railt

src/graia/amnesia/builtins/sqla/__init__.py

@@ -1,11 +1,38 @@
 from __future__ import annotations
 
+import logging
+
+from loguru import logger
+
 try:
-    import sqlalchemy as sa
+    from sqlalchemy.log import Identified, _qual_logger_name_for_cls
 except ImportError:
     raise ImportError(
         "dependency 'sqlalchemy' is required for sqlalchemy service\nplease install it or install 'graia-amnesia[sqla]'"
     )
 
 from .model import Base as Base
 from .service import SqlalchemyService as SqlalchemyService
+from .utils import LoguruHandler, get_subclasses
+
+
+def pacth_logger(log_level: str | int = "INFO", sqlalchemy_echo: bool = False) -> None:
+    handler = LoguruHandler()
+    logging.getLogger("sqlalchemy").addHandler(handler)
+
+    if isinstance(log_level, str):
+        log_level = logger.level(log_level).no
+
+    echo_log_level = log_level if sqlalchemy_echo else logging.WARNING
+
+    levels = {
+        "alembic": log_level,
+        "sqlalchemy": log_level,
+        **{
+            _qual_logger_name_for_cls(cls): echo_log_level
+            for cls in set(get_subclasses(Identified))
+        },
+    }
+
+    for name, level in levels.items():
+        logging.getLogger(name).setLevel(level)

src/graia/amnesia/builtins/sqla/manager.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, cast
+from typing import Any, Callable, TYPE_CHECKING, ClassVar
 
 from sqlalchemy import MetaData
 from sqlalchemy.ext.asyncio import AsyncAttrs
@@ -20,7 +20,7 @@ def _setup_bind(cls: type["Base"]) -> None:
     if bind_key is None:
         bind_key = ""
 
-    cast(Table, cls.__table__).info["bind_key"] = bind_key
+    cls.__table__.info["bind_key"] = bind_key
 
 
 _callbacks = []
@@ -46,6 +46,9 @@ class Base(AsyncAttrs, DeclarativeBase):
     __abstract__ = True
     metadata = MetaData(naming_convention=_NAMING_CONVENTION)
 
+    if TYPE_CHECKING:
+        __table__: ClassVar[Table]  # type: ignore
+
     def __init_subclass__(cls, **kwargs):
         for callback in _callbacks:
             callback(cls)

src/graia/amnesia/builtins/sqla/model.py

@@ -1,36 +1,48 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Literal, ClassVar
+from typing import Literal, ClassVar, Any, TypeVar, cast
 
 from launart import Launart, Service
 from loguru import logger
-from sqlalchemy.engine.result import Result
-from sqlalchemy.engine.url import URL
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from sqlalchemy import Table
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.ext.asyncio.engine import AsyncEngine
 from sqlalchemy.sql.base import Executable
 from sqlalchemy.sql.selectable import TypedReturnsRows
+from sqlalchemy.engine.result import Result
+from sqlalchemy.engine.url import URL
 from sqlalchemy.orm import DeclarativeBase
 
-from .manager import DatabaseManager, T_Row
 from .model import Base
 from .types import EngineOptions
+from .utils import get_subclasses
+
+T_Row = TypeVar("T_Row", bound=DeclarativeBase)
 
 
 class SqlalchemyService(Service):
     id: str = "database/sqlalchemy"
-    db: DatabaseManager
-    get_session: async_sessionmaker[AsyncSession]
     base_class: ClassVar[type[DeclarativeBase]] = Base
+    engines: dict[str, AsyncEngine]
+    session_factory: async_sessionmaker[AsyncSession]
 
     def __init__(
         self,
         url: str | URL,
         engine_options: EngineOptions | None = None,
+        session_options: dict[str, Any] | None = None,
+        binds: dict[str, str | URL] | None = None,
         create_table_at: Literal["preparing", "prepared", "blocking"] = "preparing"
     ) -> None:
-        self.db = DatabaseManager(url, engine_options)
+        if engine_options is None:
+            engine_options = {"echo": "debug", "pool_pre_ping": True}
+        self.engines[""] = create_async_engine(url, **engine_options)
+        for key, bind_url in (binds or {}).items():
+            self.engines[key] = create_async_engine(bind_url, **engine_options)
         self.create_table_at = create_table_at
+        self.session_options = session_options or {"expire_on_commit": False}
         super().__init__()
 
     @property
@@ -41,51 +53,96 @@ def required(self) -> set[str]:
     def stages(self) -> set[Literal["preparing", "blocking", "cleanup"]]:
         return {"preparing", "blocking", "cleanup"}
 
+    async def initialize(self):
+        binds = {}
+
+        for model in set(get_subclasses(self.base_class)):
+            table: Table | None = getattr(model, "__table__", None)
+
+            if table is None or (bind_key := table.info.get("bind_key")) is None:
+                continue
+
+            binds[model] = self.engines.get(bind_key, self.engines[""])
+
+        self.session_factory = async_sessionmaker(self.engines[""], binds=binds, **self.session_options)
+        return binds
+
+    def get_session(self, **local_kw):
+        return self.session_factory(**local_kw)
+
     async def launch(self, manager: Launart):
+        binds: dict[type[Base], AsyncEngine] = {}
+
         async with self.stage("preparing"):
             logger.info("Initializing database...")
-            await self.db.initialize()
-            self.get_session = self.db.session_factory
-            logger.success("Database initialized!")
             if self.create_table_at == "preparing":
-                async with self.db.engine.begin() as conn:
-                    await conn.run_sync(self.base_class.metadata.create_all)
-                    logger.success("Database tables created!")
+                binds = await self.initialize()
+                logger.success("Database initialized!")
+                for model, engine in binds.items():
+                    async with engine.begin() as conn:
+                        await conn.run_sync(model.__table__.create, checkfirst=True)
+                logger.success("Database tables created!")
 
+        if self.create_table_at != "preparing":
+            binds = await self.initialize()
+            logger.success("Database initialized!")
         if self.create_table_at == "prepared":
-            async with self.db.engine.begin() as conn:
-                await conn.run_sync(self.base_class.metadata.create_all)
-                logger.success("Database tables created!")
+            for model, engine in binds.items():
+                async with engine.begin() as conn:
+                    await conn.run_sync(model.__table__.create, checkfirst=True)
+            logger.success("Database tables created!")
 
         async with self.stage("blocking"):
             if self.create_table_at == "blocking":
-                async with self.db.engine.begin() as conn:
-                    await conn.run_sync(self.base_class.metadata.create_all)
-                    logger.success("Database tables created!")
+                for model, engine in binds.items():
+                    async with engine.begin() as conn:
+                        await conn.run_sync(model.__table__.create, checkfirst=True)
+                logger.success("Database tables created!")
             await manager.status.wait_for_sigexit()
         async with self.stage("cleanup"):
-            await self.db.stop()
+            for engine in self.engines.values():
+                await engine.dispose(close=True)
 
     async def execute(self, sql: Executable) -> Result:
-        return await self.db.execute(sql)
+        """执行 SQL 命令"""
+        async with self.get_session() as session:
+            return await session.execute(sql)
 
     async def select_all(self, sql: TypedReturnsRows[tuple[T_Row]]) -> Sequence[T_Row]:
-        return await self.db.select_all(sql)
+        async with self.get_session() as session:
+            result = await session.scalars(sql)
+        return result.all()
 
     async def select_first(self, sql: TypedReturnsRows[tuple[T_Row]]) -> T_Row | None:
-        return await self.db.select_first(sql)
+        async with self.get_session() as session:
+            result = await session.scalars(sql)
+        return cast("T_Row | None", result.first())
 
-    async def add(self, row: Base):
-        return await self.db.add(row)
+    async def add(self, row: Base) -> None:
+        async with self.get_session() as session:
+            session.add(row)
+            await session.commit()
+            await session.refresh(row)
 
     async def add_many(self, rows: Sequence[Base]):
-        return await self.db.add_many(rows)
+        async with self.get_session() as session:
+            session.add_all(rows)
+            await session.commit()
+            for row in rows:
+                await session.refresh(row)
 
     async def update_or_add(self, row: Base):
-        return await self.db.update_or_add(row)
+        async with self.get_session() as session:
+            await session.merge(row)
+            await session.commit()
+            await session.refresh(row)
 
     async def delete_exist(self, row: Base):
-        return await self.db.delete_exist(row)
+        async with self.get_session() as session:
+            await session.delete(row)
 
     async def delete_many_exist(self, rows: Sequence[Base]):
-        return await self.db.delete_many_exist(rows)
+        async with self.get_session() as session:
+            for row in rows:
+                await session.delete(row)
+

src/graia/amnesia/builtins/sqla/service.py

@@ -0,0 +1,29 @@
+import logging
+import sys
+
+from loguru import logger
+
+
+class LoguruHandler(logging.Handler):
+    def emit(self, record: logging.LogRecord):
+        try:
+            level = logger.level(record.levelname).name
+            if record.levelno <= logging.INFO:
+                level = {"DEBUG": "TRACE", "INFO": "DEBUG"}.get(level, level)
+        except ValueError:
+            level = record.levelno
+
+        frame, depth = sys._getframe(6), 6
+        while frame and frame.f_code.co_filename == logging.__file__:
+            frame = frame.f_back
+            depth += 1
+
+        logger.opt(depth=depth, exception=record.exc_info).log(
+            level, record.getMessage()
+        )
+
+
+def get_subclasses(cls):
+    yield from cls.__subclasses__()
+    for subclass in cls.__subclasses__():
+        yield from get_subclasses(subclass)