✨ feat(alembic): sync

Author: ProgramRipper

nonebot_plugin_orm/__init__.py

@@ -4,17 +4,17 @@
 import logging
 from asyncio import gather
 from operator import methodcaller
-from functools import wraps, partial
 from typing import Any, AsyncGenerator
+from functools import wraps, partial, lru_cache
 
 from nonebot.params import Depends
-from nonebot.plugin import PluginMetadata
-import sqlalchemy.ext.asyncio as sa_asyncio
+import sqlalchemy.ext.asyncio as sa_async
 from sqlalchemy.util import greenlet_spawn
 from nonebot.matcher import current_matcher
-from nonebot import logger, require, get_driver
+from nonebot.plugin import Plugin, PluginMetadata
 from sqlalchemy import URL, Table, MetaData, make_url
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from nonebot import logger, require, get_driver, get_plugin_by_module_name
 
 from . import migrate
 from .config import Config
@@ -55,6 +55,7 @@
     "merge",
     "upgrade",
     "downgrade",
+    "sync",
     "show",
     "history",
     "heads",
@@ -77,7 +78,8 @@
 _binds: dict[type[Model], AsyncEngine]
 _engines: dict[str, AsyncEngine]
 _metadatas: dict[str, MetaData]
-_session_factory: sa_asyncio.async_sessionmaker[AsyncSession]
+_plugins: dict[str, Plugin]
+_session_factory: sa_async.async_sessionmaker[AsyncSession]
 
 _driver = get_driver()
 
@@ -90,34 +92,33 @@ async def init_orm() -> None:
         if plugin_config.alembic_startup_check:
             await greenlet_spawn(migrate.check, alembic_config)
         else:
-            logger.warning("跳过启动检查，直接创建所有表并标记数据库为最新修订版本")
-            await migrate._upgrade_fast(alembic_config)
-            await greenlet_spawn(migrate.stamp, alembic_config)
+            logger.warning("跳过启动检查，正在同步数据库模式...")
+            await greenlet_spawn(migrate.sync, alembic_config)
 
 
 def _init_orm():
     global _session_factory
 
     _init_engines()
     _init_table()
-    _session_factory = sa_asyncio.async_sessionmaker(
+    _session_factory = sa_async.async_sessionmaker(
         _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
     )
 
 
 @wraps(lambda: None)  # NOTE: for dependency injection
-def get_session(**local_kw: Any) -> sa_asyncio.AsyncSession:
+def get_session(**local_kw: Any) -> sa_async.AsyncSession:
     try:
         return _session_factory(**local_kw)
     except NameError:
         raise RuntimeError("nonebot-plugin-orm 未初始化") from None
 
 
-AsyncSession = Annotated[sa_asyncio.AsyncSession, Depends(get_session)]
+AsyncSession = Annotated[sa_async.AsyncSession, Depends(get_session)]
 
 
 async def get_scoped_session() -> (
-    AsyncGenerator[sa_asyncio.async_scoped_session[AsyncSession], None]
+    AsyncGenerator[sa_async.async_scoped_session[AsyncSession], None]
 ):
     try:
         scoped_session = async_scoped_session(
@@ -131,7 +132,7 @@ async def get_scoped_session() -> (
 
 
 async_scoped_session = Annotated[
-    sa_asyncio.async_scoped_session[AsyncSession], Depends(get_scoped_session)
+    sa_async.async_scoped_session[sa_async.AsyncSession], Depends(get_scoped_session)
 ]
 
 
@@ -186,20 +187,21 @@ def _init_engines():
 
 
 def _init_table():
-    global _binds, _metadatas
+    global _binds, _metadatas, _plugins
 
     _binds = {}
+    _plugins = {}
 
-    if len(_engines) == 1:  # NOTE: common case: only default engine
-        _metadatas = {"": Model.metadata}
-        return
-
+    _get_plugin_by_module_name = lru_cache(None)(get_plugin_by_module_name)
     for model in Model.__subclasses__():
         table: Table | None = getattr(model, "__table__", None)
 
         if table is None or (bind_key := table.info.get("bind_key")) is None:
             continue
 
+        if plugin := _get_plugin_by_module_name(model.__module__):
+            _plugins[plugin.name] = plugin
+
         _binds[model] = _engines.get(bind_key, _engines[""])
         table.to_metadata(_metadatas.get(bind_key, _metadatas[""]))
 

nonebot_plugin_orm/__main__.py

@@ -89,7 +89,7 @@ def init(*args, **kwargs) -> None:
 @orm.command()
 @click.option("-m", "--message", help="描述")
 @click.option("--sql", is_flag=True, help="以 SQL 的形式输出修订脚本")
-@click.option("--head", default="head", help="基准版本")
+@click.option("--head", help="基准版本")
 @click.option("--splice", is_flag=True, help="允许非头部修订作为基准版本")
 @click.option("--branch-label", help="分支标签")
 @click.option(
@@ -128,11 +128,6 @@ def merge(*args, **kwargs) -> Iterable[Script]:
 @click.argument("revision", required=False)
 @click.option("--sql", is_flag=True, help="以 SQL 的形式输出修订脚本")
 @click.option("--tag", help="一个任意的字符串, 可在自定义的 env.py 中使用")
-@click.option(
-    "--fast",
-    is_flag=True,
-    help="快速升级到最新版本，不运行修订脚本，直接创建当前的表（只应该在数据库为空、修订较多且只有表结构更改时使用）",
-)
 @click.pass_obj
 def upgrade(*args, **kwargs) -> None:
     """升级到较新版本。"""
@@ -149,6 +144,14 @@ def downgrade(*args, **kwargs) -> None:
     return migrate.downgrade(*args, **kwargs)
 
 
+@orm.command()
+@click.argument("revision", required=False)
+@click.pass_obj
+def sync(*args, **kwargs) -> None:
+    """同步数据库模式 (仅用于开发)。"""
+    return migrate.sync(*args, **kwargs)
+
+
 @orm.command()
 @click.argument("revs", nargs=-1)
 @click.pass_obj