Add option in auth manager interface to define FastAPI api (#45009)

airflow/api_connexion/endpoints/asset_endpoint.py

@@ -43,14 +43,14 @@
     queued_event_collection_schema,
     queued_event_schema,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.assets.manager import asset_manager
 from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel
 from airflow.utils import timezone
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.db import get_query_count
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.www.decorators import action_logging
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/dag_endpoint.py

@@ -37,14 +37,14 @@
     dag_schema,
     dags_collection_schema,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.exceptions import AirflowException, DagNotFound
 from airflow.models.dag import DagModel, DagTag
 from airflow.utils.airflow_flask_app import get_airflow_app
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.db import get_query_count
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.www.decorators import action_logging
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/dag_parsing.py

@@ -26,12 +26,12 @@
 
 from airflow.api_connexion import security
 from airflow.api_connexion.exceptions import NotFound, PermissionDenied
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagDetails
 from airflow.models.dag import DagModel
 from airflow.models.dagbag import DagPriorityParsingRequest
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/dag_run_endpoint.py

@@ -59,6 +59,7 @@
     TaskInstanceReferenceCollection,
     task_instance_reference_collection_schema,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagAccessEntity
 from airflow.exceptions import ParamValidationError
 from airflow.models import DagModel, DagRun
@@ -71,7 +72,6 @@
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 from airflow.www.decorators import action_logging
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/dag_source_endpoint.py

@@ -26,12 +26,12 @@
 from airflow.api_connexion import security
 from airflow.api_connexion.exceptions import NotFound, PermissionDenied
 from airflow.api_connexion.schemas.dag_source_schema import dag_source_schema
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagAccessEntity, DagDetails
 from airflow.models.dag import DagModel
 from airflow.models.dag_version import DagVersion
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/dag_stats_endpoint.py

@@ -25,12 +25,12 @@
 from airflow.api_connexion.schemas.dag_stats_schema import (
     dag_stats_collection_schema,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagAccessEntity
 from airflow.models.dag import DagRun
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import DagRunState
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/import_error_endpoint.py

@@ -29,12 +29,12 @@
     import_error_collection_schema,
     import_error_schema,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import AccessView, DagDetails
 from airflow.models.dag import DagModel
 from airflow.models.errors import ParseImportError
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/task_instance_endpoint.py

@@ -47,6 +47,7 @@
     task_instance_schema,
 )
 from airflow.api_connexion.security import get_readable_dags
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagAccessEntity, DagDetails
 from airflow.exceptions import TaskNotFound
 from airflow.models.dagrun import DagRun as DR
@@ -58,7 +59,6 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.state import DagRunState, TaskInstanceState
 from airflow.www.decorators import action_logging
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/endpoints/xcom_endpoint.py

@@ -31,13 +31,13 @@
     xcom_schema_native,
     xcom_schema_string,
 )
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import DagAccessEntity
 from airflow.models import DagRun as DR, XCom
 from airflow.settings import conf
 from airflow.utils.api_migration import mark_fastapi_migration_done
 from airflow.utils.db import get_query_count
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

airflow/api_connexion/security.py

@@ -22,6 +22,7 @@
 from flask import Response, g
 
 from airflow.api_connexion.exceptions import PermissionDenied, Unauthenticated
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import (
     AccessView,
     AssetDetails,
@@ -33,7 +34,6 @@
     VariableDetails,
 )
 from airflow.utils.airflow_flask_app import get_airflow_app
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from airflow.auth.managers.base_auth_manager import ResourceMethod

airflow/api_fastapi/app.py

@@ -68,9 +68,9 @@ def create_app(apps: str = "all") -> FastAPI:
         init_dag_bag(app)
         init_views(app)
         init_plugins(app)
+        init_auth_manager(app)
         init_flask_plugins(app)
         init_error_handlers(app)
-        init_auth_manager()
 
     if "execution" in apps_list or "all" in apps_list:
         task_exec_api_app = create_task_execution_api_app(app)
@@ -112,34 +112,28 @@ def get_auth_manager_cls() -> type[BaseAuthManager]:
     return auth_manager_cls
 
 
-def init_auth_manager() -> BaseAuthManager:
-    """
-    Initialize the auth manager.
-
-    Import the user manager class and instantiate it.
-    """
+def create_auth_manager() -> BaseAuthManager:
+    """Create the auth manager."""
     global auth_manager
     auth_manager_cls = get_auth_manager_cls()
     auth_manager = auth_manager_cls()
-    auth_manager.init()
     return auth_manager
 
 
+def init_auth_manager(app: FastAPI | None = None) -> BaseAuthManager:
+    """Initialize the auth manager."""
+    am = create_auth_manager()
+    am.init()
+
+    if app and (auth_manager_fastapi_app := am.get_fastapi_app()):
+        app.mount("/auth", auth_manager_fastapi_app)
+
+    return am
+
+
 def get_auth_manager() -> BaseAuthManager:
     """Return the auth manager, provided it's been initialized before."""
     global auth_manager
-    if auth_manager is None:
-        """
-        The auth manager can be init in the main Flask application but also in the mini Flask application
-        in Fab provider.
-        This is temporary, the goal is to remove the main Flask application from core Airflow. Once that done,
-        we'll be able to remove this if because the auth manager will be only init in the min Flask
-        application defined in Fab provider.
-        """
-        from airflow.www.extensions.init_auth_manager import get_auth_manager as get_auth_manager_flask
-
-        if auth_manager_flask := get_auth_manager_flask():
-            auth_manager = auth_manager_flask
 
     if auth_manager is None:
         raise RuntimeError(

airflow/auth/managers/base_auth_manager.py

@@ -36,6 +36,7 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
+    from fastapi import FastAPI
     from flask import Blueprint
     from sqlalchemy.orm import Session
 
@@ -55,7 +56,6 @@
         VariableDetails,
     )
     from airflow.cli.cli_config import CLICommand
-    from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
     from airflow.www.security_manager import AirflowSecurityManagerV2
 
 ResourceMethod = Literal["GET", "POST", "PUT", "DELETE", "MENU"]
@@ -68,14 +68,8 @@ class BaseAuthManager(Generic[T], LoggingMixin):
     Class to derive in order to implement concrete auth managers.
 
     Auth managers are responsible for any user management related operation such as login, logout, authz, ...
-
-    :param appbuilder: the flask app builder
     """
 
-    def __init__(self, appbuilder: AirflowAppBuilder | None = None) -> None:
-        super().__init__()
-        self.appbuilder = appbuilder
-
     def init(self) -> None:
         """
         Run operations when Airflow is initializing.
@@ -440,7 +434,7 @@ def security_manager(self) -> AirflowSecurityManagerV2:
         """
         from airflow.www.security_manager import AirflowSecurityManagerV2
 
-        return AirflowSecurityManagerV2(self.appbuilder)
+        return AirflowSecurityManagerV2(getattr(self, "appbuilder"))
 
     @staticmethod
     def get_cli_commands() -> list[CLICommand]:
@@ -453,6 +447,15 @@ def get_cli_commands() -> list[CLICommand]:
 
     def get_api_endpoints(self) -> None | Blueprint:
         """Return API endpoint(s) definition for the auth manager."""
+        # TODO: Remove this method when legacy Airflow 2 UI is gone
+        return None
+
+    def get_fastapi_app(self) -> FastAPI | None:
+        """
+        Specify a sub FastAPI application specific to the auth manager.
+
+        This sub application, if specified, is mounted in the main FastAPI application.
+        """
         return None
 
     def register_views(self) -> None:

airflow/auth/managers/simple/simple_auth_manager.py

@@ -43,6 +43,7 @@
         PoolDetails,
         VariableDetails,
     )
+    from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
 
 
 class SimpleAuthManagerRole(namedtuple("SimpleAuthManagerRole", "name order"), Enum):
@@ -78,6 +79,9 @@ class SimpleAuthManager(BaseAuthManager[SimpleAuthManagerUser]):
     # Cache containing the password associated to a username
     passwords: dict[str, str] = {}
 
+    # TODO: Needs to be deleted when Airflow 2 legacy UI is gone
+    appbuilder: AirflowAppBuilder | None = None
+
     @staticmethod
     def get_generated_password_file() -> str:
         return os.path.join(

airflow/auth/managers/simple/views/auth.py

@@ -21,12 +21,12 @@
 from flask import redirect, request, session, url_for
 from flask_appbuilder import expose
 
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.simple.user import SimpleAuthManagerUser
 from airflow.configuration import conf
 from airflow.utils.jwt_signer import JWTSigner
 from airflow.utils.state import State
 from airflow.www.app import csrf
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 from airflow.www.views import AirflowBaseView
 
 

airflow/cli/cli_parser.py

@@ -36,6 +36,7 @@
 import lazy_object_proxy
 from rich_argparse import RawTextRichHelpFormatter, RichHelpFormatter
 
+from airflow.api_fastapi.app import get_auth_manager_cls
 from airflow.cli.cli_config import (
     DAG_CLI_DICT,
     ActionCommand,
@@ -47,7 +48,6 @@
 from airflow.exceptions import AirflowException
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.utils.helpers import partition
-from airflow.www.extensions.init_auth_manager import get_auth_manager_cls
 
 if TYPE_CHECKING:
     from airflow.cli.cli_config import (

airflow/www/auth.py

@@ -30,6 +30,7 @@
     PERMISSION_PREFIX,
 )
 
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.auth.managers.models.resource_details import (
     AccessView,
     ConnectionDetails,
@@ -40,7 +41,6 @@
 )
 from airflow.configuration import conf
 from airflow.utils.net import get_hostname
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from airflow.auth.managers.base_auth_manager import ResourceMethod

airflow/www/decorators.py

@@ -29,10 +29,10 @@
 from flask import after_this_request, request
 from pendulum.parsing.exceptions import ParserError
 
+from airflow.api_fastapi.app import get_auth_manager
 from airflow.models import Log
 from airflow.utils.log import secrets_masker
 from airflow.utils.session import create_session
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 
 T = TypeVar("T", bound=Callable)
 

airflow/www/extensions/init_appbuilder.py

@@ -39,8 +39,8 @@
 from flask_appbuilder.views import IndexView, UtilView
 
 from airflow import settings
+from airflow.api_fastapi.app import create_auth_manager, get_auth_manager
 from airflow.configuration import conf
-from airflow.www.extensions.init_auth_manager import get_auth_manager, init_auth_manager
 
 if TYPE_CHECKING:
     from flask import Flask
@@ -208,7 +208,9 @@ def init_app(self, app, session):
 
         self._addon_managers = app.config["ADDON_MANAGERS"]
         self.session = session
-        auth_manager = init_auth_manager(self)
+        auth_manager = create_auth_manager()
+        auth_manager.appbuilder = self
+        auth_manager.init()
         self.sm = auth_manager.security_manager
         self.bm = BabelManager(self)
         self._add_global_static()