🐛 fix(alembic): patch session during migration

ProgramRipper · ProgramRipper · commit 010ca017e68a · 2025-07-13T01:08:55.000+08:00
diff --git a/nonebot_plugin_orm/__init__.py b/nonebot_plugin_orm/__init__.py
@@ -2,11 +2,14 @@
 
 import logging
 from argparse import Namespace
-from functools import wraps, lru_cache
+from functools import cache, wraps
+from collections.abc import Generator
+from contextlib import contextmanager
 from typing_extensions import Any, Annotated
 
 import click
 from nonebot.rule import Rule
+from alembic.op import get_bind
 import sqlalchemy.ext.asyncio as sa_async
 from nonebot.permission import Permission
 from sqlalchemy.util import greenlet_spawn
@@ -16,8 +19,8 @@
 from nonebot.plugin import Plugin, PluginMetadata
 from sqlalchemy.log import Identified, _qual_logger_name_for_cls
 from nonebot.matcher import Matcher, current_event, current_matcher
-from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from nonebot import logger, require, get_driver, get_plugin_by_module_name
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncConnection, create_async_engine
 
 from . import migrate
 from .param import ORMParam
@@ -101,22 +104,6 @@ async def init_orm() -> None:
             await greenlet_spawn(migrate.sync, alembic_config)
 
 
-def _init_orm():
-    global _session_factory, _scoped_sessions
-
-    _init_engines()
-    _init_table()
-    _session_factory = sa_async.async_sessionmaker(
-        _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
-    )
-    _scoped_sessions = sa_async.async_scoped_session(
-        _session_factory,
-        lambda: (id(current_event.get(None)), current_matcher.get(None)),
-    )
-
-    run_postprocessor(_scoped_sessions.remove)
-
-
 def get_session(**local_kw: Any) -> sa_async.AsyncSession:
     try:
         return _session_factory(**local_kw)
@@ -149,6 +136,26 @@ def get_scoped_session() -> sa_async.async_scoped_session[sa_async.AsyncSession]
 ]
 
 
+@contextmanager
+def _patch_migrate_session() -> Generator[None, Any, None]:
+    global _session_factory, _scoped_sessions
+
+    session_factory, scoped_sessions = _session_factory, _scoped_sessions
+
+    _session_factory = sa_async.async_sessionmaker(
+        AsyncConnection._retrieve_proxy_for_target(get_bind()),
+        **plugin_config.sqlalchemy_session_options,
+    )
+    _scoped_sessions = sa_async.async_scoped_session(
+        _session_factory,
+        lambda: (id(current_event.get(None)), current_matcher.get(None)),
+    )
+
+    yield
+
+    _session_factory, _scoped_sessions = session_factory, scoped_sessions
+
+
 def _create_engine(engine: str | URL | dict[str, Any] | AsyncEngine) -> AsyncEngine:
     if isinstance(engine, AsyncEngine):
         return engine
@@ -200,7 +207,7 @@ def _init_table():
     _binds = {}
     _plugins = {}
 
-    _get_plugin_by_module_name = lru_cache(None)(get_plugin_by_module_name)
+    _get_plugin_by_module_name = cache(get_plugin_by_module_name)
     for model in set(get_subclasses(Model)):
         table: Table | None = getattr(model, "__table__", None)
 
@@ -214,6 +221,22 @@ def _init_table():
         table.to_metadata(_metadatas.get(bind_key, _metadatas[""]))
 
 
+def _init_orm():
+    global _session_factory, _scoped_sessions
+
+    _init_engines()
+    _init_table()
+    _session_factory = sa_async.async_sessionmaker(
+        _engines[""], binds=_binds, **plugin_config.sqlalchemy_session_options
+    )
+    _scoped_sessions = sa_async.async_scoped_session(
+        _session_factory,
+        lambda: (id(current_event.get(None)), current_matcher.get(None)),
+    )
+
+    run_postprocessor(_scoped_sessions.remove)
+
+
 def _init_logger():
     handler = LoguruHandler()
     logging.getLogger("alembic").addHandler(handler)
diff --git a/nonebot_plugin_orm/migrate.py b/nonebot_plugin_orm/migrate.py
@@ -21,11 +21,13 @@
 from alembic.config import Config
 from sqlalchemy.util import asbool
 from nonebot import logger, get_plugin
+from nonebot.matcher import current_matcher
 from sqlalchemy import MetaData, Connection
 from alembic.util.editor import open_in_editor
 from alembic.script import Script, ScriptDirectory
 from alembic.util.langhelpers import rev_id as _rev_id
 from alembic.operations.ops import UpgradeOps, DowngradeOps
+from sqlalchemy.ext.asyncio import AsyncConnection, async_sessionmaker
 from alembic.migration import StampStep, RevisionStep, MigrationContext
 from alembic.runtime.environment import EnvironmentContext, ProcessRevisionDirectiveFn
 from alembic.autogenerate.api import (
@@ -121,21 +123,17 @@ def __init__(
             stdout,
             cmd_opts,
             {
-                **{
-                    "script_location": script_location,
-                    "prepend_sys_path": ".",
-                    "revision_environment": "true",
-                    "version_path_separator": "os",
-                },
-                **config_args,
-            },
+                "script_location": script_location,
+                "prepend_sys_path": ".",
+                "revision_environment": "true",
+                "version_path_separator": "os",
+            }
+            | dict(config_args),
             {
-                **{
-                    "engines": _engines,
-                    "metadatas": _metadatas,
-                },
-                **attributes,
-            },
+                "engines": _engines,
+                "metadatas": _metadatas,
+            }
+            | attributes,
         )
 
         self._init_post_write_hooks()
@@ -646,7 +644,11 @@ def upgrade(
 
     @return_progressbar
     def upgrade(rev, _) -> Iterable[StampStep | RevisionStep]:
-        yield from script._upgrade_revs(revision, rev)
+        from . import _patch_migrate_session
+
+        with _patch_migrate_session():
+            yield from script._upgrade_revs(revision, rev)
+
         _move_run_scripts(config, script, revision)
 
     with EnvironmentContext(
@@ -691,7 +693,11 @@ def downgrade(
 
     @return_progressbar
     def downgrade(rev, _) -> Iterable[StampStep | RevisionStep]:
-        yield from script._downgrade_revs(revision, rev)
+        from . import _patch_migrate_session
+
+        with _patch_migrate_session():
+            yield from script._downgrade_revs(revision, rev)
+
         _move_run_scripts(config, script, revision)
 
     with EnvironmentContext(
diff --git a/nonebot_plugin_orm/templates/generic/env.py b/nonebot_plugin_orm/templates/generic/env.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import cast
+from typing import Any, cast
 
 from alembic import context
 from sqlalchemy import Connection
@@ -36,34 +36,26 @@ def run_migrations_offline() -> None:
     在这里调用 context.execute() 会将给定的字符串写入到脚本输出.
     """
 
-    context.configure(
-        **{
-            **dict(
-                url=engine.url,
-                dialect_opts={"paramstyle": "named"},
-                target_metadata=target_metadata,
-                literal_binds=True,
-            ),
-            **plugin_config.alembic_context,
-        }
-    )
+    opts: dict[str, Any] = {
+        "url": engine.url,
+        "dialect_opts": {"paramstyle": "named"},
+        "target_metadata": target_metadata,
+        "literal_binds": True,
+    } | plugin_config.alembic_context
+    context.configure(**opts)
 
     with context.begin_transaction():
         context.run_migrations()
 
 
 def do_run_migrations(connection: Connection) -> None:
-    context.configure(
-        **{
-            **dict(
-                connection=connection,
-                render_as_batch=True,
-                target_metadata=target_metadata,
-                include_object=no_drop_table,
-            ),
-            **plugin_config.alembic_context,
-        }
-    )
+    opts: dict[str, Any] = {
+        "connection": connection,
+        "render_as_batch": True,
+        "target_metadata": target_metadata,
+        "include_object": no_drop_table,
+    } | plugin_config.alembic_context
+    context.configure(**opts)
 
     with context.begin_transaction():
         context.run_migrations()
diff --git a/nonebot_plugin_orm/templates/multidb/env.py b/nonebot_plugin_orm/templates/multidb/env.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import cast
+from typing import Any, cast
 
 from alembic import context
 from sqlalchemy.util import await_only
@@ -51,37 +51,31 @@ def run_migrations_offline() -> None:
         config.print_stdout(f"迁移数据库 {name or '<default>'} 中 ...")
         file_ = f"{name}.sql"
         with open(file_, "w") as buffer:
-            context.configure(
-                **{
-                    **dict(
-                        url=engine.url,
-                        dialect_opts={"paramstyle": "named"},
-                        output_buffer=buffer,
-                        target_metadata=target_metadatas["name"],
-                        literal_binds=True,
-                    ),
-                    **plugin_config.alembic_context,
-                }
-            )
+            opts: dict[str, Any] = {
+                "url": engine.url,
+                "dialect_opts": {"paramstyle": "named"},
+                "output_buffer": buffer,
+                "target_metadata": target_metadatas[name],
+                "literal_binds": True,
+            } | plugin_config.alembic_context
+            context.configure(**opts)
+
             with context.begin_transaction():
                 context.run_migrations(name=name)
             config.print_stdout(f"将输出写入到 {file_}")
 
 
 def do_run_migrations(conn: Connection, name: str, metadata: MetaData) -> None:
-    context.configure(
-        **{
-            **dict(
-                connection=conn,
-                render_as_batch=True,
-                target_metadata=metadata,
-                include_object=no_drop_table,
-                upgrade_token=f"{name}_upgrades",
-                downgrade_token=f"{name}_downgrades",
-            ),
-            **plugin_config.alembic_context,
-        }
-    )
+    opts: dict[str, Any] = {
+        "connection": conn,
+        "render_as_batch": True,
+        "target_metadata": metadata,
+        "include_object": no_drop_table,
+        "upgrade_token": f"{name}_upgrades",
+        "downgrade_token": f"{name}_downgrades",
+    } | plugin_config.alembic_context
+    context.configure(**opts)
+
     context.run_migrations(name=name)
 
 
