Effortlessly Manage Multiple Databases in FastAPI with SQLAlchemy and Alembic

docker run --name mysqldb -d -p 3306:3306 -e MYSQL_ROOT_PASSWORD=root -e MYSQL_USER=dev -e MYSQL_PASSWORD=root -e MYSQL_DATABASE=dev mysql:latest

docker exec -it mysqldb /bin/bash

mysql -u root -p
Enter password: 

mysql> show databases;
+--------------------+
| Database           |
+--------------------+
| dev                |
| information_schema |
| mysql              |
| performance_schema |
| sys                |
+--------------------+
5 rows in set (0.02 sec)

mysql> create database test;
Query OK, 1 row affected (0.03 sec)

mysql> create database qa;
Query OK, 1 row affected (0.02 sec)

mysql> create database prod;
Query OK, 1 row affected (0.00 sec)

mysql> show databases;
+--------------------+
| Database           |
+--------------------+
| dev                |
| information_schema |
| mysql              |
| performance_schema |
| prod               |
| qa                 |
| sys                |
| test               |
+--------------------+
8 rows in set (0.01 sec)

DB_USER=username
DB_PASSWORD=password
DB_HOST=localhost
DB_PORT=3306

pip install alembic sqlalchemy fastapi aiomysql python-dotenv

alembic init alembic

#!/usr/bin/env python
# alembic/env.py
import asyncio
from logging.config import fileConfig
from sqlalchemy import engine_from_config, pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import AsyncEngine

from alembic import context

from app.database.base import Base

# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
    fileConfig(config.config_file_name)
# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.


def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode.
    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.
    Calls to context.execute() here emit the given string to the
    script output.
    """
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )
    with context.begin_transaction():
        context.run_migrations()


def do_run_migrations(connection: Connection) -> None:
    context.configure(connection=connection, target_metadata=target_metadata)

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


async def run_migrations_online() -> None:
    """Run migrations in 'online' mode.
    In this scenario we need to create an Engine
    and associate a connection with the context.
    """
    configuration = config.get_section(config.config_ini_section)
    configuration["sqlalchemy.url"] = config.get_main_option("sqlalchemy.url")

    connectable = AsyncEngine(
        engine_from_config(
            configuration,
            prefix="sqlalchemy.",
            poolclass=pool.NullPool,
            future=True,
        )
    )

    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)

    await connectable.dispose()


if context.is_offline_mode():
    run_migrations_offline()
else:
    asyncio.run(run_migrations_online())

# alemibic_cli.py is a script that allows you to run Alembic migrations for multiple databases.
from datetime import datetime
import os
import argparse
from alembic.config import Config
from alembic import command

from dotenv import load_dotenv

load_dotenv()

# Define databases
DATABASES = {'dev', 'test', 'qa', 'prod'}
print(DATABASES)


def generate_db_url(db_name):
    return f"mysql+aiomysql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}@{os.getenv('DB_HOST')}:{os.getenv('DB_PORT')}/{db_name}"


def create_revision(db_name, config_path, message):
    """Create a new revision for a specific db_name."""
    versions_dir = os.path.join("alembic", "versions", db_name)
    os.makedirs(versions_dir, exist_ok=True)

    config = Config(config_path)
    config.set_main_option("version_locations", versions_dir)
    config.set_main_option(
        "sqlalchemy.url",
        generate_db_url(db_name),
    )
    print(f"revision for {db_name}: {config.get_main_option('sqlalchemy.url')}")
    command.revision(
        config, message=f"{message} on {datetime.now()}", autogenerate=True
    )


def upgrade_head(db_name, config_path):
    """Upgrade to head for a specific db_name."""
    config = Config(config_path)
    config.set_main_option(
        "sqlalchemy.url",
        generate_db_url(db_name),
    )
    config.set_main_option(
        "version_locations", os.path.join("alembic", "versions", db_name)
    )
    print(f"upgrade for {db_name}: {config.get_main_option('sqlalchemy.url')}")
    command.upgrade(config, "head")

def downgrade(db_name, revision, config_path):
    """Downgrade to a specific revision for a customer."""
    config = Config(config_path)
    config.set_main_option('sqlalchemy.url', generate_db_url(db_name))
    config.set_main_option('version_locations', os.path.join('alembic', 'versions', db_name))
    print(f"downgrade for {db_name}: {config.get_main_option('sqlalchemy.url')}")
    command.downgrade(config=config, revision=revision)

def migrate_database(db_name, action, config_path, message, revision=None):
    """Run migration for a specific db_name."""
    if action == "revision":
        create_revision(db_name, config_path, message)
    elif action == "upgrade":
        upgrade_head(db_name, config_path)
    elif action == 'downgrade' and revision:
        downgrade(db_name, revision, config_path)
    else:
        print(f"Error: Invalid action or missing revision for downgrade.")


def migrate_all(action, config_path, message):
    """Run migrations for all db_names."""
    for db_name in DATABASES:
        migrate_database(db_name, action, config_path, message)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Run Alembic migrations for db_names."
    )
    parser.add_argument(
        "action",
        choices=["revision", "upgrade", "downgrade"],
        help="Action to perform: 'revision' or 'upgrade', or 'downgrade'",
    )
    parser.add_argument(
        "--db", help="Specify a db_name to run the migration for"
    )
    parser.add_argument(
        "--message", help="Revision message for the migration", default="rev"
    )
    parser.add_argument('--revision', help="Specify the revision for downgrade (required for 'downgrade' action)")
    args = parser.parse_args()

    config_path = "./alembic.ini"
    
    if args.action == 'downgrade' and not args.revision:
        print("Error: Missing revision for downgrade.")
        exit(1)

    if args.db:
        if args.db in DATABASES:
            migrate_database(args.db, args.action, config_path, args.message, args.revision)
        else:
            print(
                f"Error: Database '{args.db}' not found in the list of databases."
            )
    else:
        migrate_all(args.action, config_path, args.message)

python alembic_cli.py <action> [--db <db_name>] [--message <message>] [--revision <revision>]

python alembic_cli.py revision --message "initial"

python alembic_cli.py revision --db dev --message "initial"

python alembic_cli.py upgrade

python alembic_cli.py upgrade --db prod

python alembic_cli.py downgrade --db test --revision <revision>