[Airflow] @Task.branch로 분기처리하기

1. Task.branch 사용 예제 코드

• 함수에 @task.branch로 래핑해서 함수 자체만으로 task를 사용할 수 있음

'''
dags_python_with_branch_decorator.py
'''


from airflow import DAG
from datetime import datetime
from airflow.operators.python import PythonOperator
from airflow.decorators import task


with DAG(
    dag_id='dags_python_with_branch_decorator',
    start_date=datetime(2023,4,1),
    schedule=None,
    catchup=False
) as dag:
    @task.branch(task_id='python_branch_task')
    def select_random():
        import random
        
        item_lst = ['A', 'B', 'C']
        selected_item = random.choice(item_lst)
        if selected_item == 'A':
            return 'task_a' # task 1개만 실행시키고 싶은 경우
        elif selected_item in ['B', 'C']:
            return ['task_b', 'task_c'] # task 2개 이상 실행시키고 싶은 경우

    def common_func(**kwargs):
        print(kwargs['selected'])

    task_a = PythonOperator(
        task_id = 'task_a',
        python_callable = common_func,
        op_kwargs = {'selected':'A'}
    )

    task_b = PythonOperator(
        task_id = 'task_b',
        python_callable = common_func,
        op_kwargs = {'selected':'B'}
    )

    task_c = PythonOperator(
        task_id = 'task_c',
        python_callable = common_func,
        op_kwargs = {'selected':'C'}
    )

    select_random() >> [task_a, task_b, task_c]

 

 

실행결과 Graph

 

 

실행결과 Log

 

실행결과 Xcom