Running migrations for all databases using different metadata

[notypecheck]

Hello, I need to generate migrations for multiple databases and upgrade them.
Ideally I would want to be able to do autogeneration and upgrades with a single command (alembic revision --autogenerate, etc) without having to run that separately for each database.
To achieve that I've tried to run run_migrations function in a for loop with different arguments (db connection, metadata, etc):

def run_migrations(migration_config: _MigrationConfig, config: Config) -> None:
    if not config.get_main_option("sqlalchemy.url"):
        config.set_main_option("sqlalchemy.url", migration_config.settings.url)

    script_location = config.get_section_option(
        config.config_ini_section,
        "script_location",
    )
    version_locations = f"{script_location}/versions/{migration_config.name}"
    context.script.version_locations = [version_locations]

    connectable = context.config.attributes.get("connection")
    if not connectable:
        if (section := config.get_section(config.config_ini_section)) is None:
            raise ValueError

        connectable = AsyncEngine(
            engine_from_config(
                section,
                poolclass=pool.NullPool,
                future=True,
            ),
        )

    if isinstance(connectable, AsyncEngine):
        asyncio.run(
            do_run_migrations_async(
                connectable=connectable,
                target_metadata=migration_config.metadata,
            ),
        )
    else:
        do_run_migrations(connectable, target_metadata=migration_config.metadata)

here I'm trying to override version_locations to a different subdirectory depending on the config, but it seems like alembic only generates revisions after env.py is done executing.
To work around that I decided to use the -x argument and pass db I need to migrate into it to select a _MigrationConfig I need. This isn't quite the desired behavior, and I don't think built-in multidb template solves it either (since I want each db to have it's own Metadata), so I'm wondering if there's a solution for this problem. It seems like with this configuration migrations themselves should work fine, but revision autogeneration does not.

[zzzeek]

if you are using the -x argument, that implies you are running the CLI multiple times with different x arguments, is that right?

i dont have a clear picture of what you have so far since your example does not show any "for loop" or what _MigrationConfig looks like, but overall, for the top level goal "I want to run one alembic command and have it process multiple metadatas / sets of revision files", there are two general ways to do that.

One is, use the command API directly and make your own front-end that does the things you want, in this case just runs the desired command multiple times with different configs This is the safe / easy / maybe a little tedious way to go but it's very simple and Alembic isn't made to jump through any kind of hoops it was not designed for.

The other is more like what you have already, which I think can work, but what you need to do specifically for the autogenerate case (and likely not for the upgrade/downgrade case) is set version_locations at config time to have all of the target directories present at once in a list; this just makes it possible to refer to any specific version_location later on. Then, when each autogenerate runs, you need to set the MigrationScript.version_script to the path you want for each autogenerate run. You can intercept MigrationScript using a process_revision_directives hook. then when the "revision" command does that single step to write all the files, it will deposit them into the individual directories named in each MigrationScript.

[notypecheck]

_MigrationConfig is just a dataclass with a metadata that I want to use, db name and connection settings:

@dataclasses.dataclass(slots=True, kw_only=True, frozen=True)
class _MigrationConfig:
    name: str
    metadata: MetaData
    settings: DatabaseSettings

As I said Ideally I would like to just run alembic revision --autogenerate and migrations belonging to one metadata would be generated in one path, and to second one into different path.
I will try to use the hook you mentioned 😃

[notypecheck]

@zzzeek Do I understand correctly that in order for the latter option to work I would have to execute alembic multiple times?
This is what I have now and it seems to only generate revision for the latest version_path.
Is that because process_revision_directives.directives argument contains all revisions?

import asyncio
import dataclasses
import importlib
import pkgutil
from collections.abc import Iterable, Sequence
from logging.config import fileConfig

from alembic import context
from alembic.config import Config
from alembic.operations import MigrationScript
from alembic.runtime.migration import MigrationContext
from sqlalchemy import MetaData, engine_from_config, pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import AsyncEngine

from app.connectors.db.base import DocumentBase, UserBase
from app.settings import DatabaseSettings, db_documents, db_user_data

config = context.config

if config.config_file_name is not None:
    fileConfig(config.config_file_name)


@dataclasses.dataclass(slots=True, kw_only=True, frozen=True)
class _MigrationConfig:
    name: str
    metadata: MetaData
    settings: DatabaseSettings


def do_run_migrations(
    connection: Connection,
    migration_config: _MigrationConfig, 
    version_location: str,
) -> None:
    def process_revision_directives(
        context: MigrationContext,  # noqa: ARG001
        revision: str | Iterable[str | None] | Iterable[str],  # noqa: ARG001
        directives: list[MigrationScript],
    ) -> None:
        for directive in directives:
            directive.version_path = version_location

    context.configure(
        connection=connection,
        target_metadata=migration_config.metadata,
        compare_server_default=True,
        process_revision_directives=process_revision_directives,
    )

    with context.begin_transaction():
        context.run_migrations()


async def do_run_migrations_async(
    connectable: AsyncEngine,
    migration_config: _MigrationConfig,
    version_location: str,
) -> None:
    async with connectable.connect() as conn:
        await conn.run_sync(
            do_run_migrations,
            migration_config=migration_config,
            version_location=version_location,
        )


def run_migrations(
    migration_configs: Sequence[_MigrationConfig],
    migration_config: _MigrationConfig,
    config: Config,
) -> None:
    if not config.get_main_option("sqlalchemy.url"):
        config.set_main_option("sqlalchemy.url", migration_config.settings.url)

    script_location = config.get_section_option(
        config.config_ini_section,
        "script_location",
    )
    context.script.version_locations = [
        f"{script_location}/versions/{cfg.name}" for cfg in migration_configs
    ]
    version_location = f"{script_location}/versions/{migration_config.name}"

    connectable = context.config.attributes.get("connection")
    if not connectable:
        if (section := config.get_section(config.config_ini_section)) is None:
            raise ValueError

        connectable = AsyncEngine(
            engine_from_config(
                section,
                poolclass=pool.NullPool,
                future=True,
            ),
        )

    if isinstance(connectable, AsyncEngine):
        asyncio.run(
            do_run_migrations_async(
                connectable=connectable,
                migration_config=migration_config,
                version_location=version_location,
            ),
        )
    else:
        do_run_migrations(
            connectable=connectable,
            migration_config=migration_config,
            version_location=version_location,
        )


def main() -> None:
    migration_configs = [
        _MigrationConfig(
            metadata=UserBase.metadata,
            settings=db_documents,
            name="users",
        ),
        _MigrationConfig(
            metadata=DocumentBase.metadata,
            settings=db_user_data,
            name="documents",
        ),
    ]
    for migration_config in migration_configs:
        run_migrations(
            migration_configs=migration_configs,
            migration_config=migration_config,
            config=config,
        )


main()

[zzzeek]

I think you need to set those version_locations in the config directly using set_section_option, or otherwise you need to look into ScriptDirectory and reset the memoized _version_locations attribute.

this is not a developed usecase of alembic so I'd use print / pdb to see how it proceeds. if some obvious hooks to make this easier stand out we can add them

[notypecheck]

Assigning to version_locations seems to work, since I'm not getting an error related to missing locations.
I didn't manage to get your suggestion to work:

def process_revision_directives(
    context: MigrationContext,  # noqa: ARG001
    revision: str | Iterable[str | None] | Iterable[str],  # noqa: ARG001
    directives: list[MigrationScript],
) -> None:
    print(directives, version_location)
    for directive in directives:
        print(directive.upgrade_ops_list)
        directive.version_path = version_location

seems like with the way I'm trying to do it the same MigrationScript object is used and is extended with new operations, so I can't simply assign a new location to it:

$ alembic revision --autogenerate
# print(context.script.version_locations, version_location)
['db/migrations/versions/documents', 'db/migrations/versions/users'] db/migrations/versions/documents
Context impl PostgresqlImpl.
Will assume transactional DDL.
Detected added table 'message'
[<alembic.operations.ops.MigrationScript object at 0x0000027404549F50>] db/migrations/versions/documents
[<alembic.operations.ops.UpgradeOps object at 0x0000027404549E50>]

# print(context.script.version_locations, version_location)
['db/migrations/versions/documents', 'db/migrations/versions/users'] db/migrations/versions/users     
Context impl PostgresqlImpl.
Will assume transactional DDL.
Detected added table 'user'
[<alembic.operations.ops.MigrationScript object at 0x0000027404549F50>] db/migrations/versions/users
[<alembic.operations.ops.UpgradeOps object at 0x0000027404549E50>, <alembic.operations.ops.UpgradeOps object at 0x000002740856D650>]
Generating db\migrations\versions\documents\2024-01-20__3d1cfc370dd6.py ...  done

[zzzeek]

there should be multiple MigrationScript instances. feel free to make new ones and append them to the list inside the hook