[Python] Airflow + Keycloak SSO: User with "Admin" role gets "Viewer" role after login

We’re using Apache Airflow 3.0.2 deployed on Kubernetes with the official Airflow Helm chart (1.17.0). We’ve integrated Keycloak for SSO using OAuth. Users are able to log in successfully via SSO, but even if they have the "Admin" role in Keycloak, they are logged in with only the Viewer role in Airflow.

What we’ve tried:


In webserver_config.py, we've enabled OAuth and set up role mapping as follows:

import os
from flask_appbuilder.security.manager import AUTH_OAUTH

AUTH_TYPE = AUTH_OAUTH
OAUTH_PROVIDERS = [{
'name': 'keycloak',
'token_key': 'access_token',
'icon': 'fa-circle-o',
'remote_app': {
'client_id': os.getenv("AIRFLOW_OAUTH_CLIENT_ID"),
'client_secret': os.getenv("AIRFLOW_OAUTH_CLIENT_SECRET"),
'api_base_url': 'https://<keycloak-domain>/realms/<realm-name>/protocol/openid-connect',
'client_kwargs': {
'scope': 'openid profile email'
},
'access_token_url': 'https://<keycloak-domain>/realms/<realm-name>/protocol/openid-connect/token',
'authorize_url': 'https://<keycloak-domain>/realms/<realm-name>/protocol/openid-connect/auth',
'request_token_url': None,
'server_metadata_url': os.getenv("AIRFLOW_METADATA_URL"),
}
'role_keys': 'realm_access.roles',

}]
AUTH_ROLES_MAPPING = {
"Viewer": ["Viewer"],
"Admin": ["Admin"],
"User": ["User"],
"Public": ["Public"],
"Op": ["Op"],
}

# Uncomment to setup Full admin role name
AUTH_ROLE_ADMIN = "Admin"
# Uncomment to setup Public role name, no authentication needed
AUTH_ROLE_PUBLIC = "Public"
# Will allow user self registration
AUTH_USER_REGISTRATION = True
# The default user self registration role
AUTH_USER_REGISTRATION_ROLE = "Viewer"
AUTH_ROLES_SYNC_AT_LOGIN = True



We also tried customising the security manager to override the role mapping like this:

from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
from base64 import b64decode
import jwt
import requests
from cryptography.hazmat.primitives import serialization


OIDC_ISSUER = os.getenv("OIDC_ISSUER")
# Fetch public key
req = requests.get(OIDC_ISSUER)
key_der_base64 = req.json()["public_key"]
key_der = b64decode(key_der_base64.encode())
public_key = serialization.load_der_public_key(key_der)

class CustomSecurityManager(FabAirflowSecurityManagerOverride):
def get_oauth_user_info(self, provider, response):
if provider == "keycloak":
token = response["access_token"]
me = jwt.decode(token, public_key, algorithms=["HS256", "RS256"], audience=os.getenv("SSO_CLIENT_ID"))

# Extract roles from resource access
realm_access = me.get("realm_access", {})
groups = realm_access.get("roles", [])

log.info("groups: {0}".format(groups))

if not groups:
groups = ["Viewer"]

userinfo = {
"username": me.get("preferred_username"),
"email": me.get("email"),
"first_name": me.get("given_name"),
"last_name": me.get("family_name"),
"role_keys": groups,
}

log.info("user info: {0}".format(userinfo))

return userinfo
else:
return {}


SECURITY_MANAGER_CLASS = "my_module.CustomSecurityManager"



And added it in webserver_config.py:

Each SSO user is assigned the Admin role, but there are two records for user role mapping in the PostgreSQL table ab_user_role: one for Admin and another for Viewer.

What we want:


If a user has the "Admin" role in Keycloak, they should get the Admin role in Airflow.

Additional info:

• The Keycloak token contains roles under realm_access.roles
• Example decoded token snippet:

"realm_access": {
"roles": ["offline_access", "uma_authorization", "admin"]
}

Questions:

1. How can we extract the Keycloak realm roles correctly from the token and map them to Airflow roles?
2. Is there a better way to define role mapping with Keycloak and Airflow?
3. Do we need to modify FAB (Flask AppBuilder) role sync or extend the OAuthUserInfo object?

Any help or working example would be appreciated!

Continue reading...