Repository: apache/airflow
File: airflow-core/src/airflow/api_fastapi/app.py (Python)

Code before:

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 create_auth_manager() -> BaseAuthManager:
    """Create the auth manager."""
    auth_manager_cls = get_auth_manager_cls()
    if _AuthManagerState.instance is None or not isinstance(_AuthManagerState.instance, auth_manager_cls):
        with _AuthManagerState._lock:
            if _AuthManagerState.instance is None or not isinstance(
                _AuthManagerState.instance, auth_manager_cls
            ):
                _AuthManagerState.instance = auth_manager_cls()
    return _AuthManagerState.instance


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

    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."""
    if _AuthManagerState.instance is None:
        raise RuntimeError(
            "Auth Manager has not been initialized yet. "
            "The `init_auth_manager` method needs to be called first."
        )
    return _AuthManagerState.instance


Review comment (security):
We can even avoid calling ‎`get_auth_manager_cls` on every ‎`create_auth_manager` call: https://github.com/apache/airflow/pull/61310#discussion_r2751276702

Code after:

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 create_auth_manager() -> BaseAuthManager:
    """Create the auth manager."""
    if _AuthManagerState.instance is not None:
        return _AuthManagerState.instance
    with _AuthManagerState._lock:
        if _AuthManagerState.instance is None:
            auth_manager_cls = get_auth_manager_cls()
            _AuthManagerState.instance = auth_manager_cls()
    return _AuthManagerState.instance


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

    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."""
    if _AuthManagerState.instance is None:
        raise RuntimeError(
            "Auth Manager has not been initialized yet. "
            "The `init_auth_manager` method needs to be called first."
        )
    return _AuthManagerState.instance