已配置Airflow UI的Azure AD OAuth，如何适配REST API认证？

核心思路 #

Airflow的API认证独立于UI的OAuth配置，需基于flask-appbuilder扩展自定义认证后端，实现Azure AD JWT令牌的验证与角色解析。

具体实现步骤 #

1. 安装依赖 #

安装令牌解析与密钥集处理工具：

pip install pyjwt python-jose[cryptography]

2. 自定义Azure AD认证类 #

创建azure_ad_api_auth.py文件，实现令牌验证逻辑：

from flask import request
from jose import jwt
from jose.utils import base64url_decode
import requests

class AzureADAPIAuth:
    def __init__(self, tenant_id, client_id, audience):
        self.tenant_id = tenant_id
        self.client_id = client_id
        self.audience = audience
        self.jwks_url = f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
        self.jwks = requests.get(self.jwks_url).json()

    def authenticate(self):
        auth_header = request.headers.get('Authorization', '')
        if not auth_header.startswith('Bearer '):
            return None
        
        token = auth_header.split(' ')[1]
        try:
            # 解析令牌头获取密钥ID
            headers = jwt.get_unverified_headers(token)
            kid = headers['kid']
            # 匹配对应公钥
            rsa_key = next((key for key in self.jwks['keys'] if key['kid'] == kid), None)
            if not rsa_key:
                return None
            
            # 验证令牌签名、受众、发行方
            payload = jwt.decode(
                token,
                rsa_key,
                algorithms=["RS256"],
                audience=self.audience,
                issuer=f"https://login.microsoftonline.com/{self.tenant_id}/v2.0"
            )
            
            # 提取Azure AD分配的应用角色
            roles = payload.get('roles', [])
            if not roles:
                return None
            
            # 构造Airflow兼容的用户对象
            from flask_appbuilder.security.sqla.models import User
            user = User()
            user.username = payload.get('appid', 'azure-service-principal')
            user.roles = roles
            return user
        except Exception as e:
            print(f"Token validation failed: {str(e)}")
            return None

3. 对接Airflow API认证后端 #

创建custom_auth.py文件，将自定义认证类接入Airflow的API认证体系：

from airflow.api.auth.backend.base import BaseAuthBackend
from azure_ad_api_auth import AzureADAPIAuth

# 替换为你的Azure AD配置信息
AZURE_TENANT_ID = "你的租户ID"
AZURE_CLIENT_ID = "你的应用注册ID"
AZURE_AUDIENCE = "你的API受众（通常为应用注册ID或自定义标识符）"

auth = AzureADAPIAuth(AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_AUDIENCE)

class CustomAuthBackend(BaseAuthBackend):
    def authenticate(self, request):
        user = auth.authenticate()
        if user:
            return user.username, user.roles
        return None

修改airflow.cfg配置，指定自定义认证后端：

[api]
auth_backend = airflow.api.auth.backend.custom_auth

4. 角色映射配置 #

将Azure AD的应用角色映射到Airflow内置角色（如Admin/User/Viewer），在AzureADAPIAuth类的authenticate方法中添加映射逻辑：

ROLE_MAPPING = {
    "Airflow.Admin": "Admin",
    "Airflow.User": "User",
    "Airflow.Viewer": "Viewer"
}

# 替换原roles赋值逻辑
mapped_roles = [ROLE_MAPPING.get(role, role) for role in roles]
user.roles = mapped_roles

测试验证 #

使用服务主体获取Azure AD令牌并调用Airflow API：

# 获取访问令牌
TOKEN=$(curl -X POST "https://login.microsoftonline.com/你的租户ID/oauth2/v2.0/token" \
  -H "Content-Type: application/x-www-form-urlencoded" \
  -d "client_id=你的应用注册ID" \
  -d "scope=https://你的API受众/.default" \
  -d "client_secret=你的服务主体密钥" \
  -d "grant_type=client_credentials" | jq -r .access_token)

# 触发指定DAG
curl -X POST "http://你的Airflow地址/api/v1/dags/目标DAGID/dagRuns" \
  -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -d '{"conf": {}}'

注意事项 #

• 确保Azure AD应用注册中已配置对应API权限，且服务主体已被分配应用角色。
• 令牌的aud字段必须与认证类中配置的AZURE_AUDIENCE完全一致，否则验证失败。
• 生产环境建议添加JWKS密钥集缓存，减少重复请求开销。

内容的提问来源于stack exchange，提问作者Kombajn zbożowy