Airflow Get value from table in Bigquery and use the return value to pass into next task

user_16557673 0
Viewed 540

I am trying to Take data from BigQuery Dataset and pass the result value to bash_command so that it will execute commands to remove files in Cloud storage.

When I execute 'SELECT commandshell_variable_1 FROM sys_process.outbound_flat_file_config where file_number=1' the result is ..... gsutil rm -r gs://A/loop_member_*.csv

I want to use the result of below and Pass it to bash_command in next task ... Thank you.

DAG Code

with DAG(
    dag_id='kodz_Automation',
    description='kodz_Automation',
    schedule_interval=None,
    catchup= False,
    default_args=DEFAULT_DAG_ARGS) as dag:

def get_data_from_bq(**kwargs):
    hook = BigQueryHook(bigquery_conn_id='bigquery_default', delegate_to=None, use_legacy_sql=False)
    conn = hook.get_conn()
    cursor = conn.cursor()
    cursor.execute('SELECT commandshell_variable_1 FROM sys_process.outbound_flat_file_config where  file_number=1')
    result = cursor.fetchall()
    print('result', result)
    return result

fetch_data = PythonOperator(
    task_id='fetch_data_public_dataset',
    provide_context=True,
    python_callable=get_data_from_bq,
    dag=dag
    )
also_run_this = bash_operator.BashOperator(
    task_id='also_run_this',
    bash_command=python_callable
)

fetch_data >> also_run_this

airflow
1 Answers

To send data from one task to another you can use Airflow XCOM feature.

Using PythonOperator, the returned value will be stored in XCOM by default, so all you need to do is add a xcom_pull in the BashOperator, something like this:

also_run_this = bash_operator.BashOperator(
   task_id='also_run_this',
   bash_command="<you command> {{ ti.xcom_pull(task_ids=[\'fetch_data_public_dataset\']) }}'"
)

Learn more of XCOM here

But if you will return a lot of data, I recommend saving this in some storage (like S3, GCS, etc) and then sending the link address to the bash command.

user_5895900
Related