Apache Airflow Sample BranchPythonOperator DAG file

sample_branching_dag.py

from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.empty import EmptyOperator
from airflow.decorators import task
from airflow.decorators.branch_python import branch_task


@task(task_id='beginning')
def beginning() -> int:
    """airflow.decorators.python
    https://github.com/apache/airflow/blob/main/airflow/decorators/python.py
    """
    import random
    return random.randint(1, 3)


@branch_task(task_id='branching')
def branching(random_num: int) -> list:
    """airflow.decorators.branch_python
    https://github.com/apache/airflow/blob/main/airflow/decorators/branch_python.py
    """
    if random_num == 1:
        return ['task_1']
    elif random_num == 2:
        return ['task_2']
    elif random_num == 3:
        return ['task_3']


with DAG(
    dag_id='sample_branching',
    description='A Sample BranchPythonOperator DAG',
    schedule_interval=timedelta(days=1),
    start_date=datetime(2022, 6, 1),
    catchup=False,
    tags=['sample'],
) as dag:

    # airflow.operators.empty
    # https://github.com/apache/airflow/blob/main/airflow/operators/empty.py
    task_1 = EmptyOperator(task_id='task_1')
    task_2 = EmptyOperator(task_id='task_2')
    task_3 = EmptyOperator(task_id='task_3')

    # airflow.utils.trigger_rule
    # https://github.com/apache/airflow/blob/main/airflow/utils/trigger_rule.py
    finishing = EmptyOperator(
        task_id='finishing',
        trigger_rule='none_failed_min_one_success'
    )

    # setting up dependencies.
    branching(beginning()) >> [task_1, task_2, task_3] >> finishing