如何在Airflow DAG中用Python Operator而非BigQuery Operator插入数据？

完整替换方案 #

下面是用PythonOperator替代GoogleCloudStorageToBigQueryOperator实现相同功能的DAG代码：

from airflow import DAG
from airflow.operators.python import PythonOperator
from google.cloud import bigquery
from datetime import datetime

def load_gcs_to_bq():
    # 初始化BigQuery客户端
    client = bigquery.Client()
    
    # 对应原Operator的参数
    bucket = 'cloud-samples-data'
    source_objects = ['bigquery/us-states/us-states.csv']
    destination_table = 'airflow_test.gcs_to_bq_table'
    schema_fields = [
        {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'post_abbr', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
    write_disposition = 'WRITE_TRUNCATE'
    
    # 构建GCS数据源URI（多个文件用逗号分隔）
    source_uris = [f"gs://{bucket}/{obj}" for obj in source_objects]
    
    # 配置加载参数
    job_config = bigquery.LoadJobConfig(
        schema=schema_fields,
        write_disposition=write_disposition,
        # CSV含表头时启用该配置，与原Operator默认逻辑对齐
        skip_leading_rows=1,
        source_format=bigquery.SourceFormat.CSV
    )
    
    # 执行加载任务
    load_job = client.load_table_from_uri(
        source_uris,
        destination_table,
        job_config=job_config,
        location="US"  # 需与目标表所在区域一致
    )
    
    # 等待任务完成
    load_job.result()
    print(f"数据已成功导入到表 {destination_table}")

# 定义DAG
with DAG(
    dag_id='gcs_to_bq_python_operator',
    schedule_interval=None,
    start_date=datetime(2024, 1, 1),
    catchup=False
) as dag:
    # 替换原GoogleCloudStorageToBigQueryOperator为PythonOperator
    load_csv_task = PythonOperator(
        task_id='gcs_to_bq_example',
        python_callable=load_gcs_to_bq
    )

load_csv_task

关键逻辑说明 #

1. 客户端初始化：使用google.cloud.bigquery.Client()直接初始化客户端，Airflow会自动读取环境内置的服务账号凭证（需确保账号拥有GCS读权限和BigQuery写权限）。
2. 参数映射：原Operator的所有配置参数都对应到Python函数中，保证功能完全匹配。
3. 加载配置：通过LoadJobConfig设置表结构、写入策略、文件格式等，和原Operator的配置逻辑完全对齐。
4. 任务执行：调用load_table_from_uri发起加载任务，并用result()阻塞等待任务完成，确保数据导入成功后再推进后续流程。

注意事项 #

• 确保Airflow工作节点已安装依赖库：pip install google-cloud-bigquery。
• 若CSV无表头，需移除skip_leading_rows=1配置。
• location参数必须与目标BigQuery表的区域保持一致，否则会触发区域不匹配错误。

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