simple_dag.py - Examples

  1"""
  2Airflow 기본 DAG 예제
  3
  4이 DAG는 기본적인 Airflow 워크플로우를 보여줍니다:
  5- PythonOperator로 Python 함수 실행
  6- BashOperator로 쉘 명령 실행
  7- Task 의존성 정의
  8
  9실행: airflow dags test simple_dag 2024-01-01
 10"""
 11
 12from datetime import datetime, timedelta
 13from airflow import DAG
 14from airflow.operators.python import PythonOperator
 15from airflow.operators.bash import BashOperator
 16from airflow.operators.empty import EmptyOperator
 17
 18
 19# 기본 인자 설정
 20default_args = {
 21    'owner': 'data_team',
 22    'depends_on_past': False,
 23    'email': ['data-alerts@example.com'],
 24    'email_on_failure': False,
 25    'email_on_retry': False,
 26    'retries': 1,
 27    'retry_delay': timedelta(minutes=5),
 28}
 29
 30
 31# Python 함수 정의
 32def print_hello():
 33    """인사 메시지 출력"""
 34    print("Hello from Airflow!")
 35    return "hello_returned"
 36
 37
 38def print_date(**context):
 39    """실행 날짜 출력"""
 40    execution_date = context['ds']
 41    print(f"Execution date: {execution_date}")
 42    return execution_date
 43
 44
 45def process_data(value: int, multiplier: int = 2, **context):
 46    """데이터 처리 예시"""
 47    result = value * multiplier
 48    print(f"Processing: {value} * {multiplier} = {result}")
 49
 50    # XCom으로 결과 저장
 51    context['ti'].xcom_push(key='processed_value', value=result)
 52    return result
 53
 54
 55def summarize(**context):
 56    """이전 Task 결과 요약"""
 57    ti = context['ti']
 58
 59    # XCom에서 값 가져오기
 60    hello_result = ti.xcom_pull(task_ids='hello_task')
 61    processed_value = ti.xcom_pull(task_ids='process_task', key='processed_value')
 62
 63    print(f"Summary:")
 64    print(f"  - Hello result: {hello_result}")
 65    print(f"  - Processed value: {processed_value}")
 66    print(f"  - Execution date: {context['ds']}")
 67
 68
 69# DAG 정의
 70with DAG(
 71    dag_id='simple_dag',
 72    default_args=default_args,
 73    description='간단한 Airflow DAG 예제',
 74    schedule_interval='@daily',  # 매일 실행
 75    start_date=datetime(2024, 1, 1),
 76    catchup=False,  # 과거 실행 건너뛰기
 77    tags=['example', 'tutorial'],
 78) as dag:
 79
 80    # Task 정의
 81    start = EmptyOperator(task_id='start')
 82
 83    hello_task = PythonOperator(
 84        task_id='hello_task',
 85        python_callable=print_hello,
 86    )
 87
 88    date_task = PythonOperator(
 89        task_id='date_task',
 90        python_callable=print_date,
 91    )
 92
 93    bash_task = BashOperator(
 94        task_id='bash_task',
 95        bash_command='echo "Current time: $(date)" && sleep 2',
 96    )
 97
 98    process_task = PythonOperator(
 99        task_id='process_task',
100        python_callable=process_data,
101        op_kwargs={'value': 10, 'multiplier': 5},
102    )
103
104    summary_task = PythonOperator(
105        task_id='summary_task',
106        python_callable=summarize,
107    )
108
109    end = EmptyOperator(task_id='end')
110
111    # Task 의존성 정의
112    #     ┌─ hello_task ─┐
113    # start ─┤             ├─ process_task ─ summary_task ─ end
114    #     └─ date_task ──┘
115    #             └─ bash_task ──┘
116
117    start >> [hello_task, date_task]
118    hello_task >> process_task
119    date_task >> [bash_task, process_task]
120    [bash_task, process_task] >> summary_task >> end
121
122
123if __name__ == "__main__":
124    # 로컬 테스트
125    dag.test()