diff --git a/src/plugins/nonebot_bison/config/db.py b/src/plugins/nonebot_bison/config/db.py index 94d2cf4..47d2baa 100644 --- a/src/plugins/nonebot_bison/config/db.py +++ b/src/plugins/nonebot_bison/config/db.py @@ -3,8 +3,12 @@ from pathlib import Path import nonebot from alembic import command from alembic.config import Config +from alembic.runtime.environment import EnvironmentContext +from alembic.script.base import ScriptDirectory from nonebot_plugin_datastore import PluginData, create_session, db +from .db_model import Base + DATA = PluginData("bison") @@ -16,3 +20,20 @@ async def upgrade_db(): ) alembic_cfg.set_main_option("sqlalchemy.url", "") command.upgrade(alembic_cfg, "head") + + script = ScriptDirectory.from_config(alembic_cfg) + + def upgrade(rev, context): + return script._upgrade_revs("head", rev) + + with EnvironmentContext( + alembic_cfg, + script, + fn=upgrade, + as_sql=False, + starting_rev=None, + destination_rev="head", + tag=None, + target_metadata=Base.metadata, + ): + script.run_env() diff --git a/src/plugins/nonebot_bison/config/migrate/env.py b/src/plugins/nonebot_bison/config/migrate/env.py index fef0464..6ab47e5 100644 --- a/src/plugins/nonebot_bison/config/migrate/env.py +++ b/src/plugins/nonebot_bison/config/migrate/env.py @@ -24,13 +24,13 @@ import nonebot try: nonebot.get_driver() __as_plugin = True + target_metadata = None except: __as_plugin = False nonebot.init() + from nonebot_bison.config.db_model import Base -from nonebot_bison.config.db_model import Base - -target_metadata = Base.metadata + target_metadata = Base.metadata # other values from the config, defined by the needs of env.py, # can be acquired: @@ -63,7 +63,10 @@ def run_migrations_offline(): def do_run_migration(connection: Connection): - context.configure(connection=connection, target_metadata=target_metadata) + if __as_plugin: + context.configure(connection=connection) + else: + context.configure(connection=connection, target_metadata=target_metadata) with context.begin_transaction(): context.run_migrations()