Login delay causes false authentication error messages on UI Airflow >=3.1.1 #57981

#
## https://gist.github.com/wallyhall/915fedb4dfc766b61f442a32c95e1c29#file-apache_airflow_sso_howto-md
## https://gist.github.com/wallyhall/915fedb4dfc766b61f442a32c95e1c29?permalink_comment_id=5253211#gistcomment-5253211
#

from __future__ import annotations
import os

# Airflow 3.1.x and FAB provider >=2.6.0
#! https://gist.github.com/wallyhall/915fedb4dfc766b61f442a32c95e1c29?permalink_comment_id=5623661#gistcomment-5623661
from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
from airflow.utils.log.logging_mixin import LoggingMixin
from flask_appbuilder.security.manager import AUTH_OAUTH

basedir = os.path.abspath(os.path.dirname(__file__))

print("------------------------------------STARTING AUTH------------------------------------")

# ------------------------------------------------------------------------------
# General Flask AppBuilder / Airflow Auth Config
# ------------------------------------------------------------------------------
WTF_CSRF_ENABLED = True
WTF_CSRF_TIME_LIMIT = None

# ------------------------------------------------------------------------------
# Azure AD OAuth Settings
# ------------------------------------------------------------------------------
AAD_TENANT_ID = os.getenv("AAD_TENANT_ID")
AAD_CLIENT_ID = os.getenv("AAD_CLIENT_ID")
AAD_CLIENT_SECRET = os.getenv("AAD_CLIENT_SECRET")

AUTH_TYPE = AUTH_OAUTH

OAUTH_PROVIDERS = [
    {
        "name": "azure",
        "token_key": "access_token",
        "icon": "fa-windows",
        "remote_app": {
            "api_base_url": f"https://login.microsoftonline.com/{AAD_TENANT_ID}",
            "request_token_url": None,
            "request_token_params": {"scope": "openid profile email"},
            "access_token_url": f"https://login.microsoftonline.com/{AAD_TENANT_ID}/oauth2/v2.0/token",
            "authorize_url": f"https://login.microsoftonline.com/{AAD_TENANT_ID}/oauth2/v2.0/authorize",
            "client_id": f"{AAD_CLIENT_ID}",
            "client_secret": f"{AAD_CLIENT_SECRET}",
            "client_kwargs": {"scope": "openid profile email"},
            "jwks_uri": "https://login.microsoftonline.com/common/discovery/v2.0/keys",
        },
    }
]

# ------------------------------------------------------------------------------
# User registration and role sync
# ------------------------------------------------------------------------------
AUTH_USER_REGISTRATION = True
AUTH_USER_REGISTRATION_ROLE = "Public"
AUTH_ROLES_SYNC_AT_LOGIN = True

# This gets injected via Helm values usually, First you MUST create a role like"Admin with value Admin" in the AppRegistration "App Roles" section in the Azure Portal under Microsoft EntraID.
# Then groups MUST be linked from the Microsoft Entra ID "EnterpriseApplication" section in the Azure Portal under the "Users and Groups" section.
# Each groups or users MUST be assigned a role e.g.: Admin, Op, Viewer in the "Users and Groups"

# https://airflow.apache.org/docs/apache-airflow/2.4.3/security/access-control.html
# AUTH_ROLES_MAPPING = {
#     "airflow-admin": ["Admin"],
#     "airflow-user": ["User"],
# }
PLACEHOLDER_FOR_AUTH_ROLES_MAPPING


# ------------------------------------------------------------------------------
# Custom Security Manager
# ------------------------------------------------------------------------------
class AzureCustomSecurity(FabAirflowSecurityManagerOverride, LoggingMixin):
    """
    Custom FAB security manager for Azure AD OAuth.
    Compatible with Airflow 3.1.x (API server).
    """

    def oauth_user_info_save(self, userinfo):
        """
        Override Airflow 3.1+ user creation to reuse existing users by email.
        """
        email = userinfo.get("email")
        if not email:
            self.log.warning("OAuth userinfo missing email — cannot authenticate user.")
            return None

        existing_user = self.find_user(email=email)
        if existing_user:
            self.log.debug(f"User {email} already exists, reusing record instead of creating new one.")
            return existing_user

        self.log.debug(f"No existing user found for {email}, creating new user.")
        return super().oauth_user_info_save(userinfo)

    def get_oauth_user_info(self, provider, response=None):
        self.log.debug(f"Parsing JWT token for provider: {provider}")
        try:
            me = super().get_oauth_user_info(provider, response)
        except Exception as e:
            import traceback

            traceback.print_exc()
            self.log.debug(e)
            return None

        self.log.debug("_____________________role_keys start__________________________________")
        self.log.debug(f"Parsed JWT token: {me}")
        self.log.debug(f"Role Keys: {me.get('role_keys')}")
        self.log.debug("_____________________role_keys end____________________________________")

        #! Example output from: self.log.debug(f"Parse JWT token : {me}")
        # DEBUG - Parse JWT token : {'email': 'aj@a.com',
        # 'first_name': 'A',
        # 'last_name': 'J',
        # 'username': 'ebxxxxxxxxxxxxxxxxxxxxxxxxxxae7',
        # 'role_keys': ['airflow-user','airflow-admin']}

        return {
            "name": f"{me.get('first_name', '')} {me.get('last_name', '')}",
            "email": me.get("email"),
            "first_name": me.get("first_name"),
            "last_name": me.get("last_name"),
            "id": me.get("username"),
            "username": me.get("email"),  # email as username avoids duplicates
            "role_keys": me.get("role_keys"),
        }


# ------------------------------------------------------------------------------
# Register the Auth Manager for Airflow 3.1.x API Server
# ------------------------------------------------------------------------------
FAB_SECURITY_MANAGER_CLASS = "webserver_config.AzureCustomSecurity"
SECURITY_MANAGER_CLASS = AzureCustomSecurity