This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f90a678357c Fix FAB migration failing on MySQL under the pymysql
driver (#73220)
f90a678357c is described below
commit f90a678357ce829389023294efabdb9fa4d7eef9
Author: rjgoyln <[email protected]>
AuthorDate: Mon Sep 21 21:28:42 2026 +0800
Fix FAB migration failing on MySQL under the pymysql driver (#73220)
Revision 02ca36b0235b guards its MySQL DDL with throwaway stored
procedures, each emitted as a single multi-statement string. mysqlclient
turns CLIENT_MULTI_STATEMENTS on and accepts that; pymysql leaves it off,
so the server rejects the script and `airflow db migrate` dies part-way
through. MySQL DDL is not transactional, so the revision is left
unstamped over a partly-changed schema, and the Helm upgrade hook that
runs it reports only a timeout.
A live connection can answer the same existence questions by
introspection, which is what the rest of this revision already does for
indexes and foreign-key names. The stored procedures are only needed
offline, where there is no connection to ask.
---
.../0001_3_5_0_fix_fab_db_inconsistencies.py | 70 ++++++++++----
.../tests/unit/fab/auth_manager/models/test_db.py | 102 +++++++++++++++++++++
2 files changed, 153 insertions(+), 19 deletions(-)
diff --git
a/providers/fab/src/airflow/providers/fab/migrations/versions/0001_3_5_0_fix_fab_db_inconsistencies.py
b/providers/fab/src/airflow/providers/fab/migrations/versions/0001_3_5_0_fix_fab_db_inconsistencies.py
index 3913ae1b923..52c1046996d 100644
---
a/providers/fab/src/airflow/providers/fab/migrations/versions/0001_3_5_0_fix_fab_db_inconsistencies.py
+++
b/providers/fab/src/airflow/providers/fab/migrations/versions/0001_3_5_0_fix_fab_db_inconsistencies.py
@@ -49,7 +49,20 @@ _naming_convention = {
}
+def _get_live_bind():
+ """Return the live database connection, or ``None`` in offline (``--sql``)
mode."""
+ # Do not reach for op.get_bind() directly: in --sql mode alembic swaps the
connection for a
+ # MockConnection that writes to the output buffer, so it is never None
there and sa.inspect()
+ # on it raises NoInspectionAvailable. as_sql is the only reliable
discriminator.
+ return None if op.get_context().as_sql else op.get_bind()
+
+
def _mysql_run_procedure(procedure_name: str, body: str) -> str:
+ """Wrap ``body`` in a throwaway procedure so MySQL can guard DDL behind an
``IF``."""
+ # Offline (--sql) only. The result is a multi-statement script: drivers
that leave
+ # CLIENT_MULTI_STATEMENTS off (pymysql) reject it, and the `mysql` client
needs a DELIMITER
+ # around the procedure body. A live connection has nothing to guess at, so
the callers
+ # below introspect the schema and emit plain single statements instead.
return f"""
DROP PROCEDURE IF EXISTS {procedure_name};
CREATE PROCEDURE {procedure_name}()
@@ -147,17 +160,28 @@ def
_mysql_drop_unique_constraints_on_ab_register_user_email() -> str:
)
+def _find_unique_constraint_names(bind, table_name: str, column_name: str) ->
list[str]:
+ """Names of the unique constraints on ``table_name`` that cover
``column_name``."""
+ return [
+ uq["name"]
+ for uq in sa.inspect(bind).get_unique_constraints(table_name)
+ if uq["name"] is not None and column_name in uq["column_names"]
+ ]
+
+
def _drop_unique_constraint_if_exists(table_name: str, constraint_name: str)
-> None:
dialect_name = op.get_context().dialect.name
+ bind = _get_live_bind()
if dialect_name == "postgresql":
op.execute(sa.text(f'ALTER TABLE "{table_name}" DROP CONSTRAINT IF
EXISTS "{constraint_name}"'))
elif dialect_name == "mysql":
- op.execute(
- sa.text(
- _mysql_run_procedure(
- "DropUniqueIfExists",
- f"""
+ if bind is None:
+ op.execute(
+ sa.text(
+ _mysql_run_procedure(
+ "DropUniqueIfExists",
+ f"""
IF EXISTS (
SELECT 1
FROM information_schema.TABLE_CONSTRAINTS
@@ -172,9 +196,11 @@ def _drop_unique_constraint_if_exists(table_name: str,
constraint_name: str) ->
SELECT 1;
END IF;
""",
+ )
)
)
- )
+ elif any(uq["name"] == constraint_name for uq in
sa.inspect(bind).get_unique_constraints(table_name)):
+ op.drop_constraint(constraint_name, table_name, type_="unique")
else:
with op.batch_alter_table(table_name, schema=None) as batch_op:
with contextlib.suppress(ValueError):
@@ -192,13 +218,15 @@ def _resolve_fk_name(bind, table_name: str, column_name:
str, default: str) -> s
def _drop_index_if_exists(table_name: str, index_name: str) -> None:
dialect_name = op.get_context().dialect.name
+ bind = _get_live_bind()
if dialect_name == "mysql":
- op.execute(
- sa.text(
- _mysql_run_procedure(
- "DropIndexIfExists",
- f"""
+ if bind is None:
+ op.execute(
+ sa.text(
+ _mysql_run_procedure(
+ "DropIndexIfExists",
+ f"""
IF EXISTS (
SELECT 1
FROM information_schema.STATISTICS
@@ -210,16 +238,18 @@ def _drop_index_if_exists(table_name: str, index_name:
str) -> None:
DROP INDEX `{index_name}` ON `{table_name}`;
END IF;
""",
+ )
)
)
- )
+ elif any(idx["name"] == index_name for idx in
sa.inspect(bind).get_indexes(table_name)):
+ op.drop_index(index_name, table_name=table_name)
else:
op.drop_index(index_name, table_name=table_name, if_exists=True)
def upgrade() -> None:
dialect_name = op.get_context().dialect.name
- bind = op.get_bind()
+ bind = _get_live_bind()
if dialect_name == "postgresql":
op.create_index(
"idx_ab_user_username",
@@ -286,18 +316,20 @@ def upgrade() -> None:
)
# Drop any existing unique constraint on email, regardless of its name.
- # Raw SQL is used so this works in both online and offline (--sql) mode.
if dialect_name == "postgresql":
op.execute(sa.text(_postgresql_drop_unique_constraints_on_ab_register_user_email()))
elif dialect_name == "mysql":
-
op.execute(sa.text(_mysql_drop_unique_constraints_on_ab_register_user_email()))
+ if bind is None:
+
op.execute(sa.text(_mysql_drop_unique_constraints_on_ab_register_user_email()))
+ else:
+ for name in _find_unique_constraint_names(bind,
"ab_register_user", "email"):
+ op.drop_constraint(name, "ab_register_user", type_="unique")
elif dialect_name == "sqlite" and bind is not None:
# SQLite: batch mode rewrites the table; requires a live connection.
# Offline mode for SQLite is not supported by Airflow.
- for uq in sa.inspect(bind).get_unique_constraints("ab_register_user"):
- if "email" in uq["column_names"] and uq["name"] is not None:
- with op.batch_alter_table("ab_register_user", schema=None) as
batch_op:
- batch_op.drop_constraint(uq["name"], type_="unique")
+ for name in _find_unique_constraint_names(bind, "ab_register_user",
"email"):
+ with op.batch_alter_table("ab_register_user", schema=None) as
batch_op:
+ batch_op.drop_constraint(name, type_="unique")
with op.batch_alter_table("ab_register_user", schema=None) as batch_op:
batch_op.create_unique_constraint(batch_op.f("ab_register_user_email_uq"),
["email"])
diff --git a/providers/fab/tests/unit/fab/auth_manager/models/test_db.py
b/providers/fab/tests/unit/fab/auth_manager/models/test_db.py
index 452e5522fb7..50073d28126 100644
--- a/providers/fab/tests/unit/fab/auth_manager/models/test_db.py
+++ b/providers/fab/tests/unit/fab/auth_manager/models/test_db.py
@@ -21,11 +21,13 @@ from unittest import mock
import pytest
import sqlalchemy as sa
+from alembic import command
from alembic.autogenerate import compare_metadata
from alembic.migration import MigrationContext
from sqlalchemy import MetaData
import airflow.providers.fab as provider_fab
+from airflow import settings
from airflow.settings import engine
from airflow.utils.db import (
compare_server_default,
@@ -33,6 +35,16 @@ from airflow.utils.db import (
)
pytestmark = [pytest.mark.db_test]
+
+
+def _email_unique_constraint_names(bind) -> list[str]:
+ return sorted(
+ uq["name"]
+ for uq in sa.inspect(bind).get_unique_constraints("ab_register_user")
+ if "email" in uq["column_names"]
+ )
+
+
try:
from airflow.providers.fab.auth_manager.models.db import FABDBManager
@@ -118,6 +130,46 @@ try:
with pytest.raises(SystemExit, match="Offline migration not
supported for SQLite"):
FABDBManager(session).upgradedb(from_revision=None,
to_revision=None, show_sql_only=True)
+ @pytest.mark.parametrize(
+ ("direction", "revision_range", "expected_statements"),
+ [
+ (
+ "upgrade",
+ "6709f7a774b9:02ca36b0235b",
+ (
+ "CREATE PROCEDURE CreateIdxPermissionViewId()",
+ "CREATE PROCEDURE CreateIdxRoleId()",
+ "CREATE PROCEDURE DropEmailUqIfExists()",
+ "UPDATE alembic_version_fab SET
version_num='02ca36b0235b'",
+ ),
+ ),
+ (
+ "downgrade",
+ "02ca36b0235b:6709f7a774b9",
+ (
+ "CREATE PROCEDURE DropUniqueIfExists()",
+ "CREATE PROCEDURE DropIndexIfExists()",
+ "UPDATE alembic_version_fab SET
version_num='6709f7a774b9'",
+ ),
+ ),
+ ],
+ )
+ @mock.patch("airflow.settings.SQL_ALCHEMY_CONN",
"mysql+pymysql://user:pass@host/airflow")
+ def test_offline_mysql_sql_generation(
+ self, session, capsys, direction, revision_range,
expected_statements
+ ):
+ # Offline mode hands the migration a MockConnection rather than
None, so every
+ # introspection has to be gated on the context's as_sql flag or it
raises
+ # NoInspectionAvailable and no script is produced at all. The
version stamp is
+ # emitted last, so asserting it proves the whole revision ran.
+ config = FABDBManager(session=session).get_alembic_config()
+
+ getattr(command, direction)(config, revision_range, sql=True)
+
+ script = capsys.readouterr().out
+ for statement in expected_statements:
+ assert statement in script
+
@mock.patch("alembic.command.upgrade")
@mock.patch.object(FABDBManager, "create_db_from_orm")
@mock.patch.object(FABDBManager, "_has_existing_manager_tables",
return_value=False)
@@ -292,5 +344,55 @@ try:
if original_revision and current_revision != original_revision:
manager.upgradedb(to_revision=original_revision)
+ @pytest.mark.backend("mysql")
+ def test_upgradedb_and_downgrade_mysql_run_under_pymysql(self,
session):
+ # pymysql leaves CLIENT_MULTI_STATEMENTS off (mysqlclient, which
CI uses, turns it
+ # on), so the server rejects any revision that packs several
statements into one
+ # op.execute(). Drive the revision over a pymysql connection to
keep it split.
+ pytest.importorskip("pymysql")
+
+ manager = FABDBManager(session=session)
+ original_revision = manager.get_current_revision()
+ pymysql_url =
sa.engine.make_url(settings.SQL_ALCHEMY_CONN).set(drivername="mysql+pymysql")
+ pymysql_engine = sa.create_engine(pymysql_url)
+
+ try:
+ manager.downgrade(to_revision="6709f7a774b9")
+
+ # An Airflow 2.x database reaches this revision with FAB's own
unique index on
+ # email, named after the column. That is what the revision has
to find and drop.
+ with pymysql_engine.begin() as setup:
+ for name in _email_unique_constraint_names(setup):
+ setup.execute(sa.text(f"ALTER TABLE `ab_register_user`
DROP INDEX `{name}`"))
+ setup.execute(sa.text("ALTER TABLE `ab_register_user` ADD
UNIQUE KEY `email` (`email`)"))
+
+ config = manager.get_alembic_config()
+ with pymysql_engine.connect() as connection:
+ config.attributes["connection"] = connection
+ assert connection.dialect.driver == "pymysql"
+
+ command.upgrade(config, revision="02ca36b0235b")
+ assert _email_unique_constraint_names(connection) ==
["ab_register_user_email_uq"]
+
+ command.downgrade(config, revision="6709f7a774b9")
+ assert _email_unique_constraint_names(connection) == []
+ index_names = {
+ index["name"]
+ for index in
sa.inspect(connection).get_indexes("ab_permission_view_role")
+ }
+ assert "idx_permission_view_id" not in index_names
+ assert "idx_role_id" not in index_names
+
+ # The downgrade left no unique constraint on email at all,
so this second
+ # upgrade runs the drops against objects that are already
gone — the state a
+ # database is in after the revision failed part-way
through.
+ command.upgrade(config, revision="02ca36b0235b")
+ assert _email_unique_constraint_names(connection) ==
["ab_register_user_email_uq"]
+ finally:
+ pymysql_engine.dispose()
+ current_revision = manager.get_current_revision()
+ if original_revision and current_revision != original_revision:
+ manager.upgradedb(to_revision=original_revision)
+
except ModuleNotFoundError:
pass
