Как использовать одни и те же миграции alembic для разных баз данных

У меня есть рабочая конфигурация миграций alembic для моей базы данных postgres под названием my_db. Теперь я хочу использовать эти миграции в другой базе данных postgres под названием my_db_test при запуске моих юнит-тестов.

Это начало моего файла migrations/env.py (который работает совершенно нормально при обычном использовании):

import asyncio
from logging.config import fileConfig

from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config

from alembic import context
import os # ИЗМЕНЕН
import pkgutil  # ИЗМЕНЕН
import importlib  # ИЗМЕНЕН
from config import config as cfg  # ИЗМЕНЕН
from sqlmodel import SQLModel  # ИЗМЕНЕН

# Получите URL базы данных из конфигурации  # ИЗМЕНЕН
db_url = cfg.SQLALCHEMY_DATABASE_URI  # ИЗМЕНЕН

# Укажите путь к каталогу ваших моделей  # ИЗМЕНЕН
models_dir = os.path.join(os.path.dirname(__file__), "..", "models")  # ИЗМЕНЕН

# Динамически импортируйте все Python-файлы в каталоге моделей  # ИЗМЕНЕН
for module_info in pkgutil.iter_modules([models_dir]):  # ИЗМЕНЕН
    importlib.import_module(f"models.{module_info.name}")  # ИЗМЕНЕН

# это объект конфигурации Alembic, который предоставляет
# доступ к значениям в используемом .ini файле.
config = context.config

config.set_main_option('sqlalchemy.url', db_url)  # ИЗМЕНЕН

Идея заключалась в том, чтобы по умолчанию использовать мою обычную базу данных и переключаться на test_db, используя флаг аргумента. Я добавил аргумент парсер в этот скрипт, но, к сожалению, это не так просто, как я надеялся. Я не смог напрямую внедрить переменную в файл env.py с помощью стандартной команды, такой как: alembic upgrade head --db_url <test_db>

Поэтому я попытался создать оболочку alembic на python, которая принимает команду alembic вместе с аргументом и обновляет db_url следующим образом:

# ./alembic_wrapper.py
import sys
import argparse
from alembic import command
from alembic.config import Config
from argparse import ArgumentParser
from config import config as cfg

# Функция для парсинга аргументов командной строки
def get_db_url_from_args():
    parser = ArgumentParser(description="Переопределите URL базы данных для миграций Alembic.")
    parser.add_argument(
        "--db-url", 
        type=str, 
        help="Переопределите URL базы данных (по умолчанию: из конфигурации)"
    )
    parser.add_argument(
        "alembic_command", 
        choices=["upgrade", "downgrade", "revision", "history", "current", "stamp", "merge"],
        help="Команда Alembic для выполнения"
    )
    parser.add_argument(
        "alembic_args", 
        nargs=argparse.REMAINDER, 
        help="Аргументы для команды Alembic"
    )

    args = parser.parse_args()
    return args.db_url or cfg.SQLALCHEMY_DATABASE_URI, args.alembic_command, args.alembic_args

# Получите URL базы данных и аргументы команды Alembic
db_url, alembic_command, alembic_args = get_db_url_from_args()

# Настройте Alembic для использования соответствующего db_url
config = Config("alembic.ini")
config.set_main_option('sqlalchemy.url', db_url)
print(f"Используется URL базы данных: {db_url}")

# Выполните команду Alembic
if alembic_command == "upgrade":
    command.upgrade(config, *alembic_args)
elif alembic_command == "downgrade":
    command.downgrade(config, *alembic_args)
elif alembic_command == "revision":
    command.revision(config, *alembic_args)
elif alembic_command == "history":
    command.history(config, *alembic_args)
elif alembic_command == "current":
    command.current(config, *alembic_args)
elif alembic_command == "stamp":
    command.stamp(config, *alembic_args)
elif alembic_command == "merge":
    command.merge(config, *alembic_args)

Эта оболочка работает, если я не передаю ей аргументы. Если я просто запускаю python alembic_wrapper.py upgrade head, все работает отлично. Однако, когда я передаю db_url python alembic_wrapper.py upgrade head --db-url <db_url> (не имеет значения, test_db/prod_db или неверный URL, просто выводятся все SQL-запросы (что обычно не происходит) и фактически ничего не происходит (так как в моей базе данных нет изменений).

Вывод:

(condaenv) PS C:\Users\MyUser\Documents\visual-studio-code\my-project\backend-v2> python alembic_wrapper.py upgrade head --db-url postgresql+asyncpg://user:password@localhost:5432/my_db     
Используется URL базы данных: postgresql+asyncpg://user:password@localhost:5432/my_db    
INFO  [alembic.runtime.migration] Контекст реализует PostgresqlImpl.
INFO  [alembic.runtime.migration] Генерация статического SQL
INFO  [alembic.runtime.migration] Будет предполагаться транзакционный DDL.
BEGIN;

CREATE TABLE alembic_version (
    version_num VARCHAR(32) NOT NULL,
    CONSTRAINT alembic_version_pkc PRIMARY KEY (version_num)
);

INFO  [alembic.runtime.migration] Выполняется обновление  -> f908fe4a4afd, Первоначальная миграция, включая все таблицы
-- Выполняется обновление  -> f908fe4a4afd

-- Много SQL-запросов

COMMIT;

Я не знаю, есть ли какое-то более лучшее решение в общем, чем делать это так, но я хочу продолжать использовать только одну папку миграций для моего проекта и не хочу обновлять свой файл env.py вручную каждый раз, когда я использую test_db.