SQLLab 计算引擎对接（二次开发）

1. 实现原理

用户在 SQLLab 提交 SQL 并选定计算引擎，CubeStudio 据此发起 异步任务，真正的执行跑在 Celery worker 中，结果落到分布式存储的 CSV 文件，前端再轮询状态、拉取结果。

要点（限制）：SQLLab 只对接 查询引擎——通过网络请求把 SQL 发出去拿结果，自身 不能作为 driver 端、不能监听端口，因此无法作为 Spark 的 driver，只适合对接 ClickHouse、Presto、Impala 等查询网关，以及 MySQL/PostgreSQL 等可直连查询的库。

2. 入口视图与接口

视图类 Sqllab_Query_View，route_base = '/idex'（myapp/views/view_sqllab.py:102-103）。对外接口（全部挂在 /idex 下）：

注意：原始文档把要实现的方法写成 check_task，但视图函数叫 check_task、它真正调用的引擎方法叫 check_task_status（myapp/views/view_sqllab.py:175-181）。二开时要实现/重写的是 check_task_status。

引擎参数配置 sqllab_config

/idex/config（myapp/views/view_sqllab.py:105-157）返回界面上可选的引擎及其 db_uri 示例（db_uri_demo），用户最近 30 天成功跑过的连接也会被合并进下拉。 若配置文件里定义了 SQLLAB，则直接用它覆盖返回的 result（view_sqllab.py:154-155）。另有 SQLLAB_ARGS，会在提交时合并进请求参数（view_sqllab.py:166-167）。

3. 计算引擎的注册与对接

3.1 引擎注册表

所有引擎在 view_sqllab.py:31-40 的 engine_impls 字典里注册，新引擎在此添加：

from myapp.utils.sqllab.base_impl import Base_Impl

engine_impls = {
    'mysql': Base_Impl(),
    'presto': Base_Impl(),
    'clickhouse': Base_Impl(),
    'postgres': Base_Impl(),
    'dm': Base_Impl(),
    'gauss': Base_Impl(),
    'impala': Base_Impl(),
    'hive': Base_Impl(),
}

当前所有引擎都直接复用 Base_Impl（它在 Celery worker 里用 SQLAlchemy create_engine 直连查询，见 myapp/utils/sqllab/base_impl.py:41-109）。要给某引擎定制逻辑，就 派生 Base_Impl 并把对应 value 换成你的子类实例。

3.2 派生 Base_Impl

Base_Impl 在 myapp/utils/sqllab/base_impl.py:112 起。对接新引擎时派生它、按需重写下面几个方法即可：

返回字段要求对应 myapp/views/view_sqllab.py 各接口里的 res_keys（:171、:180、:188、:196、:204）——返回 dict 必须至少包含这些 key，否则报"返回值异常，检查引擎实现"。

from myapp.utils.sqllab.base_impl import Base_Impl

# 自定义某查询引擎的对接逻辑
class My_Engine_Impl(Base_Impl):
    # 提交查询任务，返回 {"err_msg": "", "task_id": qid}
    def submit_task(self, qid, enable_async=True):
        ...
        return {"err_msg": "", "task_id": qid}

    # 查询任务状态
    def check_task_status(self, qid):
        ...
        return {"stage": "...", "state": "...", "err_msg": "",
                "spark_log_url": "", "spark_ui_url": ""}

    # 取结果（二维数组：首行表头 + 数据行）
    def get_result(self, task_id):
        ...
        return {"err_msg": "", "result": [...]}

    # 生成下载地址
    def download_url(self, task_id):
        ...
        return {"err_msg": "", "download_url": "..."}

    # 终止任务
    def stop(self, task_id):
        ...
        return {"err_msg": ""}

然后在 engine_impls 里注册：'myengine': My_Engine_Impl()。

3.3 Base_Impl 默认行为参考

• 异步执行靠 Celery 任务 handle_task（任务名 task.idex.handle_base_task，base_impl.py:41），按 engine_arg1 用 SQLAlchemy create_engine 连库查询。
• 结果 CSV 落盘到 /data/k8s/kubeflow/global/sqllab/result/{qid}.csv（base_impl.py:67）；下载地址走 /static/global/sqllab/result/...（base_impl.py:204）。
• get_result 把 CSV 读成二维数组（首行列名 + 数据行）。
• stop 默认未实现远程 kill（base_impl.py:215-217），返回提示信息——需要时在子类里实现。
• 默认要求查询 SQL 必须包含 limit（base_impl.py:53-54）。

4. 数据模型

任务记录表 sqlab_query（注意拼写是单 l，myapp/models/model_sqllab_query.py:15）。关键字段：engine_arg1（引擎名）、engine_arg2（连接串/数据库）、qsql（SQL）、stage/status（阶段/状态）、err_msg、result_line_num、log_url/ui_url/result_url、username。

相关源码