案例 Airflow powers AI。
Airflow SSO 接入 公司 SSO 系统不是基于开源标准，而是一套自定义的方式，目前网上没有成熟的解决方案，通过查看 Flask-AppBuilder 和 Airflow 的代码发现可以扩展 flask_appbuilder.security.views.AuthRemoteUserView 并通过自定义的 SecurityManager 指定 authremoteuserview 来实现，去掉具体 SSO 逻辑后的代码如下：
from urllib.parse import urlencode from urllib.parse import urljoin import requests from flask import flash from flask import redirect from flask import request from flask_appbuilder.baseviews import expose from flask_appbuilder.security.views import AuthRemoteUserView try: from airflow.www.security import AirflowSecurityManager except ImportError: AirflowSecurityManager = None __version__ = "0.1.0" AUTHORIZE_URL = "https://example.com/sso/login" ACCESS_TOKEN_URL = "https://example.com/sso/check" class AuthComCasView(AuthRemoteUserView): def _get_redirect_uri(self): return urljoin(request.host_url, self.appbuilder.get_url_for_login) def get_authorize_params(self): return { "callback": self._get_redirect_uri(), } @expose("/login/") def login(self): token = request.args.get("token") if not token: params = self.get_authorize_params() redirect_uri = u"{}?{}".format( AUTHORIZE_URL, urlencode(params), ) return redirect(redirect_uri) data = self.exchange_token(token) if data["status"] < 0: flash("Invalid Token", "info") return "Invalid token" # Set REMOTE_USER to let user login request.environ["REMOTE_USER"] = data["username"] return super().login() @staticmethod def get_token_params(token): return { "token": token, } def exchange_token(self, token): data = self.get_token_params(token) return requests.get(ACCESS_TOKEN_URL, params=data).json() if AirflowSecurityManager is not None: class ComCasAirflowSecurityManager(AirflowSecurityManager): authremoteuserview = AuthComCasView 然后在 Airflow 的 webserver_config.py 中应用就行:
...