Debugging Airflow in a Containered with VS Code

#Dag 테스트
airflow dags test DAG_NAME YYYY-MM-DD
#Task 테스트
airflow tasks test DAG_NAME TASK_NAME YYYY-MM-DD

Debug interactively with dag.test()

디버깅 환경 구축이 필요한 이유

DAG에 코드를 작성하고 디버깅하는 과정에서 매번 schedule로 실행하여 확인할 수는 없다. Airflow scheduler에서 실행하지 않고 DAG를 vs code와 같은 에디터에서 빠르게 디버깅하고 테스트할 수 있는 방법을 알아보자.

- Airflow 2.8.0 in docker container

-  IDE : VS Code

 

Dag 파일에 test 설정

아래와 같이 dag.test()를 추가해준다.

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

 

example_dag.py 

from datetime import datetime, timedelta
from textwrap import dedent
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.postgres_operator import PostgresOperator
from plugins import slack

DAG_ID = "postgres_operator_dag" 

default_args = {
    'owner': 'airflow',
    'schedule_interval': '@hourly',
    'start_date': datetime(2023, 9, 28),
    'tags': ['shop'],
    # 'retries': 1,
    'on_failure_callback': slack.on_failure_callback, # 실패시 SLACK 함수 요청
    'on_success_callback': slack.on_failure_callback, # 실패시 SLACK 함수 요청
}

with DAG(
    dag_id=DAG_ID,
    default_args=default_args, 
    catchup=False,
    schedule_interval = '50 22 * * *', 
) as dag:
    get_products = PostgresOperator(
        task_id="get_products",
        sql="SELECT * FROM producta",
    )


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

 

Airflow debugger 환경 구축

VS code Remote Development extension 설치 및 세팅

VS Code에서 Docker container에 쉽게 접속하기 위해서 아래 extension 설치가 필요하다.

 

airflow 프로젝트 폴더에 .devcontainer 폴더를 생성한다.

mkdir .devcontainer 

// airflow에 필요한 폴더 생성하거나 기존에 있었던 것을 .devcontainer 하위로 옮긴다.
// docker-compose.yml로 세팅한 경우 필요하다면 경로에 맞게 위치시킨다.
cd .devcontainer
mkidr dags
mkdir logs
mkdir plugins

 

container 자동으로 접속하기 위해 필요한 설정 파일을 작성하여 추가한다.

 

.devcontainer/devcontainer.json

{
  "name": "airflow",
  "dockerComposeFile": "docker-compose.yml",
  "service": "airflow-worker",
  "workspaceFolder": "/opt/airflow",
  "customizations": {
  "extensions": ["ms-python.python"],
  "settings": {
    "terminal.integrated.shell.linux": "/bin/bash"
  }}
}

 

Airflow Docker container에 접속

command + shift + p > Dev Containers:Attach to Running Container..

 

🔥 Airflow Dag 파일 테스트 실행

YYYY-MM-DD : 실행날짜 (logical_date)

#Dag 테스트
airflow dags test DAG_NAME YYYY-MM-DD
#Task 테스트
airflow tasks test DAG_NAME TASK_NAME YYYY-MM-DD

 

아래와 같이 로그와 함께 실행 결과를 디버깅할 수 있다.

 

References

https://docs.astronomer.io/learn/testing-airflow

https://davidgriffiths-data.medium.com/debugging-airflow-in-a-container-with-vs-code-7cc26734444