Airflow中如何获取BigQueryOperator查询结果并传递给下游任务

首先明确：原生BigQueryOperator本身不支持直接返回查询结果集，它的设计定位是执行DML/DDL语句、将查询结果写入指定目标表，仅会返回BQ作业ID作为输出，无法直接从中获取查询到的行数据。

要实现省略临时表、GCS中转的需求，可采用以下两种可行方案：

方案1：PythonOperator直接调用BQ客户端（更推荐） #

该方案全程不依赖临时表、GCS存储，直接在Python任务中执行查询、获取结果、生成CSV，适合绝大多数场景，尤其数据量较大的场景：

from airflow.operators.python import PythonOperator
from google.cloud import bigquery
import csv

# 自定义查询并生成CSV的函数
def bq_query_to_csv(**context):
    # Cloud Composer默认继承环境服务账号权限，无需额外配置密钥
    bq_client = bigquery.Client()
    # 执行你的查询语句
    query_job = bq_client.query(query_sql)
    rows = query_job.result()
    
    # 生成本地CSV，路径用Composer共享目录即可直接给邮件附件用
    csv_path = '/home/airflow/gcs/data/audits.csv'
    with open(csv_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        # 写入表头
        writer.writerow([field.name for field in rows.schema])
        # 逐行写入查询结果
        for row in rows:
            writer.writerow(list(row))

# DAG中定义任务
with models.DAG('reporte_prueba',
    schedule_interval='0 1 * * *', # 替换成你需要的每日定时表达式
    default_args=default_dag_args) as dag:

    query_and_gen_csv = PythonOperator(
        task_id='query_and_gen_csv',
        python_callable=bq_query_to_csv,
        provide_context=True
    )

    email_summary = email_operator.EmailOperator(
        task_id='email_summary',
        to=['aa@bb.cl'],
        subject="""Reporte de Auditorías Diarias 
        Institución: {institution_report} día {date_report}
        """.format(date_report=date,institution_report=institution),
        html_content="""
        Sres.
        <br>
        Adjunto enviamos archivo con Reporte Transacciones Diarias.
        <br>
        """,
        files=['/home/airflow/gcs/data/audits.csv']
    )

    # 依赖关系仅需一行
    query_and_gen_csv >> email_summary

方案2：使用BigQueryGetDataOperator通过XCom传递结果 #

如果偏好使用官方Operator实现，可采用专门用于拉取BQ数据的BigQueryGetDataOperator，它会自动将查询结果推送到XCom供下游任务读取：

注意：XCom默认有存储大小限制（单条默认64KB，可调整但不建议超过10MB），该方案仅适合小数据量查询场景。

from airflow.providers.google.cloud.operators.bigquery import BigQueryGetDataOperator
from airflow.operators.python import PythonOperator
import csv

def gen_csv_from_xcom(**context):
    # 从XCom拉取BQ查询结果
    query_result = context['ti'].xcom_pull(task_ids='get_bq_data')
    # 生成CSV逻辑和方案1一致
    csv_path = '/home/airflow/gcs/data/audits.csv'
    with open(csv_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        # 可自行传入表头，或者从operator返回的附加信息中获取
        writer.writerow(['表头1', '表头2', '...'])
        writer.writerows(query_result)

# DAG中的任务定义
get_bq_data = BigQueryGetDataOperator(
    task_id='get_bq_data',
    sql=query_sql,
    use_legacy_sql=False,
    location='你的BQ资源所在区域，例如us-central1'
)

gen_csv = PythonOperator(
    task_id='gen_csv',
    python_callable=gen_csv_from_xcom,
    provide_context=True
)

# 依赖关系
get_bq_data >> gen_csv >> email_summary

注意事项 #

• 若查询结果超过10MB，优先选择方案1，避免XCom存储溢出导致任务失败
• Cloud Composer环境默认已预安装google-cloud-bigquery依赖，无需额外安装Python包
• 方案1可完全替代原有5个冗余任务，大幅简化DAG逻辑，降低任务失败概率

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