You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/05/03 18:12:12 UTC
(airflow) branch v3-2-test updated: [v3-2-test] Fix flaky test_migration_stairway by retrying each migration step (#66162) (#66299)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-2-test by this push:
new 8a0c50906de [v3-2-test] Fix flaky test_migration_stairway by retrying each migration step (#66162) (#66299)
8a0c50906de is described below
commit 8a0c50906de5f7ef58710de70657f603007f15f2
Author: Jason(Zhe-You) Liu <68...@users.noreply.github.com>
AuthorDate: Mon May 4 02:12:03 2026 +0800
[v3-2-test] Fix flaky test_migration_stairway by retrying each migration step (#66162) (#66299)
(cherry picked from commit 90bec8cb5af4a126d53008b49d1b14c21046c02e)
---
.../tests/unit/migrations/test_migration_utils.py | 227 ++++++++++++++++++++-
1 file changed, 219 insertions(+), 8 deletions(-)
diff --git a/airflow-core/tests/unit/migrations/test_migration_utils.py b/airflow-core/tests/unit/migrations/test_migration_utils.py
index d31d4d930d2..1ab309421fe 100644
--- a/airflow-core/tests/unit/migrations/test_migration_utils.py
+++ b/airflow-core/tests/unit/migrations/test_migration_utils.py
@@ -31,9 +31,23 @@ immediately identifiable in CI logs.
from __future__ import annotations
+import time
+from collections.abc import Callable
+from functools import partial
+
import pytest
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
from alembic.script import ScriptDirectory
+from sqlalchemy import text
+from sqlalchemy.exc import MissingGreenlet, OperationalError
+from airflow import settings
+from airflow.migrations.utils import (
+ disable_sqlite_fkeys,
+ ignore_sqlite_value_error,
+ mysql_drop_foreignkey_if_exists,
+)
from airflow.utils.db import (
_REVISION_HEADS_MAP,
_get_alembic_config,
@@ -49,6 +63,56 @@ pytestmark = pytest.mark.db_test
# 2.10.3, and keeps the number of steps manageable on slower backends.
_STAIRWAY_START_REVISION = _REVISION_HEADS_MAP["3.0.0"]
+# CI shows two transient failure modes for this test:
+# 1. PostgreSQL "too many clients already" when the metadata DB pool is
+# saturated by other parallel tests at the moment a step runs. This
+# surfaces as ``sqlalchemy.exc.OperationalError``.
+# 2. SQLAlchemy ``MissingGreenlet`` during async-engine disposal between
+# steps, which can leak into the next migration call.
+# Both are environmental; rebuilding the ORM (dispose + reconfigure with
+# NullPool) between attempts clears the leaked connections and restores
+# ``settings.Session`` so the next ``@provide_session``-wrapped call has a
+# working session.
+_STEP_RETRIES = 3
+_STEP_RETRY_BACKOFF_SECONDS = 2.0
+_TRANSIENT_STEP_ERRORS: tuple[type[BaseException], ...] = (OperationalError, MissingGreenlet)
+
+
+def _run_step_with_retry(step: Callable[[], None], *, revision_id: str, label: str) -> None:
+ """
+ Run a single migration step, retrying on known transient backend errors.
+
+ Only retries on :class:`sqlalchemy.exc.OperationalError` (e.g. PostgreSQL
+ ``too many clients already``) and :class:`sqlalchemy.exc.MissingGreenlet`
+ (async-engine disposal leaking into the next call) so genuine migration
+ regressions still surface immediately.
+
+ Rebuilds the ORM between attempts: a plain ``dispose_orm`` would leave
+ ``settings.Session`` set to ``None``, causing the next
+ ``@provide_session``-wrapped call to raise
+ ``RuntimeError("Session must be set before!")``. NullPool is reapplied so
+ the fixture's connection-leak guard is preserved across retries.
+ """
+ last_exc: BaseException | None = None
+ for attempt in range(1, _STEP_RETRIES + 1):
+ try:
+ step()
+ return
+ except _TRANSIENT_STEP_ERRORS as exc:
+ last_exc = exc
+ if attempt == _STEP_RETRIES:
+ break
+ # Free pooled connections and rebuild Session/engine before retry.
+ try:
+ settings.dispose_orm(do_log=False)
+ settings.configure_orm(disable_connection_pool=True)
+ except Exception:
+ pass
+ time.sleep(_STEP_RETRY_BACKOFF_SECONDS * attempt)
+ raise AssertionError(
+ f"Stairway test failed at revision {revision_id!r} during {label} after {_STEP_RETRIES} attempts"
+ ) from last_exc
+
def _get_revisions_in_order() -> list[str]:
"""Return revision IDs in upgrade order, starting after the stairway base."""
@@ -106,13 +170,160 @@ def test_migration_stairway(stairway_db) -> None:
for revision_id in revisions:
try:
- # Step 1: upgrade to this revision.
- upgradedb(to_revision=revision_id)
+ _run_step_with_retry(
+ partial(upgradedb, to_revision=revision_id),
+ revision_id=revision_id,
+ label="upgrade",
+ )
+ _run_step_with_retry(
+ partial(downgrade, to_revision="-1"),
+ revision_id=revision_id,
+ label="downgrade",
+ )
+ _run_step_with_retry(
+ partial(upgradedb, to_revision=revision_id),
+ revision_id=revision_id,
+ label="re-upgrade",
+ )
+ finally:
+ # upgradedb()/downgrade() each enter ``_single_connection_pool``,
+ # which calls ``reconfigure_orm()`` and resets the pool to default.
+ # Re-apply NullPool so the iteration cap is preserved.
+ settings.dispose_orm(do_log=False)
+ settings.configure_orm(disable_connection_pool=True)
+
+
+class TestDisableSqliteFkeys:
+ """Tests for :func:`disable_sqlite_fkeys`."""
+
+ def test_yields_op(self):
+ """The context manager yields the *op* object on every backend."""
+ with settings.engine.connect() as conn:
+ ctx = MigrationContext.configure(conn)
+ op = Operations(ctx)
+ with disable_sqlite_fkeys(op) as yielded:
+ assert yielded is op
+
+ @pytest.mark.backend("sqlite")
+ def test_sqlite_turns_off_fkeys(self):
+ """On SQLite, ``PRAGMA foreign_keys`` is OFF inside the context."""
+
+ with settings.engine.connect() as conn:
+ # Ensure foreign_keys is ON before entering the context manager.
+ conn.execute(text("PRAGMA foreign_keys=ON"))
+ conn.commit()
+
+ ctx = MigrationContext.configure(conn)
+ op = Operations(ctx)
+
+ with disable_sqlite_fkeys(op):
+ result = conn.execute(text("PRAGMA foreign_keys")).scalar()
+ assert result == 0, "foreign_keys should be OFF inside disable_sqlite_fkeys"
+
+ result = conn.execute(text("PRAGMA foreign_keys")).scalar()
+ assert result == 1, "foreign_keys should be restored to ON after disable_sqlite_fkeys"
+
+ @pytest.mark.backend("postgres", "mysql")
+ def test_non_sqlite_is_noop(self):
+ """On non-SQLite backends, the context manager simply yields."""
+
+ with settings.engine.connect() as conn:
+ ctx = MigrationContext.configure(conn)
+ op = Operations(ctx)
+ with disable_sqlite_fkeys(op) as yielded:
+ assert yielded is op
+
+
+class TestMysqlDropForeignkeyIfExists:
+ """Tests for :func:`mysql_drop_foreignkey_if_exists`."""
+
+ @pytest.mark.backend("mysql")
+ def test_drops_existing_fk(self):
+ """On MySQL, an existing foreign key is dropped."""
+
+ with settings.engine.connect() as conn:
+ conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_parent (id INT PRIMARY KEY)"))
+ conn.execute(
+ text(
+ """
+ CREATE TABLE IF NOT EXISTS _test_mig_child (
+ id INT PRIMARY KEY,
+ parent_id INT,
+ CONSTRAINT _test_mig_fk FOREIGN KEY (parent_id)
+ REFERENCES _test_mig_parent(id)
+ )
+ """
+ )
+ )
+ conn.commit()
+
+ try:
+ ctx = MigrationContext.configure(conn)
+ op = Operations(ctx)
+ mysql_drop_foreignkey_if_exists("_test_mig_fk", "_test_mig_child", op)
+ conn.commit()
+
+ count = conn.execute(
+ text(
+ """
+ SELECT COUNT(*) FROM information_schema.TABLE_CONSTRAINTS
+ WHERE CONSTRAINT_SCHEMA = DATABASE()
+ AND TABLE_NAME = '_test_mig_child'
+ AND CONSTRAINT_NAME = '_test_mig_fk'
+ AND CONSTRAINT_TYPE = 'FOREIGN KEY'
+ """
+ )
+ ).scalar()
+ assert count == 0, "Foreign key should have been dropped"
+ finally:
+ conn.execute(text("DROP TABLE IF EXISTS _test_mig_child"))
+ conn.execute(text("DROP TABLE IF EXISTS _test_mig_parent"))
+ conn.commit()
+
+ @pytest.mark.backend("mysql")
+ def test_noop_for_nonexistent_fk(self):
+ """On MySQL, dropping a non-existent FK does not raise."""
+
+ with settings.engine.connect() as conn:
+ conn.execute(text("CREATE TABLE IF NOT EXISTS _test_mig_nofk (id INT PRIMARY KEY)"))
+ conn.commit()
+
+ try:
+ ctx = MigrationContext.configure(conn)
+ op = Operations(ctx)
+ mysql_drop_foreignkey_if_exists("nonexistent_fk", "_test_mig_nofk", op)
+ conn.commit()
+ finally:
+ conn.execute(text("DROP TABLE IF EXISTS _test_mig_nofk"))
+ conn.commit()
+
+
+class TestIgnoreSqliteValueError:
+ """Tests for :func:`ignore_sqlite_value_error`."""
+
+ @pytest.mark.backend("sqlite")
+ def test_sqlite_suppresses_value_error(self):
+ """On SQLite, ``ValueError`` is suppressed."""
+
+ with settings.engine.connect() as conn:
+ ctx = MigrationContext.configure(conn)
+ with Operations.context(ctx), ignore_sqlite_value_error():
+ raise ValueError("should be suppressed on SQLite")
+
+ @pytest.mark.backend("postgres", "mysql")
+ def test_non_sqlite_does_not_suppress(self):
+ """On non-SQLite backends, ``ValueError`` propagates."""
- # Step 2: downgrade exactly one step back.
- downgrade(to_revision="-1")
+ with settings.engine.connect() as conn:
+ ctx = MigrationContext.configure(conn)
+ with Operations.context(ctx), pytest.raises(ValueError, match="should propagate"):
+ with ignore_sqlite_value_error():
+ raise ValueError("should propagate")
- # Step 3: re-apply so the next iteration starts from the right state.
- upgradedb(to_revision=revision_id)
- except Exception as e:
- raise AssertionError(f"Stairway test failed at revision {revision_id!r}") from e
+ def test_does_not_suppress_other_exceptions(self):
+ """Even on SQLite, only ``ValueError`` is suppressed — other exceptions propagate."""
+ with settings.engine.connect() as conn:
+ ctx = MigrationContext.configure(conn)
+ with Operations.context(ctx), pytest.raises(TypeError):
+ with ignore_sqlite_value_error():
+ raise TypeError("not a ValueError")