Airflow

세 가지 Cluster Policy

클러스터 전역에서 Dag나 Task를 검사·변형하고 싶다면 Cluster Policy를 쓴다. 대표적인 활용 사례는 Dag/태스크가 특정 표준을 지키는지 검사하기, Dag/태스크에 기본 인자를 설정하기, 커스텀 라우팅 로직 수행하기다.

``dag_policy``는 ``~airflow.models.dag.DAG`` 타입의 ``dag`` 파라미터를 받고, ``DagBag``이 Dag를 로드하는 시점에 실행된다. ``task_policy``는 ``~airflow.models.baseoperator.BaseOperator`` 타입의 ``task`` 파라미터를 받으며, ``DagBag``이 로드 시점에 태스크를 파싱하며 생성할 때 실행된다 — 이는 태스크 정의 전체를 이 정책 안에서 바꿀 수 있다는 뜻이고, 특정 DagRun에 속한 개별 태스크가 아니라 앞으로 실행될 모든 태스크 인스턴스에 적용된다. ``task_instance_mutation_hook``은 ``~airflow.models.taskinstance.TaskInstance`` 타입의 ``task_instance``와 선택적인 ``~airflow.models.dagrun.DagRun`` 타입의 ``dag_run``을 받는다 — 이는 태스크 자체가 아니라 특정 DagRun에 관련된 태스크의 인스턴스에 적용되며, Dag 파일 프로세서가 아니라 스케줄러 쪽에서 태스크 인스턴스가 생성·조정될 때 실행되고, 현재 실행 중인 그 run(인스턴스)에만 적용된다. ``dag_run`` 인자는 실행 설정(``dag_run.conf``)에 따라 정책을 라우팅할 수 있게 해주며, 초기 태스크 인스턴스 생성 단계에서는 ``None``일 수 있으므로 ``task_instance``만 선언한 훅도 문제없이 계속 동작한다. ``dag_run.conf``는 수동 트리거나 API로 트리거된 실행에만 채워지며, 스케줄된 실행은 빈 conf를 갖는다.

주의: ``task_instance_mutation_hook``은 커밋을 금지하는 스케줄러 트랜잭션 안에서 실행된다. 훅 안에서 새 DB 세션을 열거나 커밋해서는 안 되며, 특히 활성 세션을 넘기지 않고 ``task_instance.get_dagrun()``을 호출하면 안 된다 — 그 결과로 발생하는 커밋이 스케줄러를 크래시시킨다. 실행 설정을 읽어야 한다면 ``dag_run`` 인자를 대신 써야 한다.

정책의 우선순위와 예외 처리

Cluster Policy가 설정하는 추가 속성은 Dag 파일에 정의된 속성보다 우선한다 — 예를 들어 Dag 파일의 태스크에 ``sla``를 설정했더라도, Cluster Policy가 같은 태스크에 ``sla``를 다시 설정하면 Cluster Policy 쪽 값이 최종적으로 적용된다.

Dag·Task 정책은 ``~airflow.exceptions.AirflowClusterPolicyViolation`` 예외를 던져 전달받은 Dag/태스크가 규정을 지키지 않으므로 로드되면 안 된다는 것을 알릴 수 있다. 이 예외는 Airflow 웹 UI에 표시되고 메타데이터 DB의 ``import_error`` 테이블에도 기록된다. 반면 Dag를 의도적으로 스킵해야 할 때는 ``~airflow.exceptions.AirflowClusterPolicySkipDag`` 예외를 던지면 되는데, ``AirflowClusterPolicyViolation``과 달리 이 예외는 Airflow 웹 UI에 표시되지 않는다(내부적으로 메타데이터 DB의 ``import_error`` 테이블에도 기록되지 않는다).

실무적으로는, 여러 검사 항목을 적용해야 한다면 이 규칙들을 별도의 파이썬 모듈에 모아두고, 여러 커스텀 검사를 수행한 뒤 오류 메시지들을 집계해 단일 ``AirflowClusterPolicyViolation``으로 보고하는 하나의 정책/태스크 뮤테이션 훅을 두는 것이 권장된다.

두 가지 유의점도 있다. 먼저, Dag 정책은 Dag가 완전히 로드된 '이후'에 적용되므로 ``default_args`` 파라미터를 오버라이드해도 효과가 없다 — 기본 오퍼레이터 설정을 오버라이드하고 싶다면 대신 태스크 정책을 써야 한다. 또한, Cluster Policy 안에서 타입 어노테이션에 ``DAG``를 쓸 때는 임포트 순환을 피하기 위해 ``airflow``가 아니라 ``airflow.models``에서 임포트해야 한다.

정책 정의 방법 — 로컬 설정 파일 vs Pluggy entrypoint

정책 함수를 설정하는 방법은 두 가지다. 첫째, ``airflow_local_settings.py`` 파일을 파이썬 검색 경로 어딘가(``$AIRFLOW_HOME``의 ``config/`` 폴더가 기본적으로 좋은 위치)에 만들고, 위에서 다룬 정책 이름(예: ``dag_policy``)과 일치하는 콜러블들을 추가한다. 둘째(2.6에서 추가, 더 고급 방식), Pluggy 인터페이스를 이용한 setuptools entrypoint를 쓰는 방법이다 — 먼저 모듈에 ``@hookimpl`` 데코레이터가 붙은 정책 함수를 정의하고, ``pyproject.toml``의 ``[project.entry-points.'airflow.policy']``에 엔트리포인트를 등록한다. 엔트리포인트 그룹은 반드시 ``airflow.policy``여야 하고, 엔트리마다 이름이 고유해야 하며(그렇지 않으면 중복 항목은 Pluggy에 의해 무시된다), 값은 ``@hookimpl``로 표시된 모듈(또는 클래스)이어야 한다. 배포판을 Airflow 환경에 설치하고 나면 다양한 Airflow 컴포넌트가 정책 함수를 호출한다(여러 플러그인이 있을 때 정확한 호출 순서는 정해져 있지 않으므로 특정 순서에 의존하면 안 된다).

어느 방식을 쓰든 인자 이름은 문서화된 이름과 정확히 일치해야 한다.

예시로, 모든 Dag에 태그가 하나 이상 있는지 검사하는 Dag 정책, 모든 태스크에 최대 타임아웃을 강제하는 태스크 정책, Airflow owner가 없는 태스크 실행을 막는 정책, 그리고 재시도 2회 이상인 태스크를 다른 큐로 재라우팅하는 task_instance_mutation_hook 예시가 공식 테스트 코드에 포함되어 있다. 다만 마지막 예시와 관련해, priority_weight는 weight rule로 동적으로 결정되므로 뮤테이션 훅 안에서 태스크 인스턴스의 ``priority_weight``를 바꿀 수는 없다는 점에 유의해야 한다.

Metadata Engine Hooks — DB 엔진 생성 오버라이드

Cluster Policy와는 별개로, ``airflow_local_settings.py``는 Airflow가 메타데이터 DB 엔진을 만드는 방식 자체도 오버라이드할 수 있게 해준다. 이는 정적 설정만으로는 표현할 수 없는 커넥션별 로직이 필요할 때 유용하다 — 예를 들어 SQLAlchemy의 ``do_connect`` 이벤트 핸들러를 통해 단명(short-lived) JWT 토큰이나 IAM 자격증명을 주입하는 경우다.

오버라이드할 수 있는 함수는 두 가지다. ``create_metadata_engine(sql_alchemy_conn, *, engine_args, connect_args) -> Engine``은 ``configure_orm()``이 동기 메타데이터 엔진을 만들 때 호출한다. ``create_async_metadata_engine(sql_alchemy_conn_async, *, connect_args) -> AsyncEngine``은 ``_configure_async_session()``이 비동기 메타데이터 엔진을 만들 때 호출한다. 기본 구현은 Airflow가 항상 써온 것과 같은 인자로 ``sqlalchemy.create_engine`` / ``sqlalchemy.ext.asyncio.create_async_engine``을 호출하므로, 오버라이드를 제공하지 않는 한 동작 변화는 전혀 없다.

예를 들어 물리적 커넥션이 새로 생길 때마다(풀 재활용 이후 포함) JWT 토큰을 갱신하는 ``do_connect`` 핸들러를 등록하려면, ``airflow_local_settings.py``에서 ``create_metadata_engine``을 오버라이드해 ``sqlalchemy.create_engine``으로 엔진을 만든 뒤 ``event.listen(engine, "do_connect", _refresh_jwt)``로 핸들러를 등록하고 그 엔진을 반환하면 된다.