Superset SSO改造和自定义宏命令

背景
关于Superset
需要解决的问题
定制化改造
小结

背景

在最近的一个项目上，客户想要为他们的多租户（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鉴权服务并从其下的资源服务获取相关信息，以及如何添加自定义的宏命令。