Alembic migrations fail when applying on distributed databases

[cfloressuazo]

Describe the bug When running multiple migrations (more than 1 in a new database) in a single command with YugabyteDB using flask db upgrade or alembic upgrade, the operation fails. The migration script runs and creates all the tables, but the last step where it has to update the version fails with a SerializationFailure error.

Expected behaviour The flask db upgrade command should successfully apply all migrations and update the alembic_version table with the correct version number.

To Reproduce The issue can be reproduced by running multiple migrations in a single command with YugabyteDB. Here's a simplified example:

1. Use a distributed database such as YugabyteDB (https://cloud.yugabyte.com/ for a free cluster).
2. Use a simple migrations like the two below:
3. Finally run: flask db upgrade or alembic upgrade

   
   # migrations/versions/eb6e925edca4_initial_migration.py
   """Initial migration.

Revision ID: eb6e925edca4 Revises: Create Date: 2023-11-28 14:59:44.514650

""" from alembic import op import sqlalchemy as sa

revision identifiers, used by Alembic.

revision = 'eb6e925edca4' down_revision = None branch_labels = None depends_on = None

def upgrade():

commands auto generated by Alembic - please adjust!

op.create_table('roles',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('name', sa.String(length=50), server_default='', nullable=False),
sa.Column('label', sa.Unicode(length=255), server_default='', nullable=True),
sa.PrimaryKeyConstraint('id'),
sa.UniqueConstraint('name')
)
op.create_table('users',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('email', sa.Unicode(length=255), server_default='', nullable=False),
sa.Column('email_confirmed_at', sa.DateTime(), nullable=True),
sa.Column('password', sa.String(length=255), server_default='', nullable=False),
sa.Column('active', sa.Boolean(), server_default='0', nullable=False),
sa.Column('first_name', sa.Unicode(length=50), server_default='', nullable=False),
sa.Column('last_name', sa.Unicode(length=50), server_default='', nullable=False),
sa.Column('twofa_enabled', sa.Boolean(), server_default='0', nullable=False),
sa.Column('country_code', sa.Integer(), server_default='0', nullable=True),
sa.Column('phone', sa.String(length=30), nullable=True),
sa.Column('authy_id', sa.Integer(), server_default='0', nullable=False),
sa.Column('authy_status', sa.Enum('unverified', 'onetouch', 'sms', 'token', 'approved', 'denied', name='authy_statuses'), nullable=True),
sa.Column('failed_login_attempts', sa.Integer(), nullable=True),
sa.Column('address1', sa.String(length=50), nullable=True),
sa.Column('address2', sa.String(length=50), nullable=True),
sa.Column('city', sa.String(length=50), nullable=True),
sa.Column('state', sa.String(length=50), nullable=True),
sa.Column('zipcode', sa.String(length=16), nullable=True),
sa.Column('country', sa.String(length=50), nullable=True),
sa.Column('profile', sa.Unicode(length=50), server_default='default', nullable=False),
sa.PrimaryKeyConstraint('id'),
sa.UniqueConstraint('email')
)
op.create_table('users_roles',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('user_id', sa.Integer(), nullable=True),
sa.Column('role_id', sa.Integer(), nullable=True),
sa.ForeignKeyConstraint(['role_id'], ['roles.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###

def downgrade():

commands auto generated by Alembic - please adjust!

op.drop_table('users_roles')
op.drop_table('users')
op.drop_table('roles')
try:
    op.execute('DROP TYPE authy_statuses;')
except Exception as e:
    print("Possible compatibility issue with the database engine. Please check the error message below.")
    print(e)
# ### end Alembic commands ###

and
```py
# migrations/versions/5be8363eb9b1_user_trial_period.py
"""user trial period

Revision ID: 5be8363eb9b1
Revises: eb6e925edca4
Create Date: 2024-05-24 01:09:51.493417

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '5be8363eb9b1'
down_revision = 'eb6e925edca4'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('users', sa.Column('extra', sa.String(length=255), nullable=True))
    op.add_column('users', sa.Column('last_login_time', sa.DateTime(), nullable=True))

    # Update enum type
    try:
        op.execute("ALTER TYPE authy_statuses ADD VALUE 'trial_bgn'")
        op.execute("ALTER TYPE authy_statuses ADD VALUE 'trial_mid'")
        op.execute("ALTER TYPE authy_statuses ADD VALUE 'trial_end'")
    except Exception as e:
        print("Failed to update enum type. Compatibility issue most likely. See error below:")
        print(e)
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('users', 'last_login_time')
    op.drop_column('users', 'extra')

    # Downgrade enum type (optional)
    try:
        op.execute("ALTER TYPE authy_statuses DROP VALUE 'trial_bgn'")
        op.execute("ALTER TYPE authy_statuses DROP VALUE 'trial_mid'")
        op.execute("ALTER TYPE authy_statuses DROP VALUE 'trial_end'")
    except Exception as e:
        print("Failed to downgrade enum type. Compatibility issue most likely. See error below:")
        print(e)
    # ### end Alembic commands ###

Error

qdt@controller-f764c6c6c-tcmp9:~$ flask db migrate
importing everything because I am not a defined process
2024-06-04 15:16:19,834  INFO [__init__  137] dash app registered at url_base_path: /dash/
2024-06-04 15:16:20,024  INFO [__init__  246] create_app::flask
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
ERROR [root] Error: Target database is not up to date.
qdt@controller-f764c6c6c-tcmp9:~$ flask db upgrade
importing everything because I am not a defined process
2024-06-04 15:16:44,344  INFO [__init__  137] dash app registered at url_base_path: /dash/
2024-06-04 15:16:44,535  INFO [__init__  246] create_app::flask
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> eb6e925edca4, Initial migration.
FINISHED CREATING ALL THE TABLES
INFO  [alembic.runtime.migration] Running upgrade eb6e925edca4 -> 5be8363eb9b1, user trial period
Traceback (most recent call last):
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1276, in _execute_context
    self.dialect.do_execute(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/default.py", line 593, in do_execute
    cursor.execute(statement, parameters)
psycopg2.errors.SerializationFailure: Heartbeat: Transaction ad56d692-68c2-436c-8a16-11fbb9b3c684 expired or aborted by a conflict: 40001: . Errors from tablet servers: [Operation expired (yb/tablet/transaction_coordinator.cc:1752): Heartbeat: Transaction ad56d692-68c2-436c-8a16-11fbb9b3c684 expired or aborted by a conflict: 40001 (pgsql error 40001) (transaction error 1)]

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/src/quantum_ml/.local/bin/flask", line 8, in <module>
    sys.exit(main())
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask/cli.py", line 988, in main
    cli.main()
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask/cli.py", line 579, in main
    return super().main(*args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/decorators.py", line 26, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask/cli.py", line 427, in decorator
    return __ctx.invoke(f, *args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask_migrate/cli.py", line 134, in upgrade
    _upgrade(directory, revision, sql, tag, x_arg)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask_migrate/__init__.py", line 96, in wrapped
    f(*args, **kwargs)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/flask_migrate/__init__.py", line 271, in upgrade
    command.upgrade(config, revision, sql=sql, tag=tag)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/command.py", line 298, in upgrade
    script.run_env()
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/script/base.py", line 489, in run_env
    util.load_python_file(self.dir, "env.py")
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/util/pyfiles.py", line 98, in load_python_file
    module = load_module_py(module_id, path)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/util/compat.py", line 184, in load_module_py
    spec.loader.exec_module(module)
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "migrations/env.py", line 96, in <module>
    run_migrations_online()
  File "migrations/env.py", line 90, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/runtime/environment.py", line 846, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/runtime/migration.py", line 529, in run_migrations
    head_maintainer.update_to_step(step)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/runtime/migration.py", line 763, in update_to_step
    self._update_version(from_, to_)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/runtime/migration.py", line 700, in _update_version
    ret = self.context.impl._exec(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/alembic/ddl/impl.py", line 141, in _exec
    return conn.execute(construct, *multiparams, **params)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1011, in execute
    return meth(self, multiparams, params)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/sql/elements.py", line 298, in _execute_on_connection
    return connection._execute_clauseelement(self, multiparams, params)
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1124, in _execute_clauseelement
    ret = self._execute_context(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1316, in _execute_context
    self._handle_dbapi_exception(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1510, in _handle_dbapi_exception
    util.raise_(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/util/compat.py", line 182, in raise_
    raise exception
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1276, in _execute_context
    self.dialect.do_execute(
  File "/usr/src/quantum_ml/.local/lib/python3.9/site-packages/sqlalchemy/engine/default.py", line 593, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.OperationalError: (psycopg2.errors.SerializationFailure) Heartbeat: Transaction ad56d692-68c2-436c-8a16-11fbb9b3c684 expired or aborted by a conflict: 40001: . Errors from tablet servers: [Operation expired (yb/tablet/transaction_coordinator.cc:1752): Heartbeat: Transaction ad56d692-68c2-436c-8a16-11fbb9b3c684 expired or aborted by a conflict: 40001 (pgsql error 40001) (transaction error 1)]

[SQL: UPDATE alembic_version SET version_num='5be8363eb9b1' WHERE alembic_version.version_num = 'eb6e925edca4']
(Background on this error at: http://sqlalche.me/e/13/e3q8)

Versions.

• OS: MacOS
• Python: 3.9
• Alembic: 1.4.3
• SQLAlchemy: 1.3.20
• Database: YugabyteDB
• DBAPI: psycopg2

Additional context The operation to update the alembic_version table takes less than 2 seconds, so it's unlikely that the issue is caused by a timeout. The issue persists even when using a fresh new database with no other connections.

Have a nice day!

[zzzeek]

hi -

Alembic is a tool that emits SQL commands to modify structures in a database and keep track of this within a versioning table. It emits these commands correctly, and the error you have has to do with your distributed database having other traffic coming to other nodes while this goes on. This is in the scope of managing that database, setting appropriate server/client config settings or ensuring that when you send commands to it, those commands are isolated from other changes. All fully outside the scope of what Alembic could ever do on its own, not to mention we dont have support for the yugabyte database. I would seek help from yugabyte developers.

[CaselIT]

Hi,

You may try to see if if setting the flag https://alembic.sqlalchemy.org/en/latest/api/runtime.html#alembic.runtime.environment.EnvironmentContext.configure.params.transaction_per_migration helps in your case