Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix airflow db check-migrations #19597

Merged
merged 1 commit into from
Nov 18, 2021
Merged

Fix airflow db check-migrations #19597

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions airflow/utils/db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand All @@ -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}"
ashb marked this conversation as resolved.
Show resolved Hide resolved
)
ticker += 1
time.sleep(1)
Expand Down
6 changes: 5 additions & 1 deletion tests/utils/test_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)