Fix airflow db check-migrations (#19597)

airflow/utils/db.py

@@ -627,13 +627,17 @@ def check_migrations(timeout):
     :param timeout: Timeout for the migration in seconds
     :return: None
     """
-    from alembic.runtime.migration import MigrationContext
+    from alembic.runtime.environment import EnvironmentContext
     from alembic.script import ScriptDirectory
 
     config = _get_alembic_config()
     script_ = ScriptDirectory.from_config(config)
-    with settings.engine.connect() as connection:
-        context = MigrationContext.configure(connection)
+    with EnvironmentContext(
+        config,
+        script_,
+    ) as env, settings.engine.connect() as connection:
+        env.configure(connection)
+        context = env.get_context()
         ticker = 0
         while True:
             source_heads = set(script_.get_heads())
@@ -642,8 +646,8 @@ def check_migrations(timeout):
                 break
             if ticker >= timeout:
                 raise TimeoutError(
-                    f"There are still unapplied migrations after {ticker} seconds. "
-                    f"Migration Head(s) in DB: {db_heads} | Migration Head(s) in Source Code: {source_heads}"
+                    f"There are still unapplied migrations after {ticker} seconds. Migration"
+                    f"Head(s) in DB: {db_heads} | Migration Head(s) in Source Code: {source_heads}"
                 )
             ticker += 1
             time.sleep(1)

tests/utils/test_db.py

@@ -29,7 +29,7 @@
 
 from airflow.models import Base as airflow_base
 from airflow.settings import engine
-from airflow.utils.db import create_default_connections
+from airflow.utils.db import check_migrations, create_default_connections
 
 
 class TestDb(unittest.TestCase):
@@ -103,3 +103,7 @@ def test_default_connections_sort(self):
         source = inspect.getsource(create_default_connections)
         src = pattern.findall(source)
         assert sorted(src) == src
+
+    def test_check_migrations(self):
+        # Should run without error. Can't easily test the behaviour, but we can check it works
+        check_migrations(1)