diff --git a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py index 4aa03243605a5..629d9862164d7 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py +++ b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py @@ -464,6 +464,15 @@ def get_extra_menu_items(self, *, user: T) -> list[ExtraMenuItem]: """ return [] + @staticmethod + def get_db_manager() -> str | None: + """ + Specify the DB manager path needed to run the auth manager. + + This is optional and not all auth managers require a DB manager. + """ + return None + @classmethod @cache def _get_token_signer( diff --git a/airflow-core/src/airflow/utils/db_manager.py b/airflow-core/src/airflow/utils/db_manager.py index 46c4323dde539..f18528c33fc60 100644 --- a/airflow-core/src/airflow/utils/db_manager.py +++ b/airflow-core/src/airflow/utils/db_manager.py @@ -23,6 +23,7 @@ from sqlalchemy import inspect from airflow import settings +from airflow.api_fastapi.app import create_auth_manager from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.utils.log.logging_mixin import LoggingMixin @@ -144,6 +145,12 @@ def __init__(self): super().__init__() self._managers: list[BaseDBManager] = [] managers = conf.get("database", "external_db_managers").split(",") + + # Add DB manager specified by auth manager (if any) + 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) + for module in managers: manager = import_string(module) self._managers.append(manager) diff --git a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py index 8b7c5d1f8cc55..741c4962f3467 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py +++ b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py @@ -171,6 +171,9 @@ def test_get_url_logout_return_none(self, auth_manager): def test_get_extra_menu_items_return_empty_list(self, auth_manager): assert auth_manager.get_extra_menu_items(user=BaseAuthManagerUserTest(name="test")) == [] + def test_get_db_manager_return_none(self, auth_manager): + assert auth_manager.get_db_manager() is None + @patch.object(EmptyAuthManager, "filter_authorized_menu_items") def test_get_authorized_menu_items(self, mock_filter_authorized_menu_items, auth_manager): user = BaseAuthManagerUserTest(name="test") diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py index aa99ceda24c55..f0c4eefab776b 100644 --- a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py +++ b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py @@ -521,6 +521,10 @@ def get_extra_menu_items(self, *, user: User) -> list[ExtraMenuItem]: if self._is_authorized(method="MENU", resource_type=item["resource_type"], user=user) ] + @staticmethod + def get_db_manager() -> str | None: + return "airflow.providers.fab.auth_manager.models.db.FABDBManager" + def _is_authorized( self, *, diff --git a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py index 4e740bae0581c..d6db61c4492e4 100644 --- a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py +++ b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py @@ -629,3 +629,7 @@ def test_get_extra_menu_items(self, _, auth_manager_with_appbuilder, flask_app): result = auth_manager_with_appbuilder.get_extra_menu_items(user=Mock()) assert len(result) == 5 assert all(item.href.startswith(AUTH_MANAGER_FASTAPI_APP_PREFIX) for item in result) + + def test_get_db_manager(self, auth_manager): + result = auth_manager.get_db_manager() + assert result == "airflow.providers.fab.auth_manager.models.db.FABDBManager"