Conditionally pick an Airflow DAG branch using an SQL query

This article is a part of my "100 data engineering tutorials in 100 days" challenge. (33/100)

In addition to the BranchPythonOperator, which lets us execute a Python function that returns the ids of the subsequent tasks that should run, we can also use a SQL query to choose a branch. Of course, we will not do it by querying the SQL database in the Python function. There is a shorter way.

We need to add a BranchSQLOperator to our DAG. This operator is a little bit different than the BranchPythonOperator. In the case of the Python operator, the function returns the ids of the tasks to run. The SQL version of the operator expects a boolean value in the first column of the first row. Optionally, it can also return a numeric. Every non-zero value is interpreted as True. Zero = False. We can also use one of the string values with a boolean equivalent (take a look at the documentation if you want to use a string value).

If the value is true, Airflow will execute the tasks specified in the follow_task_ids_if_true parameter. Otherwise, it runs the tasks defined in the follow_task_ids_if_false argument.

1
2
3
4
5
6
7
8
operator = BranchSQLOperator(
    task_id="some_task_id",
    conn_id="sql_connection_id",
    sql="SELECT count(1) FROM a_table",
    follow_task_ids_if_true="task_to_execute_if_count_at_least_one",
    follow_task_ids_if_false="tasks_to_execute_if_counts_is_zero",
    dag=dag
)

Subscribe to the newsletter and join the free email course.


Remember to share on social media!
If you like this text, please share it on Facebook/Twitter/LinkedIn/Reddit or other social media.

If you want to contact me, send me a message on LinkedIn or Twitter.

Would you like to have a call and talk? Please schedule a meeting using this link.


Bartosz Mikulski
Bartosz Mikulski * data/machine learning engineer * conference speaker * co-founder of Software Craft Poznan & Poznan Scala User Group

Subscribe to the newsletter and get access to my free email course on building trustworthy data pipelines.

Do you want to work with me at riskmethods?

REMOTE position (available in Poland or Germany)