Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add option in auth manager interface to define FastAPI api
Browse files Browse the repository at this point in the history
vincbeck committed Dec 18, 2024
1 parent 50c2c9a commit 6bb131c
Showing 32 changed files with 98 additions and 132 deletions.
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/asset_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/dag_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/dag_parsing.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/dag_run_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/dag_source_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/dag_stats_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/import_error_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/task_instance_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/endpoints/xcom_endpoint.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/api_connexion/security.py
Original file line number Diff line number Diff line change
@@ -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
33 changes: 14 additions & 19 deletions airflow/api_fastapi/app.py
Original file line number Diff line number Diff line change
@@ -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,29 @@ 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()
return auth_manager


def init_auth_manager(app: FastAPI | None = None) -> BaseAuthManager:
"""Initialize the auth manager."""
auth_manager = create_auth_manager()
auth_manager.init()

auth_manager_fastapi_app = auth_manager.get_fastapi_app()
if app and 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(
19 changes: 11 additions & 8 deletions airflow/auth/managers/base_auth_manager.py
Original file line number Diff line number Diff line change
@@ -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:
4 changes: 4 additions & 0 deletions airflow/auth/managers/simple/simple_auth_manager.py
Original file line number Diff line number Diff line change
@@ -43,6 +43,7 @@
PoolDetails,
VariableDetails,
)
from airflow.www.extensions.init_appbuilder import AirflowAppBuilder


class SimpleAuthManagerRole(namedtuple("SimpleAuthManagerRole", "name order"), Enum):
@@ -80,6 +81,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(
2 changes: 1 addition & 1 deletion airflow/auth/managers/simple/views/auth.py
Original file line number Diff line number Diff line change
@@ -22,12 +22,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

logger = logging.getLogger(__name__)
2 changes: 1 addition & 1 deletion airflow/cli/cli_parser.py
Original file line number Diff line number Diff line change
@@ -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 (
2 changes: 1 addition & 1 deletion airflow/www/auth.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion airflow/www/decorators.py
Original file line number Diff line number Diff line change
@@ -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)

6 changes: 4 additions & 2 deletions airflow/www/extensions/init_appbuilder.py
Original file line number Diff line number Diff line change
@@ -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()
Loading

0 comments on commit 6bb131c

Please sign in to comment.