Add option in auth manager interface to define FastAPI api

airflow/api_connexion/security.py

@@ -32,8 +32,8 @@
     PoolDetails,
     VariableDetails,
 )
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 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

@@ -70,7 +70,7 @@ def create_app(apps: str = "all") -> FastAPI:
         init_plugins(app)
         init_flask_plugins(app)
         init_error_handlers(app)
-        init_auth_manager()
+        init_auth_manager(app)
 
     if "execution" in apps_list or "all" in apps_list:
         task_exec_api_app = create_task_execution_api_app(app)
@@ -112,34 +112,23 @@ 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 init_auth_manager(app: FastAPI) -> BaseAuthManager:
+    """Initialize the auth manager."""
     global auth_manager
     auth_manager_cls = get_auth_manager_cls()
     auth_manager = auth_manager_cls()
     auth_manager.init()
+
+    auth_manager_fastapi_app = auth_manager.get_fastapi_app()
+    if auth_manager_fastapi_app:
+        app.mount("/auth", auth_manager_fastapi_app)
+
     return auth_manager
 
 
 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
 
@@ -453,6 +454,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/www/auth.py

@@ -39,8 +39,8 @@
     VariableDetails,
 )
 from airflow.configuration import conf
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 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

@@ -30,9 +30,9 @@
 from pendulum.parsing.exceptions import ParserError
 
 from airflow.models import Log
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 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/security_manager.py

@@ -37,6 +37,7 @@
 )
 from airflow.exceptions import AirflowException
 from airflow.models import Connection, DagRun, Pool, TaskInstance, Variable
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.security.permissions import (
     RESOURCE_ADMIN_MENU,
     RESOURCE_ASSET,
@@ -63,7 +64,6 @@
     RESOURCE_XCOM,
 )
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.www.extensions.init_auth_manager import get_auth_manager
 from airflow.www.utils import CustomSQLAInterface
 
 EXISTING_ROLES = {

providers/src/airflow/providers/fab/auth_manager/api/auth/backend/basic_auth.py

@@ -25,8 +25,8 @@
 from flask_appbuilder.const import AUTH_LDAP
 from flask_login import login_user
 
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 
 if TYPE_CHECKING:
     from airflow.providers.fab.auth_manager.models import User

providers/src/airflow/providers/fab/auth_manager/api/auth/backend/kerberos_auth.py

@@ -26,9 +26,9 @@
 from flask import Response, current_app, g, make_response, request
 from requests_kerberos import HTTPKerberosAuth
 
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.configuration import conf
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.utils.net import getfqdn
 
 if TYPE_CHECKING:

providers/src/airflow/providers/fab/auth_manager/api/auth/backend/session.py

@@ -23,7 +23,7 @@
 
 from flask import Response
 
-from airflow.api_fastapi.app import get_auth_manager
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 
 CLIENT_AUTH: tuple[str, str] | Any | None = None
 

providers/src/airflow/providers/fab/auth_manager/api_endpoints/role_and_permission_endpoint.py

@@ -27,7 +27,6 @@
 from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound
 from airflow.api_connexion.parameters import check_limit, format_parameters
 from airflow.api_connexion.security import requires_access_custom_view
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.providers.fab.auth_manager.models import Action, Role
 from airflow.providers.fab.auth_manager.schemas.role_and_permission_schema import (
     ActionCollection,
@@ -37,6 +36,7 @@
     role_schema,
 )
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.security import permissions
 
 if TYPE_CHECKING:

providers/src/airflow/providers/fab/auth_manager/api_endpoints/user_endpoint.py

@@ -28,7 +28,6 @@
 from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound, Unknown
 from airflow.api_connexion.parameters import check_limit, format_parameters
 from airflow.api_connexion.security import requires_access_custom_view
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.providers.fab.auth_manager.models import User
 from airflow.providers.fab.auth_manager.schemas.user_schema import (
     UserCollection,
@@ -37,6 +36,7 @@
     user_schema,
 )
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.security import permissions
 
 if TYPE_CHECKING:

providers/src/airflow/providers/fab/auth_manager/decorators/auth.py

@@ -26,9 +26,9 @@
 
 from airflow.api_connexion.exceptions import PermissionDenied
 from airflow.api_connexion.security import check_authentication
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.configuration import conf
 from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.utils.airflow_flask_app import AirflowApp
 from airflow.utils.net import get_hostname
 from airflow.www.auth import _has_access

providers/src/airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -25,10 +25,12 @@
 
 import packaging.version
 from connexion import FlaskApi
+from fastapi import FastAPI
 from flask import Blueprint, g, url_for
 from packaging.version import Version
 from sqlalchemy import select
 from sqlalchemy.orm import Session, joinedload
+from starlette.middleware.wsgi import WSGIMiddleware
 
 from airflow import __version__ as airflow_version
 from airflow.auth.managers.base_auth_manager import BaseAuthManager, ResourceMethod
@@ -56,6 +58,7 @@
     USERS_COMMANDS,
 )
 from airflow.providers.fab.auth_manager.models import Permission, Role, User
+from airflow.providers.fab.www.app import create_app
 from airflow.security import permissions
 from airflow.security.permissions import (
     RESOURCE_AUDIT_LOG,
@@ -166,13 +169,36 @@ def get_cli_commands() -> list[CLICommand]:
             commands.append(GroupCommand(name="fab-db", help="Manage FAB", subcommands=DB_COMMANDS))
         return commands
 
+    def get_fastapi_app(self) -> FastAPI | None:
+        flask_blueprint = self.get_api_endpoints()
+
+        if not flask_blueprint:
+            return None
+
+        flask_app = create_app()
+        flask_app.register_blueprint(flask_blueprint)
+
+        app = FastAPI(
+            title="FAB auth manager API",
+            description=(
+                "This is FAB auth manager API. This API is only available if the auth manager used in "
+                "the Airflow environment is FAB auth manager. "
+                "This API provides endpoints to manager users and permissions managed by the FAB auth "
+                "manager."
+            ),
+        )
+        app.mount("/", WSGIMiddleware(flask_app))
+
+        return app
+
     def get_api_endpoints(self) -> None | Blueprint:
         folder = Path(__file__).parents[0].resolve()  # this is airflow/auth/managers/fab/
         with folder.joinpath("openapi", "v1.yaml").open() as f:
             specification = safe_load(f)
         return FlaskApi(
             specification=specification,
             resolver=_LazyResolver(),
+            # TODO: change to "/fab/v1" when legacy UI is gone
             base_path="/auth/fab/v1",
             options={"swagger_ui": SWAGGER_ENABLED, "swagger_path": SWAGGER_BUNDLE.__fspath__()},
             strict_validation=True,

providers/src/airflow/providers/fab/auth_manager/security_manager/override.py

@@ -73,7 +73,6 @@
 from werkzeug.security import check_password_hash, generate_password_hash
 
 from airflow import __version__ as airflow_version
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import DagBag, DagModel
@@ -107,6 +106,7 @@
     CustomUserInfoEditView,
 )
 from airflow.providers.fab.auth_manager.views.user_stats import CustomUserStatsChartView
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.security import permissions
 from airflow.www.security_manager import AirflowSecurityManagerV2
 from airflow.www.session import AirflowDatabaseSessionInterface

providers/src/airflow/providers/fab/www/app.py

@@ -33,6 +33,7 @@
 from airflow.providers.fab.www.extensions.init_manifest_files import configure_manifest_files
 from airflow.providers.fab.www.extensions.init_security import init_xframe_protection
 from airflow.providers.fab.www.extensions.init_views import init_error_handlers, init_plugins
+from airflow.www.extensions.init_security import init_api_auth
 
 app: Flask | None = None
 
@@ -41,7 +42,7 @@
 csrf = CSRFProtect()
 
 
-def create_app(config=None, testing=False):
+def create_app():
     """Create a new instance of Airflow WWW app."""
     flask_app = Flask(__name__)
     flask_app.secret_key = conf.get("webserver", "SECRET_KEY")
@@ -63,6 +64,7 @@ def create_app(config=None, testing=False):
 
     configure_logging()
     configure_manifest_files(flask_app)
+    init_api_auth(flask_app)
 
     with flask_app.app_context():
         init_appbuilder(flask_app)

providers/src/airflow/providers/fab/www/extensions/init_appbuilder.py

@@ -40,7 +40,7 @@
 
 from airflow import settings
 from airflow.configuration import conf
-from airflow.www.extensions.init_auth_manager import init_auth_manager
+from airflow.providers.fab.www.extensions.init_auth_manager import init_auth_manager
 
 if TYPE_CHECKING:
     from flask import Flask

providers/src/airflow/providers/fab/www/extensions/init_auth_manager.py

@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowConfigException
+
+if TYPE_CHECKING:
+    from airflow.auth.managers.base_auth_manager import BaseAuthManager
+    from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
+
+auth_manager: BaseAuthManager | None = None
+
+
+def get_auth_manager_cls() -> type[BaseAuthManager]:
+    """
+    Return just the auth manager class without initializing it.
+
+    Useful to save execution time if only static methods need to be called.
+    """
+    auth_manager_cls = conf.getimport(section="core", key="auth_manager")
+
+    if not auth_manager_cls:
+        raise AirflowConfigException(
+            "No auth manager defined in the config. "
+            "Please specify one using section/key [core/auth_manager]."
+        )
+
+    return auth_manager_cls
+
+
+def init_auth_manager(appbuilder: AirflowAppBuilder) -> BaseAuthManager:
+    """
+    Initialize the auth manager.
+
+    Import the user manager class and instantiate it.
+    """
+    global auth_manager
+    auth_manager_cls = get_auth_manager_cls()
+    auth_manager = auth_manager_cls(appbuilder)
+    auth_manager.init()
+    return auth_manager
+
+
+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(
+            "Auth Manager has not been initialized yet. "
+            "The `init_auth_manager` method needs to be called first."
+        )
+    return auth_manager

providers/src/airflow/providers/fab/www/extensions/init_jinja_globals.py

@@ -21,8 +21,8 @@
 import pendulum
 
 import airflow
-from airflow.api_fastapi.app import get_auth_manager
 from airflow.configuration import conf
+from airflow.providers.fab.www.extensions.init_auth_manager import get_auth_manager
 from airflow.settings import STATE_COLORS
 from airflow.utils.net import get_hostname
 from airflow.utils.platform import get_airflow_git_version