如何通过Python的BigQuery客户端限制查询执行时长至2小时

问题背景 #

我使用Python执行BigQuery查询，部分查询的执行时间长达6小时，希望将查询的最大执行时间限制为2小时：若2小时内未完成则自动终止查询并标记为失败，能在2小时内完成则正常返回结果。

现有代码 #

utility.py #

def execute_select_bq_query(query, location, jobconfig):
    from google.cloud import bigquery

    # Construct a BigQuery client object.
    client = bigquery.Client(project=queryexcproject)
    job_config = bigquery.QueryJobConfig()
    dest_table = bq_table_reference(client, jobconfig['destination'])
    job_config.destination = dest_table
    job_config.create_disposition = jobconfig['create_disposition']
    job_config.write_disposition = jobconfig['write_disposition']
    query_job = client.query(query, location=location, job_config=job_config)
    results = query_job.result()
    response = "Data loaded into BQ table {}".format(jobconfig['destination'])
    return response

main.py #

from utility import execute_select_bq_query

def get_data(request):
    job_config = {"destination": bq_project_details,
                  "create_disposition": "CREATE_IF_NEEDED",
                  "write_disposition": tbl_write_disposition}
    response = execute_select_bq_query(query, bq_location, job_config)
    print(response)
    return response()

（注：原代码中return response()存在语法错误，应为return response）

解决方案 #

通过以下两步实现超时控制：

1. 设置查询作业超时时间 #

在QueryJobConfig中添加timeout_ms参数，指定2小时对应的毫秒数（2 * 3600 * 1000 = 7200000），BigQuery会在超时后自动终止作业。

2. 捕获超时异常并处理 #

捕获google.api_core.exceptions.DeadlineExceeded异常，在异常处理逻辑中标记任务状态为失败，同时可手动确认终止作业。

修改后的代码 #

utility.py #

def execute_select_bq_query(query, location, jobconfig):
    from google.cloud import bigquery
    from google.api_core.exceptions import DeadlineExceeded

    # 2小时超时，单位毫秒
    TIMEOUT_MS = 2 * 3600 * 1000

    client = bigquery.Client(project=queryexcproject)
    job_config = bigquery.QueryJobConfig()
    dest_table = bq_table_reference(client, jobconfig['destination'])
    job_config.destination = dest_table
    job_config.create_disposition = jobconfig['create_disposition']
    job_config.write_disposition = jobconfig['write_disposition']
    # 设置超时时间
    job_config.timeout_ms = TIMEOUT_MS

    query_job = client.query(query, location=location, job_config=job_config)
    try:
        # 等待查询完成，超时会抛出异常
        results = query_job.result()
        return f"Data loaded into BQ table {jobconfig['destination']}"
    except DeadlineExceeded:
        # 手动终止作业（可选，BigQuery超时后会自动终止）
        query_job.cancel()
        # 返回失败状态
        return f"Query failed: Execution exceeded {TIMEOUT_MS//3600000} hour timeout"

main.py #

from utility import execute_select_bq_query

def get_data(request):
    job_config = {"destination": bq_project_details,
                  "create_disposition": "CREATE_IF_NEEDED",
                  "write_disposition": tbl_write_disposition}
    response = execute_select_bq_query(query, bq_location, job_config)
    print(response)
    return response

说明 #

• timeout_ms是BigQuery服务端的超时限制，超过这个时间服务端会终止查询并返回超时错误。
• 捕获DeadlineExceeded异常后，可根据业务需求添加额外的失败处理逻辑（比如记录日志、通知告警等）。
• 原代码中return response()是语法错误，已修正为return response。

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