如何在Google Composer的Airflow DAG中用BigQueryInsertJobOperator执行多查询

问题根源 #

BigQueryInsertJobOperator的sql参数不支持传入查询列表，它仅会解析并执行列表中的第一个查询且无报错，这是该Operator的设计逻辑导致的。同时由于后续查询包含DECLARE语句，无法将所有查询合并为单个字符串执行（DECLARE属于会话级语句，合并后会因上下文冲突执行失败）。

解决方案：循环创建Operator并构建依赖链 #

通过循环遍历查询列表，为每个查询创建独立的BigQueryInsertJobOperator实例，并用变量维护任务间的依赖关系，实现顺序执行。

代码示例

from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from datetime import datetime

# 定义包含DECLARE语句的独立查询列表
queries = [
    "SELECT * FROM `your-project.your-dataset.table1`",
    """
    DECLARE target_id INT64 DEFAULT 100;
    INSERT INTO `your-project.your-dataset.table2` (id, status) VALUES(target_id, 'processed');
    """,
    "UPDATE `your-project.your-dataset.table3` SET update_time = CURRENT_TIMESTAMP() WHERE id = 100"
]

default_args = {
    'start_date': datetime(2024, 1, 1),
    'retries': 1
}

with DAG('bq_sequential_queries_dag', default_args=default_args, schedule_interval='@daily') as dag:
    previous_task = None
    
    for query_idx, query_content in enumerate(queries):
        # 为每个查询创建独立的Operator
        current_task = BigQueryInsertJobOperator(
            task_id=f'execute_bq_query_{query_idx}',
            configuration={
                "query": {
                    "query": query_content,
                    "useLegacySql": False
                }
            },
            gcp_conn_id='google_cloud_default'  # 替换为你的GCP连接ID
        )
        
        # 构建顺序依赖：当前任务依赖上一个任务完成
        if previous_task:
            previous_task >> current_task
        
        # 更新上一个任务的引用，用于下一次循环
        previous_task = current_task

代码说明

• 每个查询对应一个独立的BigQueryInsertJobOperator，通过task_id后的索引区分任务，避免ID重复。
• 使用previous_task变量跟踪前一个任务实例，每次循环中让当前任务依赖前一个任务，形成任务0 >> 任务1 >> 任务2的顺序执行链。
• 每个查询作为独立的BigQuery Job执行，DECLARE语句的会话上下文独立，不会出现执行错误。

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