Skip to content

Instantly share code, notes, and snippets.

@wallyhall

wallyhall/apache_airflow_sso_howto.md

Last active October 26, 2024 18:18
Show Gist options
  • Save wallyhall/915fedb4dfc766b61f442a32c95e1c29 to your computer and use it in GitHub Desktop.
Save wallyhall/915fedb4dfc766b61f442a32c95e1c29 to your computer and use it in GitHub Desktop.
Download ZIP
Apache Airflow Azure AAD SSO howto
Raw
apache_airflow_sso_howto.md

The following instructions for enabling Azure SSO for Apache Airflow nearly take you all the way - but fall short a couple of details around the configuration of airflow itself:

https://objectpartners.com/2021/12/24/enterprise-auth-for-airflow-azure-ad

All the "Azure" instructions there can be safely followed - the resulting webserver_config.py (which can be injected into a dockerised Airflow in /opt/airflow/webserver_config.py) can be built from the following:

from __future__ import annotations

import os

from airflow.www.fab_security.manager import AUTH_OAUTH
from airflow.www.security import AirflowSecurityManager
from airflow.utils.log.logging_mixin import LoggingMixin

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

# Flask-WTF flag for CSRF
WTF_CSRF_ENABLED = True
WTF_CSRF_TIME_LIMIT = None

AUTH_TYPE = AUTH_OAUTH

OAUTH_PROVIDERS = [{
    'name':'Microsoft Azure AD',
    'token_key':'access_token',
    'icon':'fa-windows',
    'remote_app': {
        'api_base_url': "https://login.microsoftonline.com/{}".format(os.getenv("AAD_TENANT_ID")),
        'request_token_url': None,
        'request_token_params': {
            'scope': 'openid email profile'
        },
        'access_token_url': "https://login.microsoftonline.com/{}/oauth2/v2.0/token".format(os.getenv("AAD_TENANT_ID")),
        "access_token_params": {
            'scope': 'openid email profile'
        },
        'authorize_url': "https://login.microsoftonline.com/{}/oauth2/v2.0/authorize".format(os.getenv("AAD_TENANT_ID")),
        "authorize_params": {
            'scope': 'openid email profile'
        },
        'client_id': os.getenv("AAD_CLIENT_ID"),
        'client_secret': os.getenv("AAD_CLIENT_SECRET"),
        'jwks_uri': 'https://login.microsoftonline.com/common/discovery/v2.0/keys'
    }
}]

AUTH_USER_REGISTRATION_ROLE = "Public"
AUTH_USER_REGISTRATION = True
AUTH_ROLES_SYNC_AT_LOGIN = True
AUTH_ROLES_MAPPING = {
    "airflow_prod_admin": ["Admin"],
    "airflow_prod_user": ["Op"],
    "airflow_prod_viewer": ["Viewer"]
}

class AzureCustomSecurity(AirflowSecurityManager, LoggingMixin):
    def get_oauth_user_info(self, provider, response=None):
        me = self._azure_jwt_token_parse(response["id_token"])
        return {
            "name": me["name"],
            "email": me["email"],
            "first_name": me["given_name"],
            "last_name": me["family_name"],
            "id": me["oid"],
            "username": me["preferred_username"],
            "role_keys": me["roles"]
        }

# the first of these two appears to work with older Airflow versions, the latter newer.
FAB_SECURITY_MANAGER_CLASS = 'webserver_config.AzureCustomSecurity'
SECURITY_MANAGER_CLASS = AzureCustomSecurity

The above assumes environment variables are configured for the OAuth client secret, etc - and has been tested thoroughly and confirmed working.

Note the roles need to match what you configured in Azure (the example above is using airflow_prod_user etc, in deviation to the linked article above).

@vdozal
Copy link

vdozal commented Oct 25, 2024 via email

@nitinmahawadiwar
Copy link

Thank you @vdozal . This was really helpful.

I am using almost same configuration except for the environment variables (but I am ignoring it for now.) .
Can you please help me to understand an issue with role_keys given below?

Let me explain the flow. un-wanted code is removed for clarity.

....
AUTH_USER_REGISTRATION_ROLE = "Public"
....
AUTH_ROLES_MAPPING = {
        "Entra ID admin Role": ["Admin"],
        "Entra ID viewer Role": ["Viewer"],
        "Entra ID Op Role": ["Op"],
 }
...
  class AzureCustomSecurity(FabAirflowSecurityManagerOverride,.....
............
 return {
            ................
                "role_keys": me.get("role_keys", ["airflow_public"])
            }

My observation is, inside the return block, even through "role_keys" is having Admin/Ops role values from Azure , its always selecting the role assigned for AUTH_USER_REGISTRATION_ROLE = "XXX" and the UI renders accordingly.

Log statements for me object is showing all correct values that are received from Azure.

Is there anything I am missing in this case ? My understanding is, whatever the role_keys inside the return block is set, the same role should be applied while rendering the UI.

@yriveiro
Copy link

@vdozal loading the AAD values from a secret with this values and it works:

webserver:
  webserverConfigConfigMapName: webserver-config-custom
  env:
    - name: AAD_TENANT_ID
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_tenant_id
    - name: AAD_CLIENT_ID
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_client_id
    - name: AAD_CLIENT_SECRET
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_client_secret

@yriveiro
Copy link

Does anyone know how I can get the REST API working as well?

My users will not have user and password, which doesn't allow me to use a basic auth provider. session backend probably works in the browser because it does some magic with the cookies, but I would like to call the REST API using curl

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment