Airflow XCOM 拉取和推送 BigQueryInsertJobOperator 和 BigQueryOperator

Question

I am very new to airflow and I am trying to create a DAG based on the below requirement.

• Task 1 - Run a Bigquery query to get a value which I need to push to 2nd task in the dag
• Task 2 - Use the value from the above query and run another query and export the data into google cloud bucket.

I have read other answers related to this and I understand we cannot use xcom_pull or xcom_push in bigqueryoperator in airflow. So what I am doing is using a python operator where I can use jinja template variables by using "provide_context=True".

Below is the snipped of my code. Just the task 1 where I want to do "task_instance.xcom_push" in order to see the value in airflow under logs xcom.

def get_bq_operator(dag, task_id, configuration, table_params=None, trigger_rule='all_success'):
    bq_operator = BigQueryInsertJobOperator(
        task_id=task_id,
        configuration=configuration,
        gcp_conn_id=gcp_connection_id,
        dag=dag,
        params=table_params,
        trigger_rule=trigger_rule,
        **task_instance.xcom_push(key='yr_wk', value=yr_wk),**
    )
    return bq_operator


def get_bq_wm_yr_wk():
    get_bq_operator(dag,app_name,bigquery_util.get_bq_job_configuration(
                                             bq_query,
                                             query_params=None))

get_wm_yr_wk = PythonOperator(task_id='get_wm_yr_wk',
                                        python_callable=get_bq_wm_yr_wk,
                                        provide_context=True,
                                        on_failure_callback=failure_callback,
                                        on_retry_callback=failure_callback,
                                        dag=dag)

"bq_query" is the one I am passing the sql file which has my query and the query returns the value of yr_wk which I need to use in my 2nd task.

The highlighted task_instance.xcom_push(key='yr_wk', value=yr_wk), in get_bq_operator is failing and the errror i am getting is as below

raise KeyError(f'Variable {key} does not exist')

KeyError: 'Variable ei_migration_hour does not exist'

If I comment the line above, the DAG runs fine. However, how do I validate the value of yr_wk?? I want to push it so that I can view the value in logs.

Answer 1

I do not fully understand your code:), but if you want to do something with results of BigQuery query, then by far better way to approach it is to use BigQueryHook in your python callable.

Operators in Airflow are usually thin wrappers around Hooks that really provide a "complete" taks (for example you can use it run an update operation) but if you want to do something with the result of it and you already do it via Python Operator, it is far better to use Hooks directly as you do not make all the assumptions that operators have in execute method.

In your case it should be something like (and I am using here the new TaskFlow syntax which is preferred to do this kind of operations. See https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html for the tutorial on Task Flow API. Aspecially in Airflow 2 it became the de-facto default way of writing tasks.

@task(.....) 
def my_task():
   hook = BigQueryHook(....)  # initialize it with the right parameters
   result = hook.run(sql='YOUR_QUERY', ...)  # add other necessary params
   processed_result = process_result(result) # do something with the result
   return processed_result

This way you do not evey have to run xcom_push (task_flow API will do it for you automatically and other tasks will be able to use by just doing:

@task
next_task(input):
   pass

And then:

result = my_task()
next_task(result)

Then all the xcom push/pull will be handled for you automatically via TaskFlow.