diff --git a/app/db/plugin/migration.py b/app/db/plugin/migration.py index 810232d59..a41c199c9 100644 --- a/app/db/plugin/migration.py +++ b/app/db/plugin/migration.py @@ -7,10 +7,54 @@ from pathlib import Path from alembic.command import upgrade from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.script import ScriptDirectory +from sqlalchemy.engine import Connection from app.db.plugin.container import PluginDatabaseHandle -__all__ = ["run_migrations"] +__all__ = ["PluginMigrationCompatibilityError", "run_migrations"] + + +class PluginMigrationCompatibilityError(RuntimeError): + """插件库里存在随包迁移树无法识别的 revision，不能按当前版本继续升级。""" + + +def _reject_unknown_revisions( + connection: Connection, + config: Config, + directory: Path, +) -> None: + """ + 在执行升级前确认库里已记录的 revision 都能被随包迁移树识别。 + + 插件装回旧版本后，库里留着的是新版本写下的 revision，而随包回退的迁移树里没有这 + 条记录。此时 ``upgrade`` 只抛一句不带插件名也不带目录的 ``Can't locate revision``， + 排障得靠猜是哪个插件、哪个库。这里在同一批连接上先把这个事实查清楚，换成指名道姓 + 的拒绝：宁可拒绝建库，也绝不自动 ``downgrade``——降级脚本通常只负责删表删列，自动 + 执行等于把用户在新版本里产生的数据直接抹掉，且没有回头路。 + :param connection: 已指向插件库（PostgreSQL 下已限定 schema）的连接 + :param config: 已设定 script_location 的 Alembic 配置 + :param directory: 插件随包提供的迁移脚本目录 + :raise PluginMigrationCompatibilityError: 库里存在迁移树不认识的 revision + """ + known = { + script.revision + for script in ScriptDirectory.from_config(config).walk_revisions() + } + # 空库还没有 alembic_version 表，get_current_heads 返回空元组，首次建库不受影响 + unknown = sorted( + revision + for revision in MigrationContext.configure(connection).get_current_heads() + if revision not in known + ) + if not unknown: + return + raise PluginMigrationCompatibilityError( + f"插件数据库当前 revision {'、'.join(unknown)} 不属于迁移脚本目录 {directory}，" + "拒绝按当前插件版本继续升级；请换用包含该 revision 的插件版本，" + "或由管理员完成明确的数据迁移，宿主不会自动降级数据库" + ) def run_migrations(handle: PluginDatabaseHandle, directory: Path) -> None: @@ -22,6 +66,7 @@ def run_migrations(handle: PluginDatabaseHandle, directory: Path) -> None: 取用它，否则迁移会落在 public schema 而不是该插件的 schema。 :param handle: 插件数据库句柄 :param directory: 迁移脚本目录，须符合 Alembic script_location 布局 + :raise PluginMigrationCompatibilityError: 库里存在迁移树不认识的 revision """ config = Config() # 关闭 ini 插值：迁移目录里出现的 % 不应被当成插值语法（与宿主 _build_alembic_config 一致） @@ -32,9 +77,13 @@ def run_migrations(handle: PluginDatabaseHandle, directory: Path) -> None: "sqlalchemy.url", handle.engine.url.render_as_string(hide_password=False), ) + # 独占引擎下 alembic 自己建连接，兼容性检查用完即还，不与后续升级争同一把写锁 + with handle.engine.connect() as connection: + _reject_unknown_revisions(connection, config, directory) upgrade(config, "head") return with handle.engine.connect() as connection: config.attributes["connection"] = connection + _reject_unknown_revisions(connection, config, directory) upgrade(config, "head") connection.commit() diff --git a/docs/rules/10-data-and-persistent.md b/docs/rules/10-data-and-persistent.md index effaab25d..f856ddff6 100644 --- a/docs/rules/10-data-and-persistent.md +++ b/docs/rules/10-data-and-persistent.md @@ -281,6 +281,20 @@ it must be an absolute, existing directory — `ensure` raises `FileNotFoundError` before creating anything rather than leaving behind a database with neither tables nor a version stamp. +Before the upgrade runs, `run_migrations` compares the revisions already +stamped in the plugin database against the script tree shipped with the +installed plugin version. A revision that tree cannot locate means the +database was stamped by a newer build — the ordinary outcome of reinstalling +an older version — and the host raises `PluginMigrationCompatibilityError` +naming both the revision and the directory instead of letting Alembic fail +with a bare `Can't locate revision`. The host never downgrades automatically: +downgrade scripts normally drop tables and columns, so running them unattended +would destroy whatever the user produced under the newer version, with no way +back. Recovery is to install a plugin version whose migration tree contains +that revision, or to have an administrator perform an explicit data migration. +A database with no `alembic_version` table yet reports no revisions, so first +creation is never blocked by this check. + Models must inherit `app.sdk.database.plugin_declarative_base()`, which mints a fresh `MetaData` per call so a plugin's tables never collide with `app.db.base.Base.metadata` or with another plugin's same-named tables. At diff --git a/tests/test_db_plugin_framework.py b/tests/test_db_plugin_framework.py index d2eaa8928..287a65b8c 100644 --- a/tests/test_db_plugin_framework.py +++ b/tests/test_db_plugin_framework.py @@ -63,13 +63,20 @@ def _borrowed_engine_handle(engine: Any) -> PluginDatabaseHandle: ) -def _write_migration_directory(root: Path) -> Path: +def _write_migration_directory( + root: Path, + *, + name: str = "migrations", + follow_up_revision: str | None = None, +) -> Path: """ 写出一个最小可用的 Alembic script_location，只含一条建表迁移。 :param root: 承载迁移目录的父目录 + :param name: 迁移目录名，用于在同一父目录下并存新旧两棵迁移树 + :param follow_up_revision: 追加一条以 0001 为父的后续迁移，模拟插件新版本 :return: 迁移目录路径 """ - directory = root / "migrations" + directory = root / name (directory / "versions").mkdir(parents=True) (directory / "env.py").write_text( """from alembic import context @@ -116,6 +123,30 @@ def downgrade(): ''', encoding="utf-8", ) + if follow_up_revision is not None: + (directory / "versions" / f"{follow_up_revision}_add_tags.py").write_text( + f'''"""插件新版本追加的表。""" + +import sqlalchemy as sa +from alembic import op + +revision = "{follow_up_revision}" +down_revision = "0001" +branch_labels = None +depends_on = None + + +def upgrade(): + """建出新版本才有的表。""" + op.create_table("tags", sa.Column("id", sa.Integer(), primary_key=True)) + + +def downgrade(): + """回退建表。""" + op.drop_table("tags") +''', + encoding="utf-8", + ) return directory @@ -631,6 +662,48 @@ def test_run_migrations_upgrades_a_sqlite_plugin_database(plugin_data_root, sqli assert set(sa_inspect(handle.engine).get_table_names()) == tables +def test_run_migrations_rejects_a_revision_the_script_tree_cannot_locate( + plugin_data_root, + sqlite_backend, + tmp_path, +): + """装回旧版本时库里残留的新 revision 必须被明确拒绝，且不得自动降级。""" + newer = _write_migration_directory( + tmp_path, + name="migrations-new", + follow_up_revision="0002", + ) + registry_module.ensure_database("demo", (), newer) + older = _write_migration_directory(tmp_path, name="migrations-old") + + with pytest.raises(migration_module.PluginMigrationCompatibilityError) as error: + registry_module.ensure_database("demo", (), older) + + assert "0002" in str(error.value) + handle = registry_module.get_database("demo") + with handle.engine.connect() as connection: + stamped = connection.execute( + text("SELECT version_num FROM alembic_version") + ).scalar() + assert stamped == "0002" + assert "tags" in set(sa_inspect(handle.engine).get_table_names()) + + +def test_run_migrations_accepts_a_database_stamped_by_the_same_script_tree( + plugin_data_root, + sqlite_backend, + tmp_path, +): + """同一棵迁移树内的旧 revision 仍须正常升到 head，兼容检查不得误伤。""" + directory = _write_migration_directory(tmp_path, follow_up_revision="0002") + registry_module.ensure_database("demo", (), directory) + + registry_module.ensure_database("demo", (), directory) + + handle = registry_module.get_database("demo") + assert {"notes", "tags"} <= set(sa_inspect(handle.engine).get_table_names()) + + def test_run_migrations_routes_the_postgresql_connection_through_the_handle( monkeypatch, tmp_path, @@ -643,13 +716,21 @@ def test_run_migrations_routes_the_postgresql_connection_through_the_handle( captured["connection"] = config.attributes.get("connection") captured["revision"] = revision + def _record_guard(connection, _config, _directory): + """记录 revision 兼容检查落在哪条连接上。""" + captured["guarded"] = connection + monkeypatch.setattr(migration_module, "upgrade", _record_upgrade) + # 替身引擎答不出真实 revision，这里只验证路由：兼容检查必须与升级共用同一条连接， + # 否则它读到的是未被 schema_translate_map 限定的库 + monkeypatch.setattr(migration_module, "_reject_unknown_revisions", _record_guard) handle = _borrowed_engine_handle(MagicMock(name="derived_engine")) migration_module.run_migrations(handle, _write_migration_directory(tmp_path)) connection = handle.engine.connect.return_value.__enter__.return_value assert captured["connection"] is connection + assert captured["guarded"] is connection assert captured["revision"] == "head" connection.commit.assert_called_once()