[Airflow]TaskFlow API의 XCom 생성 및 DB 저장 흐름

Overview#

01. TaskFlow API에서 XCom 생성 및 DB 저장 흐름#

사전개념#

TaskFlow API는 함수의 리턴값이 자동으로 XCom에 저장됨
일반 Airflow에서는 명시적 xcom_push()가 필요하지만, TaskFlow API에서는 자동화됨

1
# 일반 Airflow 접근
2
def traditional_task(**context):
3
    value = get_data()
4
    context['ti'].xcom_push(key='my_data', value=value)
5

6
# TaskFlow API 접근
7
@task
8
def taskflow_task():
9
    value = get_data()
10
    return value

실행절차#

1. Decorator처리

@task 데코레이터가 함수를 감싸 특수 오퍼레이터로 변환

1
def task(
2
    python_callable=None,
3
    multiple_outputs=False,
4
    **kwargs
5
):
6
    def wrapper(f):
7
        # 태스크 함수를 특수 오퍼레이터로 래핑
8
        return _PythonDecoratedOperator(
9
            python_callable=f,
10
            multiple_outputs=multiple_outputs,
11
            **kwargs
12
        )
13

14
    if python_callable:
15
        return wrapper(python_callable)
16
    else:
17
        return wrapper

2. Operator 변환

데코레이터된 함수는 _PythonDecoratedOperator로 변환됨

1
# _PythonDecoratedOperator 내부 (간략화)
2
class _PythonDecoratedOperator(PythonOperator):
3
    def execute(self, context):
4
        # 부모 클래스의 execute 호출
5
        return super().execute(context)

3. 함수 실행 및 결과 캡처

PythonOperator의 execute() 메서드가 함수를 실행하고 결과를 캡처

1
def execute(self, context):
2
    # run collable
3
    result = self.execute_callable(context)
4

5
    # XCom 푸시 설정이 켜져 있으면 결과 반환 (자동 XCom 저장용)
6
    if self.do_xcom_push:
7
        return result
8

9
    return None
10

11
def execute_callable(self, context):
12
    # 사용자 함수 실행
13
    result = self.python_callable(*self.op_args, **self.op_kwargs)
14

15
    # 결과가 None이 아니면 XCom 푸시 활성화
16
    if result is not None:
17
        self.do_xcom_push = True
18
        return result
19
    return None

4. Xcom생성

execute()에서 반환된 값은 태스크 인스턴스의 xcom_push()를 통해 XCom으로 저장
Airflow 내부적으로 이 과정이 자동 진행됨

1
# 내부 동작 (직접 호출할 필요 없음)
2
context['ti'].xcom_push(key='return_value', value=result)

5. TaskInstance에서 XCom으로 변환

TaskInstance.xcom_push()는 XCom 모델의 set() 메서드를 호출
airflow/models/taskinstance.py에 정의됨

1
# TaskInstance 클래스 내부 (간략화)
2
@provide_session
3
def xcom_push(
4
    self,
5
    key: str,
6
    value: Any,
7
    session: Session = NEW_SESSION,
8
) -> None:
9
    # XCom 모델의 set() 메서드 호출
10
    XCom.set(
11
        key=key,
12
        value=value,
13
        task_id=self.task_id,
14
        dag_id=self.dag_id,
15
        execution_date=self.execution_date,
16
        session=session,
17
    )

6. 메타데이터 데이터베이스 저장 단계

XCom.set() 메서드가 실제로 값을 데이터베이스에 저장
airflow/models/xcom.py에 정의됨

1
# XCom 클래스
2
@classmethod
3
@provide_session
4
def set(
5
    cls,
6
    key: str,
7
    value: Any,
8
    task_id: str,
9
    dag_id: str,
10
    execution_date: datetime,
11
    session: Session = NEW_SESSION,
12
) -> None:
13
    # XCom 객체 생성 및 병합
14
    session.merge(
15
        XCom(
16
            key=key,
17
            value=value,
18
            task_id=task_id,
19
            dag_id=dag_id,
20
            execution_date=execution_date,
21
        )
22
    )
23
    # 변경사항 커밋
24
    session.commit()

7 데이터베이스 구조

XCom 값은 Airflow 메타데이터 데이터베이스의 xcom 테이블에 저장
일반적인 테이블 구조:
- dag_id: DAG 식별자
- task_id: 태스크 식별자
- execution_date: 실행 날짜/시간
- key: XCom 키 (TaskFlow에서는 보통 ‘return_value’)
- value: 직렬화된 반환값
- timestamp: 저장 시간