Airflow中PythonOperator内调用BigQueryOperator报task_instance键错误

我尝试通过Airflow/Python实现以下操作：

• 从XCom拉取行数（例如100000）
• 基于该行数进行循环，每次处理20000条记录：先调用BigQueryExecuteQueryOperator创建数据库表，再调用BigQueryToGCSOperator将表数据导出至GCS存储桶，完成循环

为排查问题，我简化了代码，仅保留创建表的逻辑，但在PythonOperator的_get_rowcount函数中调用task3.execute(dict())时失败，报错信息：

task3.execute(dict())
  File "/home/airflow/.local/lib/python3.10/site-packages/airflow/providers/google/cloud/operators/bigquery.py", line 663, in execute
    context['task_instance'].xcom_push(key='job_id', value=job_id)
KeyError: 'task_instance'

我的default_args中已设置'provide_context': True，相关代码如下：

def _get_rowcount(ti):

    table_size = 1000000

    rowcount = ti.xcom_pull(key='row_count', task_ids='_bq_result')
    print('BQ rowcount is:', rowcount[0][0])    
    
    for x in range(0,rowcount[0][0],table_size):
        print('Iteration is:', x)
     
        SQL =  f"Create or Replace table <project_id>.<dataset>.<table_name> as select * from `<project_id>.<dataset>.<table_name>` where row_count > {x} and row_count <= { (table_size + x)} "
        print ("SQL: " + SQL)

        task3 = BigQueryExecuteQueryOperator(
            task_id = f"create_tmp_table_{x}",
            use_legacy_sql = False,
            sql = SQL,
            gcp_conn_id = GCP_CONN_ID
        )

        task3.execute(dict())

task2 = PythonOperator(
    task_id = '_get_rowcount',
    provide_context = True,
    python_callable = _get_rowcount,
    dag = dag
    )

task2

报错原因 #

直接调用Operator的execute方法时，传入的空字典没有包含Airflow执行所需的上下文信息（比如task_instance），而BigQueryExecuteQueryOperator在执行过程中需要通过上下文获取task_instance来推送XCom，因此触发KeyError。

正确实现方式 #

Airflow不推荐在一个Operator内部直接实例化并调用另一个Operator的execute方法，更合适的做法是使用BigQuery Hook直接操作或者动态生成独立任务。

方法1：使用BigQuery Hook直接执行SQL #

在PythonOperator的函数中，直接用BigQueryHook执行SQL，避免依赖Operator的上下文要求：

from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook

def _get_rowcount(ti):
    table_size = 1000000
    rowcount = ti.xcom_pull(key='row_count', task_ids='_bq_result')
    print('BQ rowcount is:', rowcount[0][0])    
    
    # 初始化BigQuery Hook
    bq_hook = BigQueryHook(gcp_conn_id=GCP_CONN_ID, use_legacy_sql=False)
    
    for x in range(0, rowcount[0][0], table_size):
        print('Iteration is:', x)
        SQL = f"Create or Replace table <project_id>.<dataset>.<table_name> as select * from `<project_id>.<dataset>.<table_name>` where row_count > {x} and row_count <= {table_size + x}"
        print("SQL: " + SQL)
        
        # 使用Hook执行SQL
        bq_hook.run(sql=SQL)

task2 = PythonOperator(
    task_id='_get_rowcount',
    provide_context=True,
    python_callable=_get_rowcount,
    dag=dag
)

方法2：动态生成独立任务（推荐） #

如果希望每个循环步骤都作为独立的Airflow任务（便于监控和重试），可以在DAG定义阶段动态生成任务：

from airflow.providers.google.cloud.operators.bigquery import BigQueryExecuteQueryOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryToGCSOperator

# 先定义获取总行数的任务
def get_total_rowcount(**context):
    rowcount = context['ti'].xcom_pull(key='row_count', task_ids='_bq_result')
    return rowcount[0][0]

get_rowcount_task = PythonOperator(
    task_id='get_total_rowcount',
    provide_context=True,
    python_callable=get_total_rowcount,
    dag=dag
)

# 动态生成任务
table_size = 1000000
# 替换为实际获取总行数的逻辑，比如从XCom拉取（需确保DAG解析时能获取到）或使用Airflow Variable
total_rows = 1000000  

previous_task = get_rowcount_task
for x in range(0, total_rows, table_size):
    # 创建临时表任务
    create_table_task = BigQueryExecuteQueryOperator(
        task_id=f"create_tmp_table_{x}",
        use_legacy_sql=False,
        sql=f"Create or Replace table <project_id>.<dataset>.<table_name> as select * from `<project_id>.<dataset>.<table_name>` where row_count > {x} and row_count <= {table_size + x}",
        gcp_conn_id=GCP_CONN_ID,
        dag=dag
    )
    
    # 导出到GCS任务
    export_to_gcs_task = BigQueryToGCSOperator(
        task_id=f"export_table_{x}_to_gcs",
        source_project_dataset_table=f"<project_id>.<dataset>.<table_name>",
        destination_cloud_storage_uris=f"gs://your-bucket/path/{x}_data.csv",
        gcp_conn_id=GCP_CONN_ID,
        dag=dag
    )
    
    # 设置任务依赖
    previous_task >> create_table_task >> export_to_gcs_task
    previous_task = export_to_gcs_task

注意事项 #

• Airflow的Operator是设计为独立任务在调度下运行的，直接调用execute会绕过上下文管理，容易出现缺失上下文的错误。
• 动态生成任务时，需确保任务ID唯一，且DAG解析阶段能获取到所需的总行数（若为动态值，可使用Airflow Variable或在解析阶段执行查询）。

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