1、自定义Operator概述：为什么需要自定义Operator、Airflow Operator架构、BaseOperator核心方法解析

1.1 为什么需要自定义Operator

说实话，我刚接触Airflow那会儿，也觉得内置Operator够用了。PythonOperator、BashOperator、EmailOperator……官方文档列了一大堆，感觉啥都能干。

但真正上了生产环境，你就发现问题了。我印象特别深，有一次要对接一个内部的老系统，它的API认证方式特别奇葩——先拿，有效期只有5分钟，每次请求前都得刷新。用PythonOperator硬写？当然可以。但那个DAG文件变得又臭又长，而且每个Task里都重复写那段认证逻辑。

你想想看，如果这个逻辑要改，你得去翻多少个Task？维护成本直接起飞。

所以，自定义Operator的核心价值就三点：

  • 复用性：把重复的逻辑封装起来，一次编写，到处调用
  • 可读性：DAG文件里一行MyApiOperator(task_id='sync_data')，比十行PythonOperator清爽得多
  • 可维护性：逻辑变更只改Operator本身，不用动DAG

我的经验：团队里如果超过3个DAG用到了同一段逻辑，我就建议封装成自定义Operator。别等到10个DAG都写死了再重构，那代价太大了。

1.2 Airflow Operator架构

要写自定义Operator，你得先搞清楚Airflow的Operator体系长什么样。说白了，它就是个三层结构：

  1. BaseOperator：所有Operator的祖宗，定义了执行框架和生命周期
  2. 内置Operator：Airflow官方实现的，比如PythonOperator、BashOperator
  3. 自定义Operator：你继承BaseOperator写的，或者继承内置Operator扩展的

我画了张图，帮你理清这个关系：

Airflow Operator 架构层次 BaseOperator 核心框架与生命周期 内置 Operator（官方实现） PythonOperator | BashOperator | EmailOperator | ... 继承 BaseOperator，提供具体实现 自定义 Operator 继承 BaseOperator 或内置 Operator 实现业务逻辑，封装复用 💡 关键理解 自定义Operator不一定要从BaseOperator直接继承。如果内置Operator能满足80%的需求， 继承它再扩展20%的逻辑，往往更省事。我一般优先考虑继承，不行再重写。

嗯，这里要注意一个点：Operator不是Task。Operator是定义「做什么」的类，Task是DAG里的一个实例。同一个Operator可以在不同DAG里实例化出多个Task。

1.3 BaseOperator核心方法解析

写自定义Operator，你绕不开BaseOperator。我挑几个最常用的方法，给你拆开讲讲。

方法 作用 是否必须重写
__init__ 初始化Operator实例，接收参数 通常需要
execute 核心执行逻辑，Task运行时调用 必须
pre_execute execute之前执行，用于准备工作 可选
post_execute execute之后执行，用于清理工作 可选
on_kill Task被中断时调用，用于资源释放 可选

__init__：参数接收与校验

这是你自定义Operator的「门面」。所有用户传进来的参数，都在这里接收和校验。我习惯在__init__里做三件事：

  • 调用父类的__init__，把基础参数传过去
  • 接收自定义参数，赋值给实例变量
  • 做参数校验，尽早暴露问题
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

class MyCustomOperator(BaseOperator):
    
    @apply_defaults
    def __init__(
        self,
        api_endpoint: str,
        : str = None,
        timeout: int = 30,
        *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.api_endpoint = api_endpoint
        self. = 
        self.timeout = timeout
        
        # 参数校验
        if not self.api_endpoint.startswith('http'):
            raise ValueError('api_endpoint 必须以 http 开头')

一个小技巧：@apply_defaults这个装饰器，能让你在DAG里用default_args给Operator传默认参数。我建议你加上，这样用户用起来更灵活。

execute：核心执行逻辑

这是你自定义Operator的「心脏」。Airflow调度器执行Task时，调的就是这个方法。它接收一个context参数，里面包含了DAG Run、Task Instance、执行日期等信息。

def execute(self, context):
    """
    核心执行逻辑
    context 里包含：dag_run, task_instance, execution_date, ti 等
    """
    self.log.info(f"开始调用API: {self.api_endpoint}")
    
    # 你的业务逻辑
    response = self._call_api()
    
    # 可以通过 context 获取执行上下文信息
    execution_date = context['execution_date']
    self.log.info(f"执行日期: {execution_date}")
    
    # 返回值会被存入 XCom，供下游Task使用
    return response.json()

我曾经踩过一个坑：在execute里直接抛了异常，但没有做重试处理。结果网络抖动了一下，整个DAG就挂了。后来我加了重试机制，稳多了。

注意：execute方法的返回值会被自动推送到XCom。如果你不想让下游Task拿到这个值，可以返回None。但别返回敏感信息，XCom默认是明文存储的。

pre_execute 与 post_execute：钩子方法

这两个方法，说白了就是给你在execute前后「插一脚」的机会。我一般这么用：

  • pre_execute：建立数据库连接、创建临时文件、获取锁
  • post_execute：关闭连接、删除临时文件、释放锁
def pre_execute(self, context):
    """execute 之前执行"""
    self.log.info("准备资源...")
    self.temp_file = tempfile.NamedTemporaryFile(delete=False)
    
def post_execute(self, context, result=None):
    """execute 之后执行，result 是 execute 的返回值"""
    self.log.info("清理资源...")
    if self.temp_file:
        os.unlink(self.temp_file.name)

on_kill：优雅中断

这个我一开始没太在意，直到有一次Task被手动标记为失败，但后台进程还在跑。嗯，从那以后我学乖了。

def on_kill(self):
    """Task被中断时调用"""
    self.log.warning("Task被中断，清理资源...")
    if hasattr(self, 'subprocess'):
        self.subprocess.terminate()

1.4 一个完整的自定义Operator示例

光说不练假把式。我给你写个完整的例子，把上面讲的方法串起来。

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
import requests
import tempfile
import os

class ApiDataOperator(BaseOperator):
    """从API拉取数据并保存到本地文件"""
    
    @apply_defaults
    def __init__(
        self,
        endpoint: str,
        output_path: str,
        headers: dict = None,
        *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.endpoint = endpoint
        self.output_path = output_path
        self.headers = headers or {}
        self._temp_file = None
        
    def pre_execute(self, context):
        self._temp_file = tempfile.NamedTemporaryFile(
            suffix='.json', delete=False
        )
        self.log.info(f"临时文件已创建: {self._temp_file.name}")
        
    def execute(self, context):
        self.log.info(f"请求API: {self.endpoint}")
        resp = requests.get(
            self.endpoint, 
            headers=self.headers,
            timeout=30
        )
        resp.raise_for_status()
        
        # 写入临时文件
        with open(self._temp_file.name, 'w') as f:
            f.write(resp.text)
            
        self.log.info(f"数据已保存到临时文件")
        return self._temp_file.name
        
    def post_execute(self, context, result=None):
        if result and os.path.exists(result):
            # 移动到最终路径
            import shutil
            shutil.move(result, self.output_path)
            self.log.info(f"文件已移动到: {self.output_path}")
            
    def on_kill(self):
        if self._temp_file and os.path.exists(self._temp_file.name):
            os.unlink(self._temp_file.name)
            self.log.info("临时文件已清理")

我的建议：刚开始写自定义Operator，别追求一步到位。先写个能跑的版本，跑通了再慢慢加pre_execute、on_kill这些。我见过太多人一开始想得太复杂，结果连execute都没调通。

1.5 避坑指南

最后，分享几个我踩过的坑：

  • 别在__init__里做耗时操作：__init__是在DAG解析阶段调用的，不是执行阶段。如果你在里面做网络请求，整个DAG解析都会卡住。
  • 日志用self.log，别用print：self.log是Airflow自带的日志系统，会按Task组织日志文件。用print的话，日志会跑到标准输出里，查起来很麻烦。
  • 参数校验要尽早：在__init__里做参数校验，能让用户在DAG解析阶段就发现问题，而不是等到Task运行时才报错。
  • 考虑幂等性：你的Operator可能会被重试。如果每次执行都会产生副作用（比如插入数据库），要确保重试不会造成数据重复。
← 上一章 📖 返回目录 下一章 →