This is an automated email from the ASF dual-hosted git repository. rahulvats pushed a commit to branch py-client-sync in repository https://gitbox.apache.org/repos/asf/airflow.git
commit de360fc61b6309d5b29efe776f1534285764ce1f Author: Subham <[email protected]> AuthorDate: Tue Mar 24 19:01:46 2026 +0530 Fix FAB DB manager discovery in migration-only contexts (#64145) --- airflow-core/src/airflow/utils/db_manager.py | 33 +++++++++--- airflow-core/tests/unit/utils/test_db_manager.py | 68 +++++++++++++++++++++++- 2 files changed, 93 insertions(+), 8 deletions(-) diff --git a/airflow-core/src/airflow/utils/db_manager.py b/airflow-core/src/airflow/utils/db_manager.py index 2b1feff7bf9..ef0f515a945 100644 --- a/airflow-core/src/airflow/utils/db_manager.py +++ b/airflow-core/src/airflow/utils/db_manager.py @@ -196,16 +196,17 @@ class RunDBManager(LoggingMixin): """ def __init__(self): - from airflow.api_fastapi.app import create_auth_manager from airflow.providers_manager import ProvidersManager super().__init__() self._managers: list[BaseDBManager] = [] - # Start with auto-discovered DB managers from installed providers + # Start with auto-discovered DB managers from installed providers. + # ProvidersManager reads the ``db-managers`` key from each provider's + # get_provider_info() and is the primary source of truth. managers: list[str] = list(ProvidersManager().db_managers) - # Add any explicitly configured managers not already discovered + # Add any explicitly configured managers not already discovered. managers_config = conf.get("database", "external_db_managers", fallback=None) if managers_config: for m in managers_config.split(","): @@ -213,10 +214,28 @@ class RunDBManager(LoggingMixin): if stripped not in managers: managers.append(stripped) - # Add DB manager declared by the configured auth manager (existing behavior, deduplicated) - auth_manager_db_manager = create_auth_manager().get_db_manager() - if auth_manager_db_manager and auth_manager_db_manager not in managers: - managers.append(auth_manager_db_manager) + # Add the DB manager declared by the configured auth manager as a + # final fallback for backward compatibility. + # This is wrapped in a try/except because in migration-only contexts + # (e.g. the Helm migrateDatabaseJob) the auth manager may not be fully + # initializable — a Flask app context or other runtime state may be + # absent. A failure here must not silently drop the auth manager's DB + # manager from the migration list; ProvidersManager discovery above is + # the reliable path in those contexts. + try: + from airflow.api_fastapi.app import create_auth_manager + + auth_manager_db_manager = create_auth_manager().get_db_manager() + if auth_manager_db_manager and auth_manager_db_manager not in managers: + managers.append(auth_manager_db_manager) + except Exception: + self.log.debug( + "Could not retrieve DB manager from auth manager during RunDBManager " + "initialisation. This is expected in migration-only contexts where the " + "auth manager cannot be fully initialised. DB managers discovered via " + "ProvidersManager will still be used.", + exc_info=True, + ) for module in managers: manager = import_string(module.strip()) diff --git a/airflow-core/tests/unit/utils/test_db_manager.py b/airflow-core/tests/unit/utils/test_db_manager.py index efd417b12f4..61716e11c7d 100644 --- a/airflow-core/tests/unit/utils/test_db_manager.py +++ b/airflow-core/tests/unit/utils/test_db_manager.py @@ -230,10 +230,76 @@ class TestRunDBManager: def test_initdb_and_upgradedb_pass_use_migration_files_to_var_kwarg_manager(self, session): VarKwargExternalManager.initdb_kwargs = [] VarKwargExternalManager.upgradedb_kwargs = [] - run_db_manager = _create_run_db_manager(VarKwargExternalManager) run_db_manager.initdb(session=session, use_migration_files=True) run_db_manager.upgradedb(session=session, use_migration_files=False) assert VarKwargExternalManager.initdb_kwargs == [{"use_migration_files": True}] assert VarKwargExternalManager.upgradedb_kwargs == [{"use_migration_files": False}] + + @mock.patch("airflow.utils.db_manager.import_string") + def test_run_db_manager_uses_providers_manager_when_auth_manager_fails(self, mock_import): + """ + When create_auth_manager() raises in a migration-only context (e.g. the Helm + migrateDatabaseJob), RunDBManager must still load managers discovered via + ProvidersManager — the exception must not silently drop all DB managers. + """ + sentinel = object() + mock_import.return_value = sentinel + + with ( + mock.patch( + "airflow.providers_manager.ProvidersManager", + ) as mock_pm, + mock.patch( + "airflow.api_fastapi.app.create_auth_manager", + side_effect=RuntimeError("No app context"), + ), + ): + mock_pm.return_value.db_managers = ["airflow.providers.fab.auth_manager.models.db.FABDBManager"] + with mock.patch("airflow.utils.db_manager.conf") as mock_conf: + mock_conf.get.return_value = None + rdm = RunDBManager.__new__(RunDBManager) + # Manually call __init__ so we control all side-effects + RunDBManager.__init__(rdm) + + # The sentinel class returned by import_string must be in _managers + assert sentinel in rdm._managers + + @mock.patch("airflow.utils.db_manager.import_string") + def test_run_db_manager_includes_auth_manager_db_manager_when_available(self, mock_import): + """ + When create_auth_manager() succeeds and returns a DB manager class name not + already in the ProvidersManager list, it must be appended to _managers. + """ + sentinel_pm = object() + sentinel_am = object() + call_order = [] + + def _import(path): + if "fab" in path: + call_order.append("fab") + return sentinel_pm + call_order.append("auth_manager_extra") + return sentinel_am + + mock_import.side_effect = _import + + mock_am = mock.MagicMock() + mock_am.get_db_manager.return_value = "some.extra.AuthManagerDBManager" + + with ( + mock.patch("airflow.providers_manager.ProvidersManager") as mock_pm, + mock.patch( + "airflow.api_fastapi.app.create_auth_manager", + return_value=mock_am, + ), + ): + mock_pm.return_value.db_managers = ["airflow.providers.fab.auth_manager.models.db.FABDBManager"] + with mock.patch("airflow.utils.db_manager.conf") as mock_conf: + mock_conf.get.return_value = None + rdm = RunDBManager.__new__(RunDBManager) + RunDBManager.__init__(rdm) + + assert sentinel_pm in rdm._managers + assert sentinel_am in rdm._managers
