Deadlock while updating a column, updated the env.py file to replicate same set of migrations on different schemas

[jay-patel-simformsolutions]

Describe the bug

I have updated the env.py (will be attaching the updated file below) file to replicate the same set of revisions on multiple schemas. And in that if I have 2 revisions lets call them revision1 and revision2 and I create a column in revision1 and alter the column in revision2 and I am triggering the upgrading of the migrations from the code not from the cli.

env.py

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)

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
        )

revision2.py

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

This causes a deadlock when the migration is triggered via code something like this

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

Expected behavior
Migration should happen without the deadlock anyways each revision is expected to run in a different transaction

The deadlock causing queries in postgreSQL:

ALTER TABLE "23d85593-8f70-4915-b5cb-445dff4ea0b8".appointment DROP CONSTRAINT created_by_id_fk

SELECT "user".id AS user_id, "user".first_name AS user_first_name, "user".last_name AS user_last_name, "user".email AS user_email, "user".customer_id AS user_customer_id, "user".is_signed_up AS user_is_signed_up, "user".is_admin_approved AS user_is_admin_approved, "user".role AS user_role, "user".invite_accepted_at AS user_invite_accepted_at, "user".timezone AS user_timezone, "user".profile_picture_path AS user_profile_picture_path, "user".preferred_language AS user_preferred_language, "user".created_at AS user_created_at, "user".updated_at AS user_updated_at 
FROM "user" 
WHERE "user".id = '631190ac-18c5-44f7-927f-223dae31337a'::uuid::UUID

Versions.

• OS: Windows
• Python: 3.9
• Alembic:1.9.3
• SQLAlchemy:2.0.3
• Database:PostgreSQL
• DBAPI:

Additional context

I have tried adding these 2 parameters transaction_per_migration=True transactional_ddl=True in configure with these params it gave an error that the constraint created_by_id_fk does not exists

Have a nice day!

[zzzeek]

hi -

I'm not sure where the exact source of your deadlock is, however Alembic does not connect to the database, your own code in env.py is where that happens. I see you are mixing up the use of ORM sessions coming from Flask bound to some other engine somewhere, along with connections you are starting up locally in your env.py. this will cause deadlocks and this is expected. I think you'd be better off ensuring all the operations you are doing share a single connection and transaction, something like this:

        with connectable.connect() as connection:
            for schema, schema_type in customer_list:
                if schema != "public":
                    connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema}";'))
                    connection.commit()
                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)

[jay-patel-simformsolutions]

I tried but did not work

Let me give more clarity

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.

[CaselIT]

All you describe fits perfectly with what mike said.
Does removing all use of db.session from the migrations and env improves things?

[zzzeek]

I tried but did not work

Your new code example still refers to "db.session" throughout, being mixed with a Connection that you are getting yourself using get_engine(), and while this code might work sometimes, it is not in any way surprising that it deadlocks. All of that should be replaced with the use of a single Connection that you get from your .connect() call within env.py. use only a single database connection within all of env.py. pass the Connection into the get_customers_list() function. if you need to use the ORM Session interface (it does not seem that way since you have all text() statements), you can make a new Session using Session(bind=connection).

Let me give more clarity

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).

this is not really relevant to the actual error you are getting, which is that you are running statements in multiple transactions at once and getting deadlocks. Using multiple schemas does not necessitate the use of a db.session from Flask that forces itself to make a separate connection from the one you are using

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

the Flask db.session is probably still opened within a transaction. It doesn't really matter until the programming patterns in the env.py here are fixed

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.

remove the use of db.session that is making its own Connection (or if Flask supports some option, get db.session to use the same Connection that you are setting up in your env.py. use one Connection at a time, where you control where the transaction begins and ends. this is how to get rid of deadlocks.

[jay-patel-simformsolutions]

Thanks @zzzeek and @CaselIT for the help!! Means a lot

Seems Like our issue was different

This is what I was doing in the Flask API

class PostAPI(Resource):
    def post(self):
        from app import migrate_handler
        .
        . Logic to create the customer
        .
        migrate_handler(None, None, app=current_app)  # Apply migrations for the new customer
        .
        .
        .
        return "CREATED"

app.py

import awsgi
import flask_migrate
from create_feature_app import create_feature_app

app = create_feature_app()


def lambda_handler(event: dict, context: dict):
    response = awsgi.response(app, event, context, base64_content_types={"image/png"})
    return response


def migrate_handler(event: dict, context: dict, app=app):
    with app.app_context():
        flask_migrate.upgrade()
    return "Migrated"

lambda_handler was used to entertain API calls and migrate_handler was used to apply migrations (Some Serverless stuff not related to this issue)

but because I was importing the migrate_handler inside the post function (I had to because of the circular imports issue) it was creating 2 different instances of the Flask app. One used to listen to the API calls and the second one to migrate and when both happened at the same time it caused the deadlock.

But after discussing this issue here with @zzzeek and @CaselIT I feel like the env.py file is also vulnerable to such deadlocks in the future so anyways I unified and used one common db connection for all the queries and migrations. Thank you guys for the help!!!