diff --git a/airflow-core/src/airflow/utils/db_manager.py b/airflow-core/src/airflow/utils/db_manager.py index ef0f515a9452f..8c4eb11a0d4fe 100644 --- a/airflow-core/src/airflow/utils/db_manager.py +++ b/airflow-core/src/airflow/utils/db_manager.py @@ -20,7 +20,6 @@ import os from typing import TYPE_CHECKING -from alembic import command from sqlalchemy import inspect from airflow import settings @@ -50,6 +49,20 @@ def _callable_accepts_use_migration_files(callable_) -> bool: ) +def _get_alembic_command(): + from alembic import command + + return command + + +class _AlembicCommandProxy: + def __getattr__(self, name): + return getattr(_get_alembic_command(), name) + + +command = _AlembicCommandProxy() + + class BaseDBManager(LoggingMixin): """Abstract Base DB manager for external DBs.""" diff --git a/airflow-core/tests/unit/utils/test_db_manager.py b/airflow-core/tests/unit/utils/test_db_manager.py index 61716e11c7d2b..fa57fd9554ad6 100644 --- a/airflow-core/tests/unit/utils/test_db_manager.py +++ b/airflow-core/tests/unit/utils/test_db_manager.py @@ -16,6 +16,8 @@ # under the License. from __future__ import annotations +import subprocess +import sys from contextlib import nullcontext from unittest import mock @@ -97,6 +99,28 @@ def _create_run_db_manager(*managers): class TestBaseDBManager: + def test_importing_db_manager_does_not_eagerly_import_alembic(self): + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys\n" + "import airflow.utils.db_manager as module\n" + "assert hasattr(module, 'RunDBManager')\n" + "alembic_modules = sorted(\n" + " name for name in sys.modules if name == 'alembic' or name.startswith('alembic.')\n" + ")\n" + "assert not alembic_modules, alembic_modules\n" + ), + ], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr or result.stdout + @mock.patch.object(BaseDBManager, "get_alembic_config") @mock.patch.object(BaseDBManager, "get_current_revision") @mock.patch.object(BaseDBManager, "create_db_from_orm")