Airflow, Azure and OAuth

NOTE: This is an incomplete article – I will continue to publish more as I can. I have provided the needed code for “webserver_config.py” I have not included information for the “App Registration” in Azure.

This article stems from me not finding enough information, in one place, pertaining to authenticating to Apache Airflow leveraging OAuth and Azure.

I was able to piece what I needed together using documentation from different sources: GitHub, Flask, Apache, etc.

My hope is that this article provide you with the information needed to get authentication functional in your environment.

If you are using the Apache Airflow public helm chart – this is the code that will help get you going. This can be added to your “values.yaml” or “overrides.yaml” file.  If not using a helm chart, then add the python portion – starting with the first “from” to the end of the code block and add it to your “webserver_config.py” file.

webserver:
  service:
    type: NodePort
  webserverConfig: |
    
    from airflow.www.security import AirflowSecurityManager
    import logging
    from typing import Dict, Any, List, Union
    from flask_appbuilder.security.manager import AUTH_OAUTH
    import os

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

    WTF_CSRF_ENABLED = True
    # SQLALCHEMY_DATABASE_URI = conf.get("core", "SQL_ALCHEMY_CONN")

    AUTH_TYPE = AUTH_OAUTH
    AUTH_ROLES_SYNC_AT_LOGIN = True
    PERMANENT_SESSION_LIFETIME = 1800 # force users to reauth after inactivity period time in seconds
    AUTH_USER_REGISTRATION = True
    AUTH_USER_REGISTRATION_ROLE = "Public"

    class AzureRoleBasedSecurityManager(AirflowSecurityManager):
        def _get_oauth_user_info(self, provider, resp):
            if provider == "azure":
                me = self._azure_jwt_token_parse(resp["id_token"])
                return {
                    "id": me["oid"],
                    "username": me["upn"],
                    "name": me["name"],
                    "email": me["upn"],
                    "first_name": me["given_name"],
                    "last_name": me["family_name"],
                    "role_keys": me["roles"],
                }
            return {}
        oauth_user_info = _get_oauth_user_info

    SECURITY_MANAGER_CLASS = AzureRoleBasedSecurityManager

    # In order of least permissive - default is "Public" - see AUTH_USER_REGISTRATION_ROLE
    AUTH_ROLES_MAPPING = {
      "Public": ["Public"],
      "Viewer": ["Viewer"],
      "User": ["User"],
      "Op": ["Op"],
      "Admin": ["Admin"],
    }

    OAUTH_PROVIDERS = [
      {
        "name": "azure",
        "icon": "fa-windows",
        "token_key": "access_token",
        "remote_app": {
          "client_id": "<from Azure>",
          "client_secret": "<from Azure>",
          "api_base_url": "https://login.microsoftonline.com/<from Azure>/oauth2",
          "client_kwargs": {
            "scope": "User.read name preferred_username email profile upn",
            "resource": "<from Azure>"},
          "access_token_url": "https://login.microsoftonline.com/<from Azure>/oauth2/token",
          "authorize_url": "https://login.microsoftonline.com/<from Azure>/oauth2/authorize",
          "request_token_url": None,
        },
      },
    ]

Tags: CloudSecurity