如何通过Databricks Python SQL连接器获取查询ID与执行时间？

方案1：直接从Cursor获取Query ID（推荐） #

从databricks-sql-connector 2.0.0版本开始，执行cursor.execute(query)后，cursor对象会自动生成query_id属性，直接调用即可拿到当前查询的ID。拿到ID后，调用Databricks API的GET /api/2.0/sql/queries/{query_id}端点，就能获取该查询的精确执行时间。示例代码：

from databricks import sql

with sql.connect(server_hostname="your-host", http_path="your-path", access_token="your-token") as conn:
    with conn.cursor() as cursor:
        cursor.execute("SELECT * FROM your_table LIMIT 10")
        # 获取当前查询的ID
        query_id = cursor.query_id
        print(f"当前查询ID: {query_id}")
        
        # 后续通过API传入query_id获取执行时间
        # 示例API调用逻辑（需自行实现HTTP请求）
        # response = requests.get(f"https://your-host/api/2.0/sql/queries/{query_id}", headers={"Authorization": "Bearer your-token"})
        # execution_time = response.json().get("execution_time")

方案2：客户端侧计时（快速但有误差） #

如果不需要集群端的精确执行时间，仅需记录从发起请求到拿到结果的总耗时，可以直接用Python的time模块在execute前后计时：

import time
from databricks import sql

with sql.connect(...) as conn:
    with conn.cursor() as cursor:
        start_time = time.perf_counter()
        cursor.execute("SELECT * FROM your_table LIMIT 10")
        end_time = time.perf_counter()
        elapsed = end_time - start_time
        print(f"查询总耗时（含网络传输、客户端等待）: {elapsed:.2f} 秒")

注意：该时间包含网络传输、客户端等待等额外耗时，和集群实际执行查询的时间存在差异。

方案3：通过日志提取Query ID #

如果你的连接器版本较低，没有query_id属性，可以配置日志级别为DEBUG，连接器会在执行查询时打印包含Query ID的日志，通过自定义日志处理器捕获并提取：

import logging
from databricks import sql

# 配置日志
logger = logging.getLogger("databricks.sql")
logger.setLevel(logging.DEBUG)

# 自定义处理器捕获Query ID
query_id = None
class QueryIDFilter(logging.Filter):
    def filter(self, record):
        nonlocal query_id
        msg = record.getMessage()
        if "Query ID:" in msg:
            query_id = msg.split("Query ID:")[1].strip()
        return True

handler = logging.StreamHandler()
handler.addFilter(QueryIDFilter())
logger.addHandler(handler)

# 执行查询
with sql.connect(...) as conn:
    with conn.cursor() as cursor:
        cursor.execute("SELECT * FROM your_table LIMIT 10")
        print(f"提取到的Query ID: {query_id}")

方案4：通过查询注释匹配API结果 #

如果必须使用list queries API，可以在查询语句中加入唯一标识的注释（比如UUID），然后在API返回的查询列表中匹配该注释，找到对应的查询执行时间：

import uuid
from databricks import sql

# 生成唯一标识
unique_id = str(uuid.uuid4())
# 给查询添加带唯一标识的注释
query = f"/* query_unique_id: {unique_id} */ SELECT * FROM your_table LIMIT 10"

# 执行查询
with sql.connect(...) as conn:
    with conn.cursor() as cursor:
        cursor.execute(query)

# 调用list queries API，遍历结果匹配唯一标识
# 示例API调用逻辑（需自行实现HTTP请求）
# response = requests.get("https://your-host/api/2.0/sql/queries", headers={"Authorization": "Bearer your-token"})
# for q in response.json()["res"]:
#     if unique_id in q["statement_text"]:
#         print(f"查询执行时间: {q['execution_time']}")

内容的提问来源于stack exchange，提问作者manucorujo