如何使用Airflow本地BigQuery Operator执行任务及拉取BQ结果到Python变量

一、仅使用BigQuery Operator的可行性及实现 #

如果你的业务逻辑可完全封装为BigQuery支持的SQL/存储过程逻辑，不需要跨系统交互、非BQ侧的自定义计算，仅使用BigQuery Operator即可实现需求，无需引入Python Operator。

1. 前置配置 #

• GCP连接配置：在Airflow UI的Admin -> Connections页面新建GCP连接，连接类型选择Google Cloud，填入对应服务账号的密钥内容或文件路径，默认连接ID为google_cloud_default，可自定义后在Operator参数中指定。
• Airflow变量配置：在Admin -> Variables页面配置全局通用参数，避免硬编码，示例配置项：
  • key: bq_project_id，值：你的GCP项目ID
  • key: bq_dataset，值：目标数据集名称
  • key: bq_target_table，值：目标表名称
    代码中通过Variable.get()即可读取配置的变量。

2. 条件逻辑实现方案 #

BQ Operator支持执行标准SQL、BQ存储过程，所有条件分支、循环逻辑都可以封装在SQL或存储过程中，无需通过Python控制。示例DAG代码如下：

from airflow import DAG
from airflow.models import Variable
from airflow.providers.google.cloud.operators.bigquery import BigQueryExecuteQueryOperator
from datetime import datetime, timedelta

# 读取Airflow变量
BQ_PROJECT_ID = Variable.get("bq_project_id")
BQ_DATASET = Variable.get("bq_dataset")
BQ_TARGET_TABLE = Variable.get("bq_target_table")

default_args = {
    "owner": "airflow",
    "retries": 1,
    "retry_delay": timedelta(minutes=5)
}

with DAG(
    "bq_only_workflow",
    default_args=default_args,
    schedule_interval="0 2 * * *",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["bigquery"]
) as dag:
    run_bq_task = BigQueryExecuteQueryOperator(
        task_id="run_conditional_bq_logic",
        project_id=BQ_PROJECT_ID,
        sql=f"""
        -- 条件逻辑直接写在SQL中，示例：根据昨日销售额判断是否更新目标表
        DECLARE yesterday_sale_total INT64;
        SET yesterday_sale_total = (
            SELECT SUM(sale_amount) 
            FROM `{BQ_PROJECT_ID}.{BQ_DATASET}.sale_table` 
            WHERE dt = DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY)
        );
        
        IF yesterday_sale_total > 10000 THEN
          -- 满足条件执行更新逻辑
          MERGE `{BQ_PROJECT_ID}.{BQ_DATASET}.{BQ_TARGET_TABLE}` T
          USING (
            SELECT user_id, order_id, sale_amount, dt 
            FROM `{BQ_PROJECT_ID}.{BQ_DATASET}.sale_table` 
            WHERE dt = DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY)
          ) S
          ON T.user_id = S.user_id AND T.dt = S.dt
          WHEN MATCHED THEN UPDATE SET T.sale_amount = S.sale_amount
          WHEN NOT MATCHED THEN INSERT ROW;
        ELSE
          -- 不满足条件写入操作日志
          INSERT INTO `{BQ_PROJECT_ID}.{BQ_DATASET}.operate_log` (dt, operate_type, remark)
          VALUES (CURRENT_DATE(), 'skip_update', '昨日销售额不足1万，跳过目标表更新');
        END IF;
        """,
        use_legacy_sql=False,
        location="us-central1"
    )

如果逻辑更复杂，也可以把逻辑封装为BQ存储过程，Operator的sql参数直接写CALL 存储过程名(参数)即可调用。

二、结合BQ Operator与Python Operator的实现方案 #

如果你的逻辑需要调用第三方接口、生成非BQ格式的输出、和非GCP系统交互等BQ不支持的操作，仅用BQ Operator无法实现，可通过以下方案将BQ查询结果赋值给Python变量使用：

实现逻辑 #

1. 使用BigQueryGetDataOperator拉取BQ查询结果，结果会以二维列表的形式自动存入Airflow的XCom组件
2. 在Python Operator中通过ti.xcom_pull拉取对应任务的返回结果，赋值给Python变量后做自定义处理
   示例DAG代码如下：

from airflow import DAG
from airflow.models import Variable
from airflow.providers.google.cloud.operators.bigquery import BigQueryGetDataOperator
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

BQ_PROJECT_ID = Variable.get("bq_project_id")
BQ_DATASET = Variable.get("bq_dataset")

default_args = {
    "owner": "airflow",
    "retries": 1,
    "retry_delay": timedelta(minutes=5)
}

def process_bq_result(**context):
    # 拉取BQ Operator返回的结果
    bq_query_result = context["ti"].xcom_pull(task_ids="get_bq_sale_data")
    # 赋值给Python变量，示例取第一行第一列的销售额数据
    yesterday_sale_total = bq_query_result[0][0]
    # 以下可写自定义Python逻辑：调用第三方接口、生成报表、同步数据到其他系统等
    print(f"昨日销售额为{yesterday_sale_total}")

with DAG(
    "bq_python_combined_workflow",
    default_args=default_args,
    schedule_interval="0 2 * * *",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=["bigquery", "python"]
) as dag:
    # 拉取BQ数据，结果自动存入XCom
    get_bq_data = BigQueryGetDataOperator(
        task_id="get_bq_sale_data",
        project_id=BQ_PROJECT_ID,
        dataset_id=BQ_DATASET,
        table_id="sale_table",
        selected_fields="sale_amount",
        where_clause="dt = DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY)",
        max_results=1,
        location="us-central1"
    )

    # 处理BQ返回的数据
    process_data = PythonOperator(
        task_id="process_bq_data",
        python_callable=process_bq_result,
        provide_context=True
    )

    get_bq_data >> process_data

注意：如果拉取的BQ数据量较大，不要用XCom传递结果，建议先将BQ结果导出到GCS，再在Python Operator中读取GCS文件处理，避免XCom存储溢出。

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