Superset SSO改造和自定义宏命令
背景
在最近的一个项目上，客户想要为他们的多租户（Multi-tenant）系统添加一个新的报表中心。技术选型自然沿用之前的选择：Apache Superset，一款由爱彼迎贡献给开源社区的框架。
关于Superset
Superset的前端是中规中矩的React，图表功能则是使用NVD3/D3。后端没有使用万年Java，而是Python3。Web方面使用的是Flask框架，其他的框架没有过多的深入了解。
需要解决的问题
由于之前的业务原因，之前的系统在用户登录时，只能选择其中的一个租户绑定到会话中。这个模式在业务早期没有什么困扰，但随着多租户用户的增多，系统的用户更希望看到跨租户的总览数据。
为此，我们新增了一个资源服务，提供了一个接口用于查询到当前用户的租户信息。用户的认证时通过OAuth 2.0，连接到鉴权服务。
这种变化对于原来的解决方案带来了两个问题：
- Superset需要接入资源服务所用的鉴权服务，并且在OAuth 2.0鉴权后访问资源服务，通过接口获取到当前用户的租户信息。
- Superset需要在执行查询时，动态插入行级过滤条件，这个过滤条件的值是依赖当前用户的租户信息。这可以使用SQL Templating，SQL templating，内置了一些表达式（官方称之为macro，宏，下文称之为宏命令），但功能有限。
之前的做法是当通过OAuth登录Superset，登录的用户名被改为租户的ID，也就是一个租户下的多个用户在使用Superset时使用的是一个Superset用户。这是一个安全的隐患，无法准确地追踪用户的行为。另外，因为Superset的Row level security只能绑定到角色上，所以每个租户用户又有一个独有的角色。这样的影响是显而易见的：但随着业务的增长，租户相关的数据会越来越多，一定程度上造成管理上的混乱。
定制化改造
针对问题1，在现在的Superset（1.3.2）中，早已提供了对OAuth登录的支持，官方提供的教程也很详细。但是在开发过程中，还是遇到了一些小问题。
针对问题2，想办法改造这个SQL templating的文本处理逻辑，增加更多的宏命令，来获取当前用户的租户信息。对于这个功能，官方文档只提供了一个针对Presto数据库的文本处理改造方案，对于这部分功能改造的博客，网上的信息很少。但是经过摸索，还是走出了一条路。
准备环境
官方提供两种方案，一种容器化的，另一种是本地化加虚拟环境。为了调试方便，我采用了后者。
Superset默认使用sqlite，本地启动的话，sqlite文件在~/.superset/superset.db，可以使用IDEA的database面板打开。数据库schema请选择main。
教程中提到的环境变量PYTHONPATH，可以理解为Java中的CLASS_PATH（是目录，而不是具体的某个文件），用于加载外部的模块（module）。因为Python是解释型语言，所以可以在这个目录直接放入Pythone文件。Superset在启动时会加载这个目录下的superset_config.py，并根据其中的代码，加载其他模块。
改造OAuth SSO
请先阅读官方教程：传送门（英文）。
安装依赖
Superset接入OAuth SSO需要依赖库Authlib，可以通过pip安装。
pip install Authlib
对于采用容器化部署的小伙伴，要注意容器被重置时要安装下载这个依赖。
对于喜欢多个命令行窗口的小伙伴，要注意安装这个依赖时，要激活superset虚拟环境（virtualenv）。
配置SSO
根据教程，我们会在superset_config.py中选择认证方式为OAuth，并添加鉴权服务的配置，其中配置的详细说明如下：
from flask_appbuilder.security.manager import AUTH_OAUTH
AUTH_TYPE = AUTH_OAUTH # 选择认证方式，注意，这个值是引用自flask_appbuilder.security.manager
OAUTH_PROVIDERS = [
{
'name': 'spring-sso', # SSO的名字，用于展示在登录页面，格式为SIGN WITH {SSO的名字，大写}。可以配置多个SSO。
'token_key': 'access_token', # AccessToken在ResponseBody中的名字，必须指定，用于框架保存AccessToken。
'remote_app': {
'client_id': 'superset-client', # Superset在鉴权注册的id
'client_secret': 'superset', # 配套的密钥
'client_kwargs': {
'scope': 'openid' # OAuth2的scope，多个值用空格分开
},
'access_token_method': 'POST', # 请求access token接口时的HTTP方法
'access_token_params': {
# 请求access token接口附在URL上的参数，视鉴权服务的接口规范添加。可选配置。
'client_id': 'superset-client',
},
'access_token_headers': {
# 请求access token接口附在HEADER上的参数，视鉴权服务的接口规范添加。可选配置。
'Authorization': 'Basic Base64EncodedClientIdAndSecret'
},
'api_base_url': 'http://resource-server', # 资源服务API根路径，用于获取AccessToken后请求用户信息。
'authorize_url': 'http://auth-server/oauth2/authorize', # OAuth 2.0中的authorize接口
'access_token_url': 'http://auth-server/oauth2/token', # OAuth 2.0中的token接口
}
}
]
# 是否允许创建不存在的用户。通过SSO登录的用户有可能没有保存在Superset的用户表中，如果这个配置项为False，那么用户将被拒绝登录。
AUTH_USER_REGISTRATION = True
# 创建时的默认权限，只允许一个值。
AUTH_USER_REGISTRATION_ROLE = "Admin"
DEFAULT_FEATURE_FLAGS: Dict[str, bool] = {
# 当配置项AUTH_ROLES_SYNC_AT_LOGIN为True时，每次SSO登录后会将用户信息中的角色同步至Superset数据库。
# 具体做法见下一节内容。
"AUTH_ROLES_SYNC_AT_LOGIN": False,
}
添加自定义的SecurityManager
Superset默认支持OAuth 2.0的登录方式有GitHub、Twitter、LinkedIn、Google等。但如果鉴权服务是自建的话，就需要编写配套的SecurityManager，以便返回给框架正确的用户信息。
在PYTHONPATH下添加一个新的文件：custom_sso_security_manager.py，添加一个SecurityManager继承类，覆盖oauth_user_info方法：
import jwt
from flask import session
from superset.security import SupersetSecurityManager
class CustomSsoSecurityManager(SupersetSecurityManager):
def oauth_user_info(self, provider, response=None):
if provider == 'spring-sso': # 判断SSO的名字
access_token = response.get('access_token') # 从Response中获取AccessToken
decoded = jwt.decode(access_token, verify=False) # 解析JWT
sub = decoded.get('sub') # 得到OpenId
# 向资源服务请求，通过oauth_remotes调用时，框架会自动在Authorization Header添加AccessToken。
# 这个AccessToken就是通过之前配置里的token_key解析得到的。
# 这里的路径就是之前配置里的api_base_url。
# 理论上资源服务和鉴权服务是分开的，但大部分的SSO vendor提供的获取用户信息接口与token接口的根路径是一致的。
# 这里是根据业务的需要，向资源服务获取当前用户的租户信息。
user_details_resp = self.appbuilder.sm.oauth_remotes[provider].get('tenants')
# 将租户信息保存在session中。
session["tenants"] = user_details_resp.json()
# 拼接成用户信息。
# 用户信息中必须要有username或email，否则日志会抛出异常：OAUTH userinfo does not have username or email
# 用户信息可以添加role_keys列表，作为用户的角色列表。
# 当配置项AUTH_ROLES_SYNC_AT_LOGIN为True时，每次SSO登录后都将列表中的角色同步至Superset数据库。
user_info = { 'username': sub, 'first_name': sub }
return user_info
然后再superset_config.py追加以下几行：
from custom_sso_security_manager import CustomSsoSecurityManager
CUSTOM_SECURITY_MANAGER = CustomSsoSecurityManager
运行一下吧
大功告成，可以试着运行一下，看看是否可以正常接口SSO。
自定义宏命令
开启配置
为了根据用户的租户信息对查询的数据进行过滤，需要Superset的SQL Templating和Row level security两个特性的配合。在superset_config.py中打开这两个配置：
DEFAULT_FEATURE_FLAGS: Dict[str, bool] = {
# ...
"ENABLE_TEMPLATE_PROCESSING": True,
"ROW_LEVEL_SECURITY": True,
# ...
}
清先阅读下官方文档：SQL Templating和Row level security。
目前Superset的Row level security功能是比较完备的，可以在页面上配置过滤的从句（Clause）。而且过滤从句可以被SQL Templating处理，所以这里可以写入宏命令，只是注意这里不需要写上where关键字。因此Row level security无需进行任何改造。
但是对于官方提供的宏命令，还不足以支撑业务的需要（比如一个宏命令tenants()，从session中获取当前用户的租户信息）。所以需要对其进行扩展。
添加自定义宏命令
Superset在jinja_context.py下实现了SQL Templating，对于SQL语句中的宏命令的替换处理，主要是通过JinjaTemplateProcessor来实现的，对于HQL的支持是通过HiveTemplateProcessor来实现的。后者在前者的基础上添加了一些针对分区（partition）的宏命令。
对于宏命令的扩展，可以参考Superset的教程，在superset_config.py中添加CUSTOM_TEMPLATE_PROCESSORS：
from custom_template_processor import CustomTemplateProcessor
from superset.jinja_context import BaseTemplateProcessor
from typing import Type, Dict
CUSTOM_TEMPLATE_PROCESSORS: Dict[str, Type[BaseTemplateProcessor]] = {
"sqlite": CustomTemplateProcessor
}
CUSTOM_TEMPLATE_PROCESSORS是一个Dict对象，可以理解为Java中的Map。键类型为str，代表着所负责的数据库引擎类型，在我的本地环境中，数据库使用的是sqlite，所以这里的写的是sqlite。值类型是BaseTemplateProcessor的子类，这里我自定义了一个CustomTemplateProcessor，保存在同目录的custom_template_processor.py中：
from functools import partial
from flask import session
from superset.jinja_context import JinjaTemplateProcessor, safe_proxy
from typing import Any
def tenants() -> (): return session["tenants"]
# 只需继承JinjaTemplateProcessor即可。
class CustomTemplateProcessor(JinjaTemplateProcessor):
# 官方的文档中给出的列子是将宏命令的识别由{{}}改为$，所以覆盖的是process_template。
# 现在的需要是添加新的宏命令，所以只需覆盖set_context方法即可。记得执行父类的方法！
def set_context(self, **kwargs: Any) -> None:
# 执行父类的方法。
super().set_context(**kwargs)
# 更新context
self._context.update(
{
# 键值是宏命令表达式
# 值一定要写为partial(safe_proxy, func, args)，否则父类在更新context会抛出安全异常
"tenants": partial(safe_proxy, tenants),
}
)
添加后，重启服务，就可以去Row level security添加新增的宏命令了：
tenant IN ({{ "'" + "','".join(tenants()) + "'" }})
补充说明
任何TemplateProcessor都是单例模式，所以不要在这个类中保存与请求或线程相关的状态。
目前租户信息是保存在服务session（内存）中，后期也可以优化为redis，或是持久化到Superset的数据库，在每次登录时更新下。
小结
本篇博客主要是指导如何使用Superset介入OAuth 2.0鉴权服务并从其下的资源服务获取相关信息，以及如何添加自定义的宏命令。
