Why branching handles conditional logic in Apache Airflow - See It in Action

Why Branching Handles Conditional Logic in Airflow

📖 Scenario: You are building a simple Airflow workflow to decide which task to run based on a condition. This is like choosing a path in a choose-your-own-adventure story depending on a choice you make.

🎯 Goal: Create an Airflow DAG that uses branching to run different tasks based on a condition. You will first set up the tasks, then add a branching condition, and finally print which task runs.

📋 What You'll Learn

Create three tasks: start, branching, and two possible downstream tasks task_a and task_b.

Use BranchPythonOperator to decide which task to run based on a variable run_task_a.

Print the name of the task that runs after branching.

💡 Why This Matters

🌍 Real World

Branching in Airflow is used to run different tasks based on conditions like success/failure, data values, or external triggers.

💼 Career

Understanding branching helps you build flexible workflows that adapt to different situations, a key skill for workflow automation and data engineering roles.

1

Create initial Airflow tasks

Create an Airflow DAG named branching_dag with a start task using DummyOperator and two tasks task_a and task_b also using DummyOperator.

Apache Airflow

# Your code here to create DAG and three DummyOperator tasks: start, task_a, task_b

Need a hint?

Use DummyOperator to create simple placeholder tasks.

2

Add branching condition variable

Create a variable called run_task_a and set it to True. This variable will decide which task to run next.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

# Create a variable run_task_a and set it to True
# Your code here

Need a hint?

This variable will be used later to decide which task to run.

3

Add BranchPythonOperator with condition

Create a function called choose_task that returns 'task_a' if run_task_a is True, else returns 'task_b'. Then create a BranchPythonOperator named branching that uses this function. Set the task dependencies so start runs before branching, and branching runs before both task_a and task_b.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime
from airflow.operators.python import BranchPythonOperator

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

run_task_a = True

# Define choose_task function and BranchPythonOperator branching
# Set dependencies: start >> branching >> [task_a, task_b]
# Your code here

Need a hint?

Use BranchPythonOperator to decide which task runs next based on the function's return value.

4

Print which task runs after branching

Add a print statement that outputs the result of calling choose_task(). This shows which task will run after the branching decision.

Apache Airflow

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime
from airflow.operators.python import BranchPythonOperator

default_args = {'start_date': datetime(2024, 1, 1)}

dag = DAG('branching_dag', default_args=default_args, schedule_interval=None)

start = DummyOperator(task_id='start', dag=dag)
task_a = DummyOperator(task_id='task_a', dag=dag)
task_b = DummyOperator(task_id='task_b', dag=dag)

run_task_a = True

def choose_task():
    if run_task_a:
        return 'task_a'
    else:
        return 'task_b'

branching = BranchPythonOperator(
    task_id='branching',
    python_callable=choose_task,
    dag=dag
)

start >> branching >> [task_a, task_b]

# Print which task will run after branching
# Your code here

Need a hint?

Calling choose_task() returns the task id that will run next.