Deadlock while migrating the same set of revisions for multiple schemas in postgreSQL

[jay-patel-simformsolutions]

My requirement is that I have a series of migrations and I apply the same set of migrations on multiple database schemas (a typical multitenant requirement). I have updated my env.py file to do this customer behavior in this way

import logging
from logging.config import fileConfig

from alembic import context
from alembic.context import get_x_argument
from db import db
from flask import current_app
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.sql import text

x_args = get_x_argument(as_dictionary=True)


def get_customers_list():
    if "schema_name" not in x_args:
        customer_list = [("public", "PUBLIC")]
        try:
            customer = db.session.execute(text("SELECT id, customer_type FROM public.customer"))
            # customer[0] is the uuid of the customer and customer[1] is the customer_type("WAREHOUSE"/"PARTNER")
            customer_list = list(
                map(lambda customer: (str(customer[0]), customer[1]), customer)
            ) + [("public", "PUBLIC")]
        except ProgrammingError:
            pass
        return customer_list
    try:
        customer = db.session.execute(
            text(
                f"SELECT id, customer_type FROM public.customer where id = '{x_args['schema_name']}'"
            )
        )
        # customer[0] is the uuid of the customer and customer[1] is the customer_type("WAREHOUSE"/"PARTNER")
        customer_list = list(map(lambda customer: (str(customer[0]), customer[1]), customer))
        return customer_list
    except ProgrammingError:
        return None


# 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.
fileConfig(config.config_file_name)
logger = logging.getLogger("alembic.env")


def include_object(object, name, type_, reflected, compare_to):
    if type_ == "table" and reflected and compare_to is None:
        return False
    return True


def get_engine():
    try:
        # this works with Flask-SQLAlchemy<3 and Alchemical
        return current_app.extensions["migrate"].db.get_engine()
    except TypeError:
        # this works with Flask-SQLAlchemy>=3
        return current_app.extensions["migrate"].db.engine


def get_engine_url():
    try:
        return get_engine().url.render_as_string(hide_password=False).replace("%", "%%")
    except AttributeError:
        return str(get_engine().url).replace("%", "%%")


# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
config.set_main_option("sqlalchemy.url", get_engine_url())
target_db = current_app.extensions["migrate"].db

# 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 get_metadata():
    if hasattr(target_db, "metadatas"):
        return target_db.metadatas[None]
    return target_db.metadata


def run_migrations_offline():
    """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=get_metadata(),
        literal_binds=True,
        include_object=include_object,
    )

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


def run_migrations_online():
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """

    customer_list = get_customers_list()

    if not customer_list:
        return

    # this callback is used to prevent an auto-migration from being generated
    # when there are no changes to the schema
    # reference: http://alembic.zzzcomputing.com/en/latest/cookbook.html
    def process_revision_directives(context, revision, directives):
        if getattr(config.cmd_opts, "autogenerate", False):
            script = directives[0]
            if (len(script._upgrade_ops) == 0 or script._upgrade_ops[0].is_empty()) and len(
                customer_list
            ) == len(script._upgrade_ops):
                logger.info("No changes in schema detected.")
                directives[:] = []

    connectable = get_engine()

    for schema, schema_type in customer_list:
        if schema != "public":
            db.session.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema}";'))
            db.session.commit()
        with connectable.connect() as connection:
            context.configure(
                include_object=include_object,
                connection=connection,
                target_metadata=get_metadata(),
                process_revision_directives=process_revision_directives,
                version_table_schema=schema,
                **current_app.extensions["migrate"].configure_args,
            )

            with context.begin_transaction():
                context.run_migrations(schema=schema, schema_type=schema_type)


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

This works perfectly fine when I run the command python -m flask db upgrade but when the flask server is running and I want to upgrade the database migrations when a new customer is added, so I do a flask migrate in the API call as bellow

with app.app_context():
        flask_migrate.upgrade()

In this process everything works fine until there are 2 revisions lets say revision1 and revision2 and in revision1 we are creating a constraint

revision1.py

def upgrade(schema="public", schema_type="PUBLIC"):
    if schema_type == "WAREHOUSE":
        op.add_column(
            "appointment",
            sa.Column("created_by", sa.UUID(), nullable=True),
            schema=schema,
        )
        op.create_foreign_key(
            "created_by_id_fk", "appointment", "user", ["created_by"], ["id"], source_schema=schema
        )

and drop that in revision2

revision2.py

def upgrade(schema="public", schema_type="PUBLIC"):
    if schema_type == "WAREHOUSE":
        op.drop_constraint("created_by_id_fk", "appointment", schema=schema)

when I try to apply these revisions from the API it goes into deadlock but works correctly when triggered via cli with flask db upgrade. Need some help in understanding and solving it.

[miguelgrinberg]

I think this question should be directed at the Alembic team. Flask-Migrate is just a wrapper that integrates Alembic into Flask.