[Airflow] Python Operator로 외부 모듈을 Import하는 방법

1. Python Operator와 모듈 경로 이해하기

Airflow에서 PythonOperator를 사용해 작업을 정의할 때, 외부 함수가 필요하면 해당 모듈을 DAG 내에서 import할 수 있습니다. 이때 import 경로를 정확히 지정하는 것이 중요합니다.

Python 경로 설정 기본 원칙

1. 실행 중인 Python 파일과 동일한 디렉토리에 있는 모듈은 자동으로 import됩니다.
2. pip로 설치한 라이브러리들은 자동으로 sys.path에 포함됩니다.
3. sys.path 변수에 경로를 추가해, Python이 모듈을 찾을 수 있도록 설정할 수 있습니다.

Airflow의 기본 sys.path 구성

Airflow는 dags와 plugins 폴더를 자동으로 sys.path에 포함합니다. 따라서 이 두 폴더에 외부 모듈을 배치하면 별도의 설정 없이도 DAG 파일에서 이를 import할 수 있습니다.

경로를 수동으로 추가하기

경로 설정이 필요한 경우, sys.path.append('경로')로 명시적으로 경로를 추가하거나, 환경 변수 PYTHONPATH에 경로를 지정해 Airflow에서 모듈을 인식하도록 할 수 있습니다.

 

.env 파일 설정

WORKSPACE_FOLDER=~/vscode/airflow
PYTHONPATH=${WORKSPACE_FOLDER}/plugins

2. 공통 함수 작성 및 plugins 폴더 활용

모듈의 재활용성을 높이기 위해 공통 함수는 plugins 폴더에 작성하는 것이 좋습니다. 이는 DAG 파일이 더욱 깔끔하고 관리하기 쉬워지도록 해줍니다.

공통 함수와 DAG에서의 import

1. 공통 함수 파일 작성 (plugins/utilities.py):

def greet_user(name):
	print(f"Hello, {name}! Welcome to Airflow")

 

2. DAG 파일에서 함수 import 및 사용하기 (dags/greet_dag.py)

from airflow import DAG
from airflow.operators.python import PythonOperator
# plugins 폴더에 있는 utilities.py함수에 정의된 greet_user 함수를 import
from utilities import greet_user  
import pendulum

def run_greeting():
    greet_user("Airflow User")

with DAG(
    dag_id="greet_user_dag",
    schedule="0 9 * * *",
    start_date=pendulum.datetime(2024, 3, 1, tz="Asia/Seoul"),
    catchup=False,
) as dag:
    greeting_task = PythonOperator(
        task_id="greeting_task",
        python_callable=run_greeting
    )

 

• greet_user_dag라는 DAG은 매일 오전 9시에 실행됩니다.
• greeting_task라는 태스크는 run_greeting 함수를 실행하고, 이 함수는 greet_user("Airflow User")를 호출합니다.

이렇게 하면, plugins 폴더에 위치한 utilities.py 내부의 함수를 재사용할 수 있게 되어, 여러 DAG에서 같은 함수를 import하여 코드의 재활용성과 유지보수성을 높일 수 있습니다.