案例

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'