Airflow - Make http request for each row of Big Query Result Set

I would like to use Airflow DAG to execute a GCP Big Query, and for each row of the result set, I will make a request to an endpoint. With the return of this http request, I would like to store it in GCS, as external table.

I didn't find a way to pass the result set of GCP Big query to other operator and iterate over it. The only way that I find, is to use a Python Operator, but I supposed that there is a better way to do that.

Solution 1

I think for this this kind of need and use case, the most easy solution is to do the different operations in a PythonOperator as you mentioned.

In the PythonOperator, you can use Python BigQuery client to execute BigQuery job.

You can retrieve the result in a List of Dict, then launch the api requests in your endpoint, per element in the list (or maybe a more optimized way) and again use Python client to store the result in GCS as external table.

Solution 2

You can also mix usual operators with PythonOperator and use xcom to retrieve the result from the previous operator, example :

import airflow
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import PythonOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator


def call_api(**kwargs):
    ti = kwargs['ti']
    query_results = ti.xcom_pull(task_ids='your_query')

    # api call by query result


with airflow.DAG(
        "dag_id",
        default_args={},
        schedule_interval=None) as dag:
    your_query_task = BigQueryInsertJobOperator(
        task_id='your_query',
        configuration={
            "query": {
                "query": 'your_query',
                "useLegacySql": False,
            }
        },
        location='EU'
    )
    start_dag = DummyOperator(task_id='OK', dag=dag)

    api_call_task = PythonOperator(
        task_id="save_file_bq",
        op_kwargs={
            'dataset': 'dataset',
            'table': 'table'
        },
        python_callable=call_api
    )

    # The rest of DAG for upload the result in GCS as external table....

    start_dag >> api_call_task

For the rest of logic in the DAG, you can decide what is the best approach for you :

• Upload the result in GCS in the PythonOperator
• Or use an usual operator with xcom again but in this case, you should extend the built in operator