커스텀 Operator 만들기
커스텀 Operator 만들기 (Creating a custom Operator)
자신이나 팀의 요구에 맞는 새 Operator를 만드는 방법을 설명하는 문서예요. 공개 SDK 기본 클래스 BaseOperator를 확장해 커스텀 Operator를 만들고, Hooks로 외부 서비스와 통신하고, UI 표시·Jinja 템플릿·Sensors를 활용하고, 테스트하는 방법까지 살펴볼게요.
출처: 문서
본문
Airflow는 자신이나 팀의 요구에 맞는 새 Operator를 만들 수 있게 해요. 이 확장성은 Apache Airflow를 강력하게 만드는 많은 기능 중 하나예요.
공개 SDK 기본 클래스 BaseOperator를 확장해 원하는 어떤 Operator든 만들 수 있어요.
파생 클래스에서 오버라이드해야 할 두 메서드가 있어요:
- 생성자 (Constructor) — Operator에 필요한 파라미터를 정의해요. 자신의 Operator에 특화된 인자만 지정하면 돼요.
default_args는 Dag 파일에서 지정할 수 있어요. 자세한 내용은 Default args를 참고하세요. - Execute — 러너가 operator를 호출할 때 실행할 코드예요. 이 메서드는 구성 값을 읽는 데 사용할 수 있는 Airflow context를 파라미터로 포함해요.
참고 (Note)
커스텀 operator를 구현할 때
__init__메서드에서 비용이 큰 작업을 하지 마세요. Operator는 그것을 사용하는 태스크마다 스케줄러 사이클마다 한 번씩 인스턴스화되며, 데이터베이스 호출은 스케줄링을 상당히 느리게 하고 리소스를 낭비할 수 있어요.
새 파일 hello_operator.py에서 예시 HelloOperator를 구현해 봐요:
from airflow.sdk import BaseOperator
class HelloOperator(BaseOperator):
def __init__(self, name: str, **kwargs) -> None:
super().__init__(**kwargs)
self.name = name
def execute(self, context):
message = f"Hello {self.name}"
print(message)
return message
참고 (Note)
import가 동작하려면 파일을 PYTHONPATH env에 있는 디렉토리에 놓아야 해요. Airflow는 기본적으로 Airflow home의
dags/,plugins/,config/디렉토리를 PYTHONPATH에 추가해요. 예를 들어 우리 예시에서는 파일을custom_operator/디렉토리에 놓았어요. Python과 Airflow가 모듈을 관리하는 방법에 대한 자세한 내용은 Modules Management를 참고하세요.
파생된 커스텀 operator를 이제 이렇게 사용할 수 있어요:
from custom_operator.hello_operator import HelloOperator
with dag:
hello_task = HelloOperator(task_id="sample-task", name="foo_bar")
커스텀 operator를 저장하기 위해 plugins 폴더를 계속 사용할 수도 있어요. hello_operator.py 파일이 plugins 폴더 안에 있다면 operator를 이렇게 import할 수 있어요:
from hello_operator import HelloOperator
Operator가 외부 서비스(API, 데이터베이스 등)와 통신한다면 Hooks를 사용해 통신 계층을 구현하는 것이 좋아요. 이렇게 하면 구현된 로직을 다른 사용자가 다른 operator에서 재사용할 수 있어요. 이런 접근 방식은 외부 서비스마다 CustomServiceBaseOperator를 사용하는 것보다 더 나은 결합 분리(decooupling)와 통합 활용을 제공해요.
또 다른 고려 사항은 임시 상태예요. 작업에 인-메모리 상태(예: on_kill 메서드에서 요청을 취소하는 데 사용해야 하는 job id)가 필요하다면 그 상태는 hook이 아니라 operator에 유지해야 해요. 이렇게 하면 서비스 hook을 완전히 state-less하게 유지하고 작업의 전체 로직을 한 곳(operator)에 둘 수 있어요.
Hooks
Hooks는 Dag에서 외부 공유 리소스와 통신하는 인터페이스 역할을 해요. 예를 들어 Dag의 여러 태스크가 MySQL 데이터베이스에 접근해야 할 수 있어요. 태스크마다 connection을 만드는 대신 hook에서 connection을 가져와 활용할 수 있어요. Hook은 또한 Dag에 연결 인증 파라미터를 저장하지 않아도 되게 해줘요. connections를 만들고 관리하는 방법은 Managing Connections을, providers를 통해 커스텀 connection 타입을 추가하는 방법의 세부 사항은 Providers를 참고하세요.
MySQL에서 이름을 가져오도록 이전 예시를 확장해 봐요:
class HelloDBOperator(BaseOperator):
def __init__(self, name: str, mysql_conn_id: str, database: str, **kwargs) -> None:
super().__init__(**kwargs)
self.name = name
self.mysql_conn_id = mysql_conn_id
self.database = database
def execute(self, context):
hook = MySqlHook(mysql_conn_id=self.mysql_conn_id, schema=self.database)
sql = "select name from user"
result = hook.get_first(sql)
message = f"Hello {result['name']}"
print(message)
return message
Operator가 hook 객체에서 쿼리를 호출하면, connection이 없으면 새 connection이 생성돼요. hook은 username과 password 같은 인증 파라미터를 Airflow backend에서 가져와 airflow.hooks.base.BaseHook.get_connection()에 전달해요. Hook은 execute 메서드나 execute에서 호출되는 어떤 메서드에서만 생성해야 해요. 생성자는 Airflow가 Dag를 파싱할 때마다(자주 발생함) 호출되고 거기서 hook을 인스턴스화하면 불필요한 데이터베이스 연결이 많이 생겨요. execute는 Dag 실행 중에만 호출돼요.
사용자 인터페이스 (User interface)
Airflow는 개발자가 operator가 Dag UI에 어떻게 표시되는지 제어할 수도 있게 해줘요. ui_color를 오버라이드해 UI에서 operator의 배경색을 바꿔요. ui_fgcolor를 오버라이드해 레이블의 색을 바꿔요. custom_operator_name을 오버라이드해 표시 이름을 클래스명이 아닌 다른 것으로 바꿔요.
class HelloOperator(BaseOperator):
ui_color = "#ff0000"
ui_fgcolor = "#000000"
custom_operator_name = "Howdy"
# ...
템플릿 (Templating)
operator를 매개변수화하기 위해 Jinja templates를 사용할 수 있어요. Airflow는 operator를 렌더링할 때 template_fields에 있는 필드 이름을 템플릿에 고려해요.
class HelloOperator(BaseOperator):
template_fields: Sequence[str] = ("name",)
def __init__(self, name: str, world: str, **kwargs) -> None:
super().__init__(**kwargs)
self.name = name
self.world = world
def execute(self, context):
message = f"Hello {self.world} it's {self.name}!"
print(message)
return message
템플릿은 이렇게 사용할 수 있어요:
with dag:
hello_task = HelloOperator(
task_id="task_id_1",
name="{{ task_instance.task_id }}",
world="Earth",
)
이 예시에서 Jinja는 name 파라미터를 찾아 {{ task_instance.task_id }}를 task_id_1로 치환해요.
파라미터는 bash 스크립트나 SQL 파일 같은 파일 이름을 포함할 수도 있어요. template_ext에 파일의 확장자를 추가해야 해요. template_field가 template_ext에 언급된 확장자로 끝나는 문자열을 포함하면, Jinja는 파일의 내용을 읽고 템플릿을 실제 값으로 대체해요. Jinja는 args가 아니라 operator 속성을 치환한다는 점을 참고하세요.
class HelloOperator(BaseOperator):
template_fields: Sequence[str] = ("guest_name",)
template_ext = ".sql"
def __init__(self, name: str, **kwargs) -> None:
super().__init__(**kwargs)
self.guest_name = name
예시에서 template_fields는 ['name']이 아니라 ['guest_name']이어야 해요.
또한 template_fields_renderers에 웹 UI에서 템플릿 필드의 값을 어떤 스타일로 렌더링할지 정의하는 딕셔너리를 제공할 수 있어요. 예를 들어:
class MyRequestOperator(BaseOperator):
template_fields: Sequence[str] = ("request_body",)
template_fields_renderers = {"request_body": "json"}
def __init__(self, request_body: str, **kwargs) -> None:
super().__init__(**kwargs)
self.request_body = request_body
template_field 자체가 딕셔너리인 상황에서는 점으로 구분된 키 경로를 지정해 개별 요소를 적절히 추출·렌더링할 수도 있어요. 예를 들어:
class MyConfigOperator(BaseOperator):
template_fields: Sequence[str] = ("configuration",)
template_fields_renderers = {
"configuration": "json",
"configuration.query.sql": "sql",
}
def __init__(self, configuration: dict, **kwargs) -> None:
super().__init__(**kwargs)
self.configuration = configuration
그런 다음 이 템플릿을 이렇게 사용해요:
with dag:
config_task = MyConfigOperator(
task_id="task_id_1",
configuration={"query": {"job_id": "123", "sql": "select * from my_table"}},
)
이것은 UI가 configuration을 json으로 렌더링하고, 추가로 query.sql의 configuration에 포함된 값을 SQL lexer로 렌더링하게 해요.
현재 사용할 수 있는 lexer:
bash bash_command doc doc_json doc_md doc_rst doc_yaml doc_md hql html jinja json md mysql postgresql powershell py python_callable rst sql tsql yaml
존재하지 않는 lexer를 사용하면 템플릿 필드의 값이 pretty-printed 객체로 렌더링돼요.
제한 사항 (Limitations)
오용을 방지하기 위해, operator의 생성자에서 템플릿 필드를 정의하고 할당할 때 다음 제한 사항을 준수해야 해요:
- 생성자에 전달되는 템플릿 필드에 해당하는 파라미터는 필드와 정확히 같은 이름이어야 해요. 다음 예시는 생성자에 전달된 파라미터가 템플릿 필드와 같지 않아 유효하지 않아요:
class HelloOperator(BaseOperator):
template_fields = "foo"
def __init__(self, foo_id) -> None: # should be def __init__(self, foo) -> None
self.foo = foo_id # should be self.foo = foo
- 템플릿 필드의 인스턴스 멤버는 생성자의 해당 파라미터로 할당되어야 해요. 직접 할당하거나, (그 필드들이
template_fields로 정의된) 부모 생성자를 파라미터의 명시적 할당과 함께 호출해서요. 다음 예시는 인스턴스 멤버self.foo가 템플릿 필드임에도 전혀 할당되지 않아 유효하지 않아요:
class HelloOperator(BaseOperator):
template_fields = ("foo", "bar")
def __init__(self, foo, bar) -> None:
self.bar = bar
다음 예시도 MyHelloOperator의 인스턴스 멤버 self.foo가 부모 생성자에 전달된 kwargs의 일부로 암묵적으로 초기화되어 유효하지 않아요:
class HelloOperator(BaseOperator):
template_fields = "foo"
def __init__(self, foo) -> None:
self.foo = foo
class MyHelloOperator(HelloOperator):
template_fields = ("foo", "bar")
def __init__(self, bar, **kwargs) -> None: # should be def __init__(self, foo, bar, **kwargs)
super().__init__(**kwargs) # should be super().__init__(foo=foo, **kwargs)
self.bar = bar
- 생성자에서 할당하는 동안 파라미터에 작업을 적용하는 것은 허용되지 않아요. 값에 대한 어떤 작업도
execute()메서드에서 적용해야 해요. 따라서 다음 예시는 유효하지 않아요:
class HelloOperator(BaseOperator):
template_fields = "foo"
def __init__(self, foo) -> None:
self.foo = foo.lower() # assignment should be only self.foo = foo
Operator가 기본 operator에서 상속받고 자체 생성자가 없을 때는 위 제한 사항이 적용되지 않아요. 다만 템플릿 필드는 그 제한 사항에 따라 부모에서 제대로 설정되어야 해요.
따라서 다음 예시는 유효해요:
class HelloOperator(BaseOperator):
template_fields = "foo"
def __init__(self, foo) -> None:
self.foo = foo
class MyHelloOperator(HelloOperator):
template_fields = "foo"
위 제한 사항은 'validate-operators-init'이라는 prek hook에 의해 강제돼요.
서브클래싱으로 템플릿 필드 추가하기 (Add template fields with subclassing)
커스텀 operator를 만드는 흔한 용도는 기존 template_fields를 단순히 확장하는 것이에요. 사용하려는 operator가 특정 파라미터를 템플릿으로 정의하지 않는데, 인자를 Jinja 표현식으로 동적으로 전달하고 싶은 상황이 있을 수 있어요. 이는 기존 operator를 빠르게 서브클래싱해 쉽게 달성할 수 있어요.
앞서 정의한 HelloOperator를 사용하고 싶다고 가정해 봐요:
class HelloOperator(BaseOperator):
template_fields: Sequence[str] = ("name",)
def __init__(self, name: str, world: str, **kwargs) -> None:
super().__init__(**kwargs)
self.name = name
self.world = world
def execute(self, context):
message = f"Hello {self.world} it's {self.name}!"
print(message)
return message
하지만 world 인자를 동적으로 매개변수화하고 싶어요. template_fields 속성이 Sequence[str] 타입(즉 문자열의 리스트 또는 튜플)임이 보장되므로, HelloOperator를 서브클래스화해 원하는 대로 template_fields를 쉽게 수정할 수 있어요.
class MyHelloOperator(HelloOperator):
template_fields: Sequence[str] = (*HelloOperator.template_fields, "world")
이제 MyHelloOperator를 이렇게 사용할 수 있어요:
with dag:
hello_task = MyHelloOperator(
task_id="task_id_1",
name="{{ task_instance.task_id }}",
world="{{ var.value.my_world }}",
)
이 예시에서 world 인자는 Jinja 표현식을 통해 "my_world"라는 이름의 Airflow Variable 값으로 동적으로 설정돼요.
Operator extra link 정의하기 (Define an operator extra link)
operator를 위해 사용자를 외부 시스템으로 리다이렉트하는 extra link를 정의할 수 있어요. 예를 들어 사용자를 operator의 매뉴얼로 리다이렉트하는 링크를 추가할 수 있어요.
Sensors
Airflow는 성공 기준이 충족될 때까지 정기 간격으로 어떤 상태(예: 파일 존재)를 폴링하는 것이 목적인 특별한 종류의 operator를 위한 프리미티브를 제공해요.
외부 상태를 폴링하고 성공 기준을 평가하는 poke 메서드를 정의해 airflow.sensors.base.BaseSensorOperator를 확장해 원하는 어떤 sensor든 만들 수 있어요.
Sensors에는 'reschedule' 모드라는 강력한 기능이 있어요. 이는 sensor 태스크가 poke 사이에 worker 슬롯을 차단하는 대신 재스케줄될 수 있게 해요. 이는 더 긴 폴링 간격을 견딜 수 있고 오랫동안 폴링할 것으로 예상될 때 유용해요.
Reschedule 모드에는 sensor가 재스케줄 실행 사이에 내부 상태를 유지할 수 없다는 함정이 있어요. 이 경우 sensor를 airflow.sensors.base.poke_mode_only()로 데코레이션해야 해요. 이렇게 하면 사용자에게 sensor가 reschedule 모드에 적합하지 않다는 것을 알릴 수 있어요.
내부 상태를 유지하고 reschedule 모드에 사용할 수 없는 sensor의 예시는 airflow.providers.google.cloud.sensors.gcs.GCSUploadSessionCompleteSensor이에요. 프리픽스의 객체 수를 폴링하고(이 수가 sensor의 내부 상태) 객체 수가 일정 시간 동안 변하지 않으면 성공해요.
Operator 테스트하기 (Testing your operator)
Operator를 인스턴스화하고 코드가 읽는 context 키와 함께 execute()(sensor는 poke())를 호출해요. 이는 Dag도, Dag run도, 메타데이터 데이터베이스도 필요 없어요:
def test_hello_operator():
op = HelloOperator(task_id="hello", name="Bob")
assert op.execute(context={}) == "Hello Bob"
렌더링된 template fields를 검증해야 한다면 op.render_template_fields(context)를 사용하세요.
전체 Dag run을 실행해 보고 싶을 때만 dag.test()를 사용하세요. 그것은 통합 테스트로, 메타데이터 데이터베이스와 Airflow가 직렬화할 수 있는 파일에 정의된 Dag가 필요해요. deferrable operator를 포함한 전체 패턴 세트는 Unit tests를 참고하세요.