如何在Airflow中为每个任务执行通用Python逻辑且避免代码重复？

你提到的需求完全可以实现，Airflow提供了几种类似Java切面（AOP）的方案，能给任务统一添加日志、HTTP调用、XCom处理等通用功能，不用重复编写代码。下面是三种最常用的实现方式：

一、自定义通用Operator包装 #

创建一个基础Operator，把通用逻辑封装进去，所有任务都通过这个Operator来包装原任务，既保留原任务的功能，又自动添加通用逻辑。

示例代码：

from airflow.models.baseoperator import BaseOperator
from airflow.utils.decorators import apply_defaults
import logging
import requests

class BaseCustomOperator(BaseOperator):
    @apply_defaults
    def __init__(self, original_task, **kwargs):
        super().__init__(**kwargs)
        self.original_task = original_task

    def execute(self, context):
        # 通用前置逻辑：打印任务ID日志
        logging.info(f"Executing task: {self.task_id}")
        # 发起HTTP调用（可根据需求调整参数）
        requests.get("https://your-webhook-url.com/notify", params={"task_id": self.task_id})
        
        # 执行原任务的核心逻辑
        result = self.original_task.execute(context)
        
        # 通用后置逻辑：统一处理XCom参数
        ti = context['ti']
        raw_xcom = ti.xcom_pull(task_ids=self.task_id)
        if raw_xcom:
            # 示例：修改XCom内容后重新推送
            processed_xcom = {"status": "processed", "data": raw_xcom}
            ti.xcom_push(key="processed_data", value=processed_xcom)
        
        return result

使用方式：

from airflow.operators.python import PythonOperator
from datetime import datetime
from airflow import DAG

# 原任务函数
def task1_func():
    return "task1 completed"

def task2_func():
    return "task2 completed"

def task3_func():
    return "task3 completed"

with DAG(dag_id="wrapped_operator_dag", start_date=datetime(2024,1,1), schedule_interval="@daily") as dag:
    # 用自定义Operator包装每个原任务
    task1 = BaseCustomOperator(
        task_id="task1",
        original_task=PythonOperator(task_id="inner_task1", python_callable=task1_func)
    )
    task2 = BaseCustomOperator(
        task_id="task2",
        original_task=PythonOperator(task_id="inner_task2", python_callable=task2_func)
    )
    task3 = BaseCustomOperator(
        task_id="task3",
        original_task=PythonOperator(task_id="inner_task3", python_callable=task3_func)
    )
    
    # 依赖关系保持简洁
    task1 >> task2 >> task3

二、全局/任务级回调函数 #

Airflow支持设置任务执行的回调函数，相当于在任务执行前后插入切面逻辑，不需要修改任务本身，最接近Java的AOP思想。

首先定义通用回调函数：

import logging
import requests

def pre_execution_callback(context):
    # 任务执行前的通用逻辑
    task_id = context['ti'].task_id
    logging.info(f"Starting task: {task_id}")
    requests.get("https://your-webhook-url.com/pre-exec", params={"task_id": task_id})

def post_success_callback(context):
    # 任务执行成功后的通用逻辑
    ti = context['ti']
    task_id = ti.task_id
    raw_xcom = ti.xcom_pull(task_ids=task_id)
    if raw_xcom:
        ti.xcom_push(key="processed_xcom", value=f"processed_{raw_xcom}")
    logging.info(f"Task finished: {task_id}")

全局配置（所有任务自动应用） #

在DAG定义时设置全局回调，整个DAG的任务都会自动触发这些逻辑：

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def task1_func():
    return "task1"

def task2_func():
    return "task2"

def task3_func():
    return "task3"

with DAG(
    dag_id="callback_dag",
    start_date=datetime(2024,1,1),
    schedule_interval="@daily",
    # 全局设置前置、后置回调
    on_execute_callback=pre_execution_callback,
    on_success_callback=post_success_callback
) as dag:
    task1 = PythonOperator(task_id="task1", python_callable=task1_func)
    task2 = PythonOperator(task_id="task2", python_callable=task2_func)
    task3 = PythonOperator(task_id="task3", python_callable=task3_func)
    
    task1 >> task2 >> task3

单个任务配置 #

如果只想给部分任务添加通用逻辑，直接在任务参数里设置：

task1 = PythonOperator(
    task_id="task1",
    python_callable=task1_func,
    on_execute_callback=pre_execution_callback,
    on_success_callback=post_success_callback
)

三、自定义任务装饰器 #

如果你用Airflow的@task装饰器定义任务，可以写一个包装装饰器，给任务函数统一添加通用逻辑。

示例代码：

import logging
import requests
from airflow.decorators import task
from datetime import datetime
from airflow import DAG

# 自定义切面装饰器
def task_aspect(func):
    def wrapper(**context):
        # 前置逻辑
        task_id = context['ti'].task_id
        logging.info(f"Executing task: {task_id}")
        requests.get("https://your-webhook-url.com/notify", params={"task_id": task_id})
        
        # 执行原任务函数
        result = func(**context)
        
        # 后置逻辑：处理XCom
        context['ti'].xcom_push(key="processed_result", value=f"processed_{result}")
        return result
    return wrapper

# 用装饰器包装任务函数
@task
@task_aspect
def task1_func(**context):
    return "task1 done"

@task
@task_aspect
def task2_func(**context):
    return "task2 done"

@task
@task_aspect
def task3_func(**context):
    return "task3 done"

with DAG(dag_id="decorator_dag", start_date=datetime(2024,1,1), schedule_interval="@daily") as dag:
    # 依赖关系依然简洁明了
    task1_func() >> task2_func() >> task3_func()

方案选择建议 #

• 自定义Operator：适合需要包装不同类型Operator、或者需要高度定制通用逻辑的场景，灵活性最强。
• 全局回调：最省心的方案，不需要修改任务代码，适合所有任务都需要相同通用逻辑的场景。
• 自定义装饰器：适合用@task装饰器定义的任务，代码侵入性低，写法简洁。

内容的提问来源于stack exchange，提问作者Panos K