DAG 디버깅하기 — dag.test()로 IDE에서, pdb로 커맨드라인에서

파이프라인을 만들다 보면 "이 DAG가 정말 의도대로 도는 걸까?"라는 순간이 꼭 찾아와요. Airflow는 dag.test() 메서드 하나로 DAG 전체를 단일 Python 프로세스에서 그대로 실행해보게 해줘서, IDE에서 편하게 디버깅할 수 있어요.

출처: Debugging Airflow Dags — Airflow Documentation

본문

dag.test()로 DAG 테스트하기

IDE에서 DAG를 디버깅하려면 DAG 파일에 dag.test 명령을 넣고, DAG 전체를 하나의 직렬 Python 프로세스로 돌려보면 돼요.

이 방식은 어떤 지원되는 데이터베이스에서도 쓸 수 있고(로컬 SQLite 포함), 모든 태스크가 단일 프로세스에서 돌기 때문에 빠르게 실패(fail fast) 하는 장점이 있어요.

dag.test를 준비하려면 DAG 파일 맨 아래에 두 줄만 추가하면 돼요:

if __name__ == "__main__":
    dag.test()

그게 전부예요! 선택적으로 인자를 추가해서 테스트를 다듬을 수도 있는데, 주로 쓰는 예는 다음과 같아요:

execution_date: 특정 실행 날짜의 DAG run을 테스트하고 싶을 때
use_executor: executor를 사용해 DAG를 테스트하고 싶을 때. 기본적으로 dag.test는 executor 없이 모든 태스크를 로컬에서 그냥 실행해요. 이 인자를 주면 Airflow 환경에 설정된 executor로 DAG를 실행해요.

조건부로 태스크 건너뛰기

로컬 환경에서 일부 태스크(예: 의존성 확인 센서나 정리 단계)를 실행하고 싶지 않다면, mark_success_pattern 인자에 그 태스크들의 task_id 패턴을 주면 자동으로 성공 처리해줘요.

다음 예시에서 DAG를 테스트할 때는 두 상위 DAG가 끝나길 기다리지 않아요. 대신 테스트 데이터를 직접 넣고, 정리 단계도 건너뛰어서 중간 csv를 살펴볼 수 있게 해요:

with DAG("example_dag", default_args=default_args) as dag:
    sensor = ExternalTaskSensor(
        task_id="wait_for_ingestion_dag",
        external_dag_id="ingest_raw_data",
    )
    sensor2 = ExternalTaskSensor(
        task_id="wait_for_dim_dag",
        external_dag_id="ingest_dim",
    )
    collect_stats = PythonOperator(
        task_id="extract_stats_csv",
        python_callable=extract_stats_csv,
    )
    # ... run other tasks
    cleanup = PythonOperator(
        task_id="cleanup",
        python_callable=Path.unlink,
        op_args=[collect_stats.output],
    )

    [sensor, sensor2] >> collect_stats >> cleanup


if __name__ == "__main__":
    ingest_testing_data()
    run = dag.test(mark_success_pattern="wait_for_.*|cleanup")
    print(f"Intermediate csv: {run.get_task_instance('collect_stats').xcom_pull(task_id='collect_stats')}")

커맨드라인에서 DAG 디버깅하기

위에서 추가한 두 줄만 있으면 pdb로도 DAG를 디버깅할 수 있어요. python -m pdb <path_to_dag_file>.py를 실행하면 커맨드라인에서 대화형 디버깅을 해볼 수 있어요:

[Breeze:3.10.19] root@ef2c84ad4856:/opt/airflow# python -m pdb providers/standard/src/airflow/providers/standard/example_dags/example_bash_operator.py
> /opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_bash_operator.py(18)<module>()
-> """Example Dag demonstrating the usage of the BashOperator."""
(Pdb) b 45
Breakpoint 1 at /opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_bash_operator.py:45
(Pdb) c
> /opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_bash_operator.py(45)<module>()
-> task_id = "run_after_loop",
(Pdb) run_this_last
<Task (EmptyOperator): run_this_last>

IDE 설정 단계

DAG 파일 끝에 main 블록을 추가해서 파일을 실행 가능하게 만들어요:
```
if __name__ == "__main__":
    dag.test()
```
이제 DAG 파일을 실행하거나 디버그해요.

DAG 디버깅하기 — dag.test()로 IDE에서, pdb로 커맨드라인에서