案例

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 中应用就行:

from flask_appbuilder.security.manager import AUTH_REMOTE_USER

AUTH_TYPE = AUTH_REMOTE_USER
FAB_SECURITY_MANAGER_CLASS = 'fab_auth_com.ComCasAirflowSecurityManager'