feat: allow callbacks when SQLA modelbase init subclass

Author: RF-Tar-Railt

pdm.lock

@@ -3,12 +3,14 @@
 from typing import cast
 
 from launart import Launart, Service
+from launart.status import Phase
 
 try:
     from aiohttp import ClientSession, ClientTimeout
 except ImportError:
     raise ImportError(
-        "dependency 'aiohttp' is required for aiohttp client service\nplease install it or install 'graia-amnesia[aiohttp]'"
+        "dependency 'aiohttp' is required for aiohttp client service\n"
+        "please install it or install 'graia-amnesia[aiohttp]'"
     )
 
 
@@ -21,14 +23,14 @@ def __init__(self, session: ClientSession | None = None) -> None:
         super().__init__()
 
     @property
-    def stages(self):
+    def stages(self) -> set[Phase]:
         return {"preparing", "cleanup"}
 
     @property
     def required(self):
         return set()
 
-    async def launch(self, _: Launart):
+    async def launch(self, manager: Launart):
         async with self.stage("preparing"):
             if self.session is None:
                 self.session = ClientSession(timeout=ClientTimeout(total=None))

src/graia/amnesia/builtins/aiohttp.py

@@ -4,6 +4,7 @@
 import logging
 
 from launart import Launart, Service
+from launart.status import Phase
 from launart.utilles import any_completed
 from loguru import logger
 
@@ -85,7 +86,7 @@ def required(self):
         return set()
 
     @property
-    def stages(self):
+    def stages(self) -> set[Phase]:
         return {"preparing", "blocking", "cleanup"}
 
     async def launch(self, manager: Launart) -> None:

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

@@ -3,6 +3,7 @@
 from typing import cast
 
 from launart import Launart, Service
+from launart.status import Phase
 
 try:
     from httpx import AsyncClient, Timeout
@@ -21,14 +22,14 @@ def __init__(self, session: AsyncClient | None = None) -> None:
         super().__init__()
 
     @property
-    def stages(self):
+    def stages(self) -> set[Phase]:
         return {"preparing", "cleanup"}
 
     @property
     def required(self):
         return set()
 
-    async def launch(self, _: Launart):
+    async def launch(self, manager: Launart):
         async with self.stage("preparing"):
             if self.session is None:
                 self.session = AsyncClient(timeout=Timeout())

src/graia/amnesia/builtins/httpx.py

@@ -7,6 +7,7 @@
 from typing import Any
 
 from launart import Launart, Service
+from launart.status import Phase
 
 
 class Memcache:
@@ -75,7 +76,7 @@ def required(self):
         return set()
 
     @property
-    def stages(self):
+    def stages(self) -> set[Phase]:
         return {"blocking"}
 
     @property

src/graia/amnesia/builtins/memcache.py

@@ -1,6 +1,9 @@
+from typing import Any, Callable, cast
+
 from sqlalchemy import MetaData
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.schema import Table
 
 _NAMING_CONVENTION = {
     "ix": "ix_%(column_0_label)s",
@@ -11,6 +14,54 @@
 }
 
 
+def _setup_bind(cls: type["Base"]) -> None:
+    bind_key: str | None = getattr(cls, "__bind_key__", None)
+
+    if bind_key is None:
+        bind_key = ""
+
+    cast(Table, cls.__table__).info["bind_key"] = bind_key
+
+
+_callbacks = []
+
+
+def register_callback(callback: Callable[[type["Base"]], Any]) -> None:
+    """
+    Register a callback to be called when a new Base subclass is created.
+    The callback should accept a single argument, which is the subclass itself.
+    """
+    _callbacks.append(callback)
+
+
+def remove_callback(callback: Callable[[type["Base"]], Any]) -> None:
+    """
+    Remove a previously registered callback.
+    """
+    if callback in _callbacks:
+        _callbacks.remove(callback)
+
+
 class Base(AsyncAttrs, DeclarativeBase):
     __abstract__ = True
     metadata = MetaData(naming_convention=_NAMING_CONVENTION)
+
+    def __init_subclass__(cls, **kwargs):
+        for callback in _callbacks:
+            callback(cls)
+
+        if not hasattr(cls, "__tablename__") and "tablename" in kwargs:
+            cls.__tablename__ = kwargs["tablename"]
+        if not hasattr(cls, "__table_args__") and "table_args" in kwargs:
+            cls.__table_args__ = kwargs["table_args"]
+        if not hasattr(cls, "__mapper__") and "mapper" in kwargs:
+            cls.__mapper__ = kwargs["mapper"]
+        if not hasattr(cls, "__mapper_args__") and "mapper_args" in kwargs:
+            cls.__mapping_args__ = kwargs["mapper_args"]
+
+        super().__init_subclass__(**kwargs)
+
+        if not hasattr(cls, "__table__"):
+            return
+
+        _setup_bind(cls)

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Literal
+from typing import Literal, ClassVar
 
 from launart import Launart, Service
 from loguru import logger
@@ -10,6 +10,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 from sqlalchemy.sql.base import Executable
 from sqlalchemy.sql.selectable import TypedReturnsRows
+from sqlalchemy.orm import DeclarativeBase
 
 from .manager import DatabaseManager, T_Row
 from .model import Base
@@ -20,6 +21,7 @@ class SqlalchemyService(Service):
     id: str = "database/sqlalchemy"
     db: DatabaseManager
     get_session: async_sessionmaker[AsyncSession]
+    base_class: ClassVar[type[DeclarativeBase]] = Base
 
     def __init__(
         self,
@@ -47,18 +49,18 @@ async def launch(self, manager: Launart):
             logger.success("Database initialized!")
             if self.create_table_at == "preparing":
                 async with self.db.engine.begin() as conn:
-                    await conn.run_sync(Base.metadata.create_all)
+                    await conn.run_sync(self.base_class.metadata.create_all)
                     logger.success("Database tables created!")
 
         if self.create_table_at == "prepared":
             async with self.db.engine.begin() as conn:
-                await conn.run_sync(Base.metadata.create_all)
+                await conn.run_sync(self.base_class.metadata.create_all)
                 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(Base.metadata.create_all)
+                    await conn.run_sync(self.base_class.metadata.create_all)
                     logger.success("Database tables created!")
             await manager.status.wait_for_sigexit()
         async with self.stage("cleanup"):