superset 二开增加 flink 数据源连接通过flink sql 查询数据

前言

superset 目前还不支持 flink 的数据源连接，目前我们公司在探索使用数据湖那一套东西：

使用 flink 作为计算引擎
使用 paimon + oss对象存储对接 flink 作为底层存储
使用 superset 通过 flink gateway 查询 paimon 数据形成报表

增加flink数据源

界面配置

我们通过添加其他数据源连接来增加 flink 的数据源连接。

在填写 SQL_ALCHEMY_URI 的时候这里的

driver

需要注意，后边在二开代码的时候，需要根据这个 driver 识别到不同的 engine。
我们是通过 flink gateway 提供的 HTTP 接口来进行 flink sql 查询的，所以这里的 host, port 就是 flink gateway 的地址。
在添加连接的时候必须指定 catalog，不然在 superset 的 sqllab 左下侧就没法显示对应的 databases 和 tables。
如果我们的连接需要一些额外参数，可以通过右侧的进阶添加一些额外的参数，在业务代码里使用：

我这里就指定了该连接使用的 catalog, 以及每次执行 sqllab 查询的时候初始化的一些命令。

代码开发

定义 flink.py 文件

我们需要在

superset/superset/db_engine_specs

目录下新增一个

flink.py

文件包含三个类：

FlinkClient: 用于和 flink gateway 交互执行 flink sql。
FlinkEngine: 模拟 mysql 的 cursor, 在一个 cursor 实例的生命周期内，就是和 flink gateway 的session 生命周期，当cursor 结束时，就是断开 session 的时候。
FlinkEngineSpec: 继承 superset 自身的 BaseEngineSpect, superset 的业务代码需要通过该类执行 sql 和查询结果。

FlinkClient

import logging
from typing import Any, Dict, Optional, Tuple, List, Union, Set
import time
import re

import requests
import sqlparse
from sqlalchemy import types, select
from sqlalchemy.orm import Session
from sqlalchemy.sql import text
from sqlalchemy.engine import  Engine

from superset.models.core import Database
from superset.config import FLINK_HOST
from superset.db_engine_specs.base import BaseEngineSpec
from superset.models.sql_lab import Query

logger = logging.getLogger(__name__)classFlinkClient:
    result_type ={"NOT_READY":"NOT_READY",# 表明 sql 还在执行中"PAYLOAD":"PAYLOAD",# 表明 sql 已经在 flink 集群上执行了，需要 client 循环获取结果"EOS":"EOS"# 表明已经获取到 sql 执行结果了，可以退出循环}

    result_kind ={"SUCCESS_WITH_CONTENT":"SUCCESS_WITH_CONTENT",# 执行的是查询结果的 sql"SUCCESS":"SUCCESS"# 执行的是命令}def__init__(self,**kwargs):
        self.session_id =None
        self.operation_ids =[]# 添加连接时在额外参数中填写的初始化命令# 在执行 sql 前会先执行话初始化命令
        self.init_commands = kwargs.get("init_commands",[])# FLINK_HOST 就是 flink gateway 的地址，我是从环境变量中获取的
        self.get_session_url = FLINK_HOST +"/v1/sessions","POST"
        self.execute_statement_url = FLINK_HOST +"/v1/sessions/{SESSION_ID}/statements/","POST"
        self.fetch_result_url = FLINK_HOST +"/v1/sessions/{SESSION_ID}/operations/{OPERATION_ID}/result/{BATCH_NUM}","GET"

        self.kwargs = kwargs

    def__enter__(self):# 使用上下文模式，调用的时候获取 session 和执行初始化命令
        self.get_session()for c in self.init_commands:
            operation_id = self.execute(c)
            self.fetch_result(operation_id=operation_id)return self

    def__exit__(self, exc_type, exc_value, traceback):if exc_type isnotNone:
            logger.error(f"flink gateway got error: {exc_type}, {exc_value}")returnFalsedefhandle_request(self,
                       url:str,
                       method:str,
                       form_data: Dict[str, Any]=None,
                       json_data: Dict[str, Any]=None,
                       params: Dict[str, Any]=None,
                       headers: Dict[str, Any]=None,
                       timeout: Tuple[int,...]=(10,60))-> Dict[str, Any]:try:
            kwargs ={"timeout": timeout,"headers":{"Content-Type":"application/json"}}if form_data:
                kwargs["data"]= form_data
            if params:
                kwargs["params"]= params
            if json_data:
                kwargs["json"]= json_data
            if headers:
                kwargs["headers"].update(headers)# logger.info(f"request to flink gateway url: {url}")# logger.info(f"request to flink gateway kwargs: {kwargs}")if method =='GET':
                response = requests.get(url,**kwargs)elif method =='POST':
                response = requests.post(url,**kwargs)elif method =='PUT':
                response = requests.put(url,**kwargs)elif method =='DELETE':
                response = requests.delete(url,**kwargs)else:raise ValueError("Unsupported HTTP method")

            response.raise_for_status()

            res ={'status_code': response.status_code,'headers':dict(response.headers),'data': response.json()}# logger.info(f"flink gateway response: {res}")return res
        except Exception as e:
            logger.error(f"flink gateway res error: {str(e)}")raise e

    defget_session(self):
        res = self.handle_request(self.get_session_url[0], self.get_session_url[1])
        self.session_id = res['data'].get('sessionHandle')defping(self):
        operation_id = self.execute("select 1")returnTrueif self.fetch_result(operation_id=operation_id)elseFalsedefexecute(self, statement:str):# 执行 flink sql
        data ={"statement": statement}
        res = self.handle_request(
            self.execute_statement_url[0].format(SESSION_ID=self.session_id),
            self.execute_statement_url[1],
            json_data=data)
        self.operation_ids.append(res['data'].get('operationHandle'))return res['data'].get('operationHandle')deffetch_result(self, batch_num:int=0, operation_id:str=None)-> Dict[str, Any]:"""
        通过 flink gateway 获取执行结果：gateway 将 sql 提交至集群后返回 PAYLOAD 状态表示提交成功，
            否则返回 NOT_READY 状态。
        当提交至集群成功后，如果执行的 sql 是查询内容的，需要通过 batch_num（nextResultUri） 不断循环请求执行结果，
        直到 gateway 返回 EOS 状态，表示集群执行完毕，获取结果完毕。
        """
        url = self.fetch_result_url[0].format(SESSION_ID=self.session_id,
                                              OPERATION_ID=operation_id,
                                              BATCH_NUM=batch_num)
        res_data =[]
        res = self.handle_request(url, self.fetch_result_url[1])#  后续考虑是否做成从环境变量中获取超时时间，且超时后是否考虑杀死集群上执行的任务
        timeout =300# flink gateway 提交任务至 session 集群超时时间while timeout and res['data']['resultType']== self.result_type['NOT_READY']:
            time.sleep(1)
            timeout -=1
            res = self.handle_request(url, self.fetch_result_url[1])# 等待集群执行完毕，获取结果if res['data']['resultKind']== self.result_kind['SUCCESS_WITH_CONTENT']:
            timeout =3600# flink gateway 从集群获取结果超时时间while timeout and res['data']['resultType']!= self.result_type['EOS']:
                time.sleep(3)
                timeout -=1
                res_data.extend(res['data']['results']['data'])
                logger.info(f"jobID: {res['data'].get('jobID')} waiting for result")
                next_result_url = FLINK_HOST + res['data']['nextResultUri']
                res = self.handle_request(next_result_url, self.fetch_result_url[1])
            res['data']['results']['data']= res_data
        return res['data']['results']defget_schema_names(self, catalog:str)-> List[str]:
        operation_id = self.execute(f"use catalog {catalog}")
        self.fetch_result(operation_id=operation_id)
        operation_id = self.execute(f"show databases")
        res = self.fetch_result(operation_id=operation_id)return[i['fields'][0]for i in res['data']]defget_table_names(self, catalog:str, schema:str)-> List[str]:
        operation_id = self.execute(f"use catalog {catalog}")
        self.fetch_result(operation_id=operation_id)
        operation_id = self.execute(f"use {schema}")
        self.fetch_result(operation_id=operation_id)
        operation_id = self.execute("show tables")
        res = self.fetch_result(operation_id=operation_id)return[i['fields'][0]for i in res['data']]defget_columns(self, catalog:str, schema:str, table_name:str)-> List[str]:
        operation_id = self.execute(f"use catalog {catalog}")
        self.fetch_result(operation_id=operation_id)
        operation_id = self.execute(f"use {schema}")
        self.fetch_result(operation_id=operation_id)
        operation_id = self.execute(f"desc {table_name}")
        res = self.fetch_result(operation_id=operation_id)return[i['fields']for i in res['data']]

这里与 flink gateway 交互获取结果有个需要注意的地方就是，我们将执行 sql 提交至 flink gateway 后, gateway resultType 会很快返回

PAYLOAD

状态，这个时候不代表 sql 执行完了，代表的是集群在执行中了，我们可以阻塞获取执行结果了，然后我们在阻塞获取结果，当状态变为

EOS

的时候，代表我们获取到了结果了，这个时候可以退出阻塞了。
官方的流程图说明如下：

这里需要了解下我们的客户端通过 HTTP 接口与 gateway 交互的流程，不熟悉的可以先通过官方文档了解下：
https://nightlies.apache.org/flink/flink-docs-master/zh/docs/dev/table/sql-gateway/overview/

FlinkEngine
FlinkEngine 模拟的是类似 sql dialect 的 cursor, 通过该方法可以返回一个连接 flink gateway 的 client，在 cursor 的整个实例生命周期内使用的是同一个 FlinkClient 的 session。

classFlinkEngine:def__init__(self, catalog:str, schema:str=None,
                 init_commands: List[str]=None,**kwargs):
        self.catalog = catalog
        self.schema = schema
        self.init_commands = init_commands if init_commands else[]
        self.kwargs = kwargs

        self.client: Optional[FlinkClient]=None

        self.columns =None@propertydefengine(self)->'FlinkEngine':return self

    defraw_connection(self)->'FlinkEngine':# 实例化 flink clint 生成一个 session, 当 cursor 结束时重置 clientwith FlinkClient(catalog=self.catalog, init_commands=self.init_commands)as c:
            self.client = c
        return self

    defcursor(self)->'FlinkEngine':"""Return a new :py:class:`Cursor` object using the connection."""return self

    defclose(self):
        self.client =None
        self.columns =Nonedefcommit(self):"""Presto does not support transactions"""pass@propertydefdescription(self):"""This read-only attribute is a sequence of 7-item sequences.

           Each of these sequences contains information describing one result column:

           - name
           - type_code
           - display_size (None in current implementation)
           - internal_size (None in current implementation)
           - precision (None in current implementation)
           - scale (None in current implementation)
           - null_ok (always True in current implementation)

           The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the
           section below.
           """if self.columns isNone:returnNonereturn[# name, type_code, display_size, internal_size, precision, scale, null_ok(col['name'], col['type'],None,None,None,None,True)for col in self.columns
        ]

FlinkEngineSpec
该方法需要继承 superset 的

BaseEngineSpec

, 需要定义 engine 信息和 drivers 信息, 在 superset 的 sqllab 执行 sql 的时候会通过 drivers 定位到该方法执行。
所以在前边界面配置的时候需要注意连接信息中的 driver_name 要和类属性 drivers 匹配。

classFlinkEngineSpec(BaseEngineSpec):
    engine ="flink"
    engine_name ="Apache Flink"# 我们后边在业务代码中会通过判断连接的 driver_name 是否为 flink_driver 来调用该类中的方法# 因此需要注意前端界面配置是否一致
    drivers ={"flink_driver":"flink gateway engine"}
    default_driver ="flink_driver"

    client_init_commands =[]@classmethoddefget_schema_names(cls, catalog:str)-> List[str]:with FlinkClient(init_commands=cls.client_init_commands)as c:
            names = c.get_schema_names(catalog)return names

    @classmethoddefget_table_names(cls, catalog:str, schema:str, database=None)-> List[str]:with FlinkClient(init_commands=cls.client_init_commands)as c:
            names = c.get_table_names(catalog, schema)return names

    @classmethoddefget_view_names(cls, catalog:str, schema:str, database=None)-> Set[str]:returnset()@classmethoddefget_columns(cls, catalog:str, schema:str, table:str)-> List[Dict[str, Any]]:with FlinkClient(init_commands=cls.client_init_commands)as c:
            cs = c.get_columns(catalog, schema, table)
        result: List[Dict[str, Any]]=[]for column in cs:
            column_spec = cls.get_column_spec(column[1])
            column_type = column_spec.sqla_type if column_spec elseNoneif column_type isNone:
                column_type = types.String()
            c ={"name": column[0],"type": column_type,"nullable": column[2],"default":None,"key": column[3]}try:
                c.update({"comment": column[6]})except Exception:pass
            result.append(c)return result

    @classmethoddefget_pk_constraint(cls, catalog:str, schema:str, table:str)-> Dict[str, Any]:with FlinkClient(init_commands=cls.client_init_commands)as c:
            cs = c.get_columns(catalog, schema, table)
        pks ={"constrained_columns":None,"name":None}for column in cs:
            _type = column[3]ifisinstance(_type,str)and _type.startswith("PRI"):
                matches = re.findall(r'\((.*?)\)', _type)
                pks["constrained_columns"]=[field.strip()for field in matches[0].split(',')]breakreturn pks

    @classmethoddefselect_star(# pylint: disable=too-many-arguments,too-many-locals
        cls,
        database: Database,
        table_name:str,
        engine: Engine,
        schema: Optional[str]=None,
        limit:int=100,
        show_cols:bool=False,
        indent:bool=True,
        latest_partition:bool=True,
        cols: Optional[List[Dict[str, Any]]]=None,)->str:
        fields: Union[str, List[Any]]="*"
        cols = cols or[]if(show_cols or latest_partition)andnot cols:
            cols = database.get_columns(table_name, schema)if show_cols:
            fields = cls._get_fields(cols)if schema:
            full_table_name =f"{schema}.{table_name}"else:
            full_table_name =f"{table_name}"

        qry = select(fields).select_from(text(full_table_name))if limit:
            qry = qry.limit(limit)if latest_partition:
            partition_query = cls.where_latest_partition(
                table_name, schema, database, qry, columns=cols
            )if partition_query isnotNone:
                qry = partition_query

        sql =str(qry.compile(compile_kwargs={"literal_binds":True}))if indent:
            sql = sqlparse.format(sql, reindent=True)return sql

    @classmethoddefexecute(# pylint: disable=unused-argument
        cls,
        cursor: FlinkEngine,
        query:str,**kwargs: Any,)->None:"""执行 flink sql 语句"""return cursor.client.execute(query)@classmethoddefhandle_cursor(cls, cursor: FlinkClient, query: Query, session: Session)->None:"""
        在执行 flink sql 执行过程中，执行一些动作：
        记录flink sql 任务的一些关键信息
        记录一些执行日志
        sleep 等待执行结果 等
        """return@classmethoddeffetch_data(
        cls, cursor: FlinkEngine, limit: Optional[int]=None)-> List[Tuple[Any,...]]:
        res = cursor.client.fetch_result(operation_id=cursor.client.operation_ids[-1])
        cursor.columns =[{"name": i['name'],"type": i["logicalType"]["type"]}for i in
                          res.get('columns',[])]return[tuple(i['fields'])for i in res['data']]@classmethoddefhas_implicit_cancel(cls)->bool:"""
        该方法是sqllab 界面执行 sql 中点击暂停时调用的
        这里直接返回了 True, 因为 gateway 的 session 自己有过期时间
        我们也可以通过调用 gateway 的关闭 session 接口主动关闭
        """returnTrue@classmethoddefcancel_query(# pylint: disable=unused-argument
        cls,
        cursor: FlinkClient,
        query: Query,
        cancel_query_id:str,)->bool:"""
        该方法是sqllab 界面执行 sql 中点击暂停时调用的
        这里直接返回了 True, 因为 gateway 的 session 自己有过期时间
        我们也可以通过调用 gateway 的关闭 session 接口主动关闭
        """returnTrue

修改测试连接逻辑

测试连接入口方法在

superset/databases/commands/test_connection.py

下的

TestConnectionDatabaseCommand

类中的 run 方法，我们需要通过连接的 driver 来通过 FLinkClient 测试与 Flink gateway 的连接是否正常：

# flink 类型的连接走 flink gateway 验证if database.driver == FLINK_DRIVER_NAME:from superset.db_engine_specs.flink import FlinkClient
    init_commands = database.get_encrypted_extra().get("init_commands",[])with FlinkClient(init_commands=init_commands)as c:ifnot c.ping():raise Exception("ping flink gateway err")return

修改 sqllab 界面逻辑

sqllab 界面需要修改获取库表信息和执行 sql 的接口逻辑：

查询库表字段信息的接口入口类都在

superset/superset/databases/api.py

中:

api 入口的代码逻辑不需要修改。
获取库名称列表，修改

superset/superset/models/core.py

中

Database

类中的

get_all_schema_names

方法：

# flink 连接不走 sqlalchemy 的 create engine, 属于 FlinkEngineSpecif self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])return self.db_engine_spec.get_schema_names(extra['catalog'])

获取表名称列表，修改

superset/superset/models/core.py

中

Database

类中的

get_all_table_names_in_schema

方法和

get_all_view_names_in_schema

方法：

# flink 连接不走 sqlalchemy 的 create engine, 属于 FlinkEngineSpecif self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])
    tables ={(table, schema)for table in self.db_engine_spec.get_table_names(
            extra['catalog'],
            schema
        )}return tables

if self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])return{(view, schema)for view in self.db_engine_spec.get_view_names(
            extra['catalog'],
            schema
        )}

获取字段信息，修改

superset/superset/models/core.py

中

Database

类中的

get_columns

方法：

if self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])return self.db_engine_spec.get_columns(extra["catalog"], schema, table_name)

获取表 comment 信息，修改

get_table_comment

方法，这个目前还没有找到通过 flink sql 查询表 comment 信息的方法，这里直接返回空：

if self.driver == FLINK_DRIVER_NAME:return""

获取索引信息，修改

get_indexes

方法，返回空列表：

if self.driver == FLINK_DRIVER_NAME:return[]

获取主键信息，修改

get_pk_constraint

方法：

if self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])return self.db_engine_spec.get_pk_constraint(extra["catalog"], schema, table_name)

获取外键信息，修改

get_foreign_keys

方法：

if self.driver == FLINK_DRIVER_NAME:return[]

执行 sql 相关的需要修改

_get_sqla_engine

方法。

# VOYAH 如果是 flink_driver 就使用 FlinkEngineSpec.engineif self.driver == FLINK_DRIVER_NAME:
    extra = self.get_encrypted_extra()
    self.db_engine_spec.client_init_commands = extra.get("init_commands",[])from superset.db_engine_specs.flink import FlinkEngine
    return FlinkEngine(schema=schema,**params)

总结

增加其他数据源连接，主要需要修改两个文件新增一个文件：

修改 superset/databases/commands/test_connection.py中的TestConnectionDatabaseCommand 类中的 run 方法。修改 superset/superset/models/core.py中 Database类中的get_all_table_names_in_schema，get_all_view_names_in_schema，get_columns，get_table_comment， get_indexes，get_pk_constraint，get_foreign_keys，_get_sqla_engine 方法。
在 superset/superset/db_engine_specs目录下新增一个 flink.py文件。

标签： superset 自定义数据源 superset

本文转载自: https://blog.csdn.net/qq_42586468/article/details/136946220
版权归原作者 一切如来心秘密 所有，如有侵权，请联系我们删除。

superset 二开增加 flink 数据源连接通过flink sql 查询数据

前言

增加flink数据源

界面配置

代码开发

定义 flink.py 文件

修改测试连接逻辑

修改 sqllab 界面逻辑

总结

发表评论

“superset 二开增加 flink 数据源连接通过flink sql 查询数据”的评论:

关于作者

overfit同步小助手

相关阅读

文章导航