How to assign Admin/User roles based on email with Google OAuth2?

[bun-paul-kwon]

Hello Airflow community,

We are currently using AWS MWAA (Celery Executor) and considering migrating to Airflow 3.0.2 on Kubernetes (AWS EKS).

• Image: apache/airflow:3.0.2
• Helm chart: https://airflow.apache.org airflow 1.18.0

We are integrating Google OAuth2 for authentication and authorization. We've configured our webserver_config.py via the apiSever.apiServerConfig in the Helm values.yaml as follows:

apiServer:
    apiServerConfig: |
        from flask_appbuilder.security.manager import AUTH_OAUTH

        AUTH_TYPE = AUTH_OAUTH
        CSRF_ENABLED = True
        AUTH_USER_REGISTRATION = True
        AUTH_USER_REGISTRATION_ROLE = "User"
        AUTH_ROLES_SYNC_AT_LOGIN = False
        OAUTH_PROVIDERS = [
            {
                'name': 'google',
                'token_key': 'access_token',
                'icon': 'fa-google',
                'remote_app': {
                    'api_base_url': 'https://www.googleapis.com/oauth2/v2/',
                    'client_kwargs': {'scope': 'email profile'},
                    'access_token_url': 'https://accounts.google.com/o/oauth2/token',
                    'authorize_url': 'https://accounts.google.com/o/oauth2/auth',
                    'request_token_url': None,
                    'client_id': '<MY_CLIENT_ID>',
                    'client_secret': '<MY_CLIENT_SECRET>'
                }
            }
        ]

Login with Google accounts is working correctly. However, because AUTH_USER_REGISTRATION_ROLE = "User", every new user is automatically mapped to the "User" role.

Our goal is to have specific accounts (e.g., [email protected]) automatically mapped to the Admin role upon login, while all others default to "User".

We found a workaround by manually updating the role_id in the ab_user_role table in the metadata database for the specific user. This works, but we are unsure if this is the best practice for managing Admin accounts.

What is the recommended approach for managing Admin vs. general User roles when using OAuth2? Is there a configuration setting or a specific method (perhaps custom logic in webserver_config.py?) we should be using to assign roles based on the user's email address?

Thank you 🙏

[bun-paul-kwon]

I managed to resolve this issue. The key was to implement a CustomOidcSecurityManager to override the oauth_user_info method.

# values.yaml
apiServer:
    apiServerConfig: |
        from flask_appbuilder.security.manager import AUTH_OAUTH
        from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
        import logging

        log = logging.getLogger(__name__)

        ADMIN_EMAIL_LIST = {
            '[email protected]',
            '[email protected]'
        }

        class CustomOidcSecurityManager(FabAirflowSecurityManagerOverride):
            def oauth_user_info(self, provider, response):
                
                if provider != 'google':
                    return {}

                userinfo = response.get('userinfo')
                email = userinfo.get('email')
                hd = userinfo.get('hd')

                if hd is None or hd != "mycompany.com":
                    return {}
                
                if email is None:
                    return {}

                if email in ADMIN_EMAIL_LIST:
                    userinfo['role_keys'] = ['Admin']
                else:
                    userinfo['role_keys'] = ['User']

                return userinfo

        AUTH_TYPE = AUTH_OAUTH
        CSRF_ENABLED = True
        AUTH_USER_REGISTRATION = True
        # AUTH_USER_REGISTRATION_ROLE = "Viewer"
        AUTH_ROLES_SYNC_AT_LOGIN = True
        PERMANENT_SESSION_LIFETIME = 86400
        AUTH_ROLES_MAPPING = {
          "Admin": ["Admin"],
          "User": ["User"],
        }
        SECURITY_MANAGER_CLASS = CustomOidcSecurityManager
        OAUTH_PROVIDERS = [
            {
                'name': 'google',
                'token_key': 'access_token',
                'icon': 'fa-google',
                'remote_app': {
                    'api_base_url': 'https://www.googleapis.com/oauth2/v2/',
                    'client_kwargs': {'scope': 'openid profile email'},
                    'access_token_url': 'https://accounts.google.com/o/oauth2/token',
                    'authorize_url': 'https://accounts.google.com/o/oauth2/auth',
                    'request_token_url': None,
                    'client_id': '<CLIENT_ID>',
                    'client_secret': '<CLIENT_SECRET>',
                    'userinfo_endpoint': 'https://openidconnect.googleapis.com/v1/userinfo',
                    'jwks_uri': 'https://www.googleapis.com/oauth2/v3/certs',
                    'authorize_params': {'hd': 'mycompany.com'}
                }
            }
        ]