如何通过python-bigquery-sqlalchemy获取BigQuery多子作业查询结果？

当使用python-bigquery-sqlalchemy执行BigQuery多语句脚本（比如包含循环分支的BEGIN...END脚本）时，默认只会返回第一个子作业的结果集，其余子作业的结果会被丢弃。这是因为BigQuery会将这类脚本拆分为多个独立的子作业执行，而sqlalchemy的默认查询流程只获取父作业的初始结果。

解决方案思路 #

要获取所有子作业的结果，需要：

1. 触发多语句脚本执行，获取父作业ID
2. 通过BigQuery Client列出该父作业下的所有子作业
3. 遍历每个子作业，提取结果并转换为ORM模型实例
4. 合并所有结果集

具体实现代码 #

from google.cloud import bigquery
from sqlalchemy.orm import Session
# 替换为你的ORM模型和Session实例
from your_models import MyOrmTable
from your_db_setup import db

def get_multi_statement_results(db_session: Session, orm_model, query_script):
    # 1. 执行脚本，触发BigQuery作业并等待完成
    # 初始查询会返回第一个结果集，同时触发整个脚本执行
    initial_results = db_session.query(orm_model).from_statement(query_script).all()
    
    # 2. 获取底层BigQuery Client（复用sqlalchemy的配置）
    bq_client = db_session.bind.raw_connection().client
    
    # 3. 直接用client执行脚本，明确获取父job
    parent_job = bq_client.query(query_script)
    parent_job.result()  # 等待整个脚本执行完成
    
    # 4. 遍历所有子作业，收集结果
    full_results = list(initial_results)
    for child_job in bq_client.list_jobs(parent_job=parent_job.job_id):
        # 仅处理已完成的查询类型子作业
        if child_job.job_type == "QUERY" and child_job.state == "DONE":
            for row in child_job.result():
                # 将BigQuery行对象转换为ORM模型实例
                row_dict = dict(row.items())
                model_instance = orm_model(**row_dict)
                full_results.append(model_instance)
    
    return full_results

# 使用示例
multi_statement_query = """
BEGIN
    FOR record IN 
    (
    SELECT num FROM UNNEST(GENERATE_ARRAY(1, 5)) AS num
    )
    DO
    WITH numbers AS (
        SELECT num
        FROM UNNEST(GENERATE_ARRAY(1, 100)) AS num
    )
    SELECT num
    FROM numbers
    LIMIT 10;
    END FOR;
END
"""

# 获取所有子作业的结果
all_results = get_multi_statement_results(db, MyOrmTable, multi_statement_query)
for item in all_results:
    print(item.num)

关键说明 #

• BigQuery子作业机制：多语句脚本中的每个SELECT都会生成独立的子作业，结果存储在对应子作业中，而非父作业。
• Client复用：通过sqlalchemy引擎的原始连接获取BigQuery Client，确保和ORM使用相同的项目、认证配置。
• 结果转换：将BigQuery返回的行对象转为字典后，传入ORM模型构造函数，生成符合预期的实例。
• 作业过滤：只处理已完成的查询类型子作业，避免获取脚本控制语句（如循环、分支）的无效作业。

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