[Airflow] BaseBranchOperator 로 분기처리하기

1. BaseBranchOperator 사용 예제코드

'''
dags_base_branch_operator.py
'''


from airflow import DAG
import pendulum
from airflow.operators.branch import BaseBranchOperator
from airflow.operators.python import PythonOperator


with DAG(
    dag_id='dags_base_branch_operator',
    start_date=pendulum.datetime(2023,4,1, tz='Asia/Seoul'),
    schedule=None,
    catchup=False
) as dag:
	# 클래스 상속부분
    # BaseBranchOperator 상속시 choose_branch 함수를 구현해줘야함
    class CustomBranchOperator(BaseBranchOperator):
        def choose_branch(self, context):
            import random
            print(context)
            
            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개 이상 실행시키고 싶은 경우

    custom_branch_operator = CustomBranchOperator(task_id='python_branch_task')

    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'}
    )

    custom_branch_operator >> [task_a, task_b, task_c]

 

실행결과 Graph

 

실행결과 XCom

실행결과 Log