Airflow 2.3.0集成AWS Cognito：用户组与角色映射实现求助

关键修改步骤 #

1. 配置Cognito返回用户组信息

   • 进入Cognito用户池的App Client设置，在Token configuration的ID Token部分，勾选cognito:groups，确保组信息被包含在ID Token中。
   • 确认OAuth 2.0 scopes包含openid、email、profile。

2. 安装依赖库
   在Airflow Helm部署的values.yaml中添加解析JWT所需的依赖：

   webserver:
     extraPipPackages:
       - pyjwt==2.8.0
       - cryptography==41.0.7
   

3. 修改webserver_config.py代码
   更新自定义SecurityManager，解析ID Token中的Cognito组并映射到Airflow角色：

   import sys
   from airflow import configuration as conf
   from airflow.www.security import AirflowSecurityManager
   from flask_appbuilder.security.manager import AUTH_OAUTH
   import logging
   import os
   import json
   import jwt
   from jwt import PyJWKClient
   
   logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
   logger = logging.getLogger()
   
   class CognitoSecurity(AirflowSecurityManager):
   
       def oauth_user_info(self, provider, response=None):
           if provider == "aws_cognito" and response:
               logger.info(response)
               
               # 从认证响应中获取ID Token并解析用户组
               id_token = response.get('id_token')
               cognito_groups = []
               if id_token:
                   try:
                       # 通过JWKS验证并解析ID Token
                       jwks_client = PyJWKClient(JWKS_URI)
                       signing_key = jwks_client.get_signing_key_from_jwt(id_token)
                       payload = jwt.decode(
                           id_token,
                           signing_key.key,
                           algorithms=["RS256"],
                           audience=CONSUMER_KEY,
                           issuer=f"https://cognito-idp.{os.environ['AWS_REGION']}.amazonaws.com/{os.environ['COGNITO_POOL_ID']}"
                       )
                       cognito_groups = payload.get('cognito:groups', [])
                       logger.info("Cognito groups retrieved: %s", cognito_groups)
                   except jwt.exceptions.InvalidTokenError as e:
                       logger.error("Failed to decode ID Token: %s", str(e))
   
               # 获取用户基础信息
               res = self.appbuilder.sm.oauth_remotes[provider].get('oauth2/userInfo')
               if res.raw.status != 200:
                   logger.error('Failed to obtain user info: %s', res.data)
                   return
               me = json.loads(res._content)
               logger.info("User data: %s", me)
   
               # 映射Cognito组到Airflow内置/自定义角色
               group_role_mapping = {
                   # 替换为你的Cognito组名与Airflow角色的对应关系
                   'airflow-admins': 'Admin',
                   'airflow-users': 'User',
                   'airflow-operators': 'Op',
                   'airflow-viewers': 'Viewer'
               }
               airflow_roles = []
               for group in cognito_groups:
                   if group in group_role_mapping:
                       airflow_roles.append(group_role_mapping[group])
               
               # 若没有匹配组，使用默认注册角色
               if not airflow_roles:
                   airflow_roles.append(AUTH_USER_REGISTRATION_ROLE)
   
               return {
                   "username": me.get("username"), 
                   "email": me.get("email"),
                   "roles": airflow_roles
               }
           else:
               return {}
   
   AUTH_TYPE = AUTH_OAUTH
   AUTH_ROLES_SYNC_AT_LOGIN = True  # 每次登录同步角色
   AUTH_USER_REGISTRATION = True
   AUTH_USER_REGISTRATION_ROLE = "Viewer"  # 默认角色（无匹配组时使用）
   COGNITO_URL = os.environ['COGNITO_URL'] 
   CONSUMER_KEY = os.environ['CONSUMER_KEY']
   SECRET_KEY = os.environ['SECRET_KEY']
   REDIRECT_URI = os.environ['REDIRECT_URI']
   JWKS_URI = ("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json"
               % (os.environ['AWS_REGION'], os.environ['COGNITO_POOL_ID'])) 
   OAUTH_PROVIDERS = [{
       'name':'aws_cognito',
       'token_key':'access_token',
       'url': COGNITO_URL,
       'icon': 'fa-amazon',
       'remote_app': {
           'client_id': CONSUMER_KEY,
           'client_secret': SECRET_KEY,
           'base_url': os.path.join(COGNITO_URL, 'oauth2/idpresponse'),
           "api_base_url": COGNITO_URL,
           'redirect_uri' : REDIRECT_URI,
           'jwks_uri': JWKS_URI,
           'client_kwargs': {
               'scope': 'email openid profile'
           },
           'access_token_url': os.path.join(COGNITO_URL, 'oauth2/token'),
           'authorize_url': os.path.join(COGNITO_URL, 'oauth2/authorize')
       }
   }]
   SECURITY_MANAGER_CLASS = CognitoSecurity
   

4. 部署更新

   • 应用修改后的values.yaml和webserver_config.py（通过Helm的webserver.extraConfig挂载或更新配置）
   • 重启Airflow Webserver Pod，使配置生效

验证 #

登录Airflow后，进入Security > List Users，查看用户的角色是否与Cognito组匹配。每次登录时，角色会自动同步，确保权限的一致性。

内容的提问来源于stack exchange，提问作者tfleischer