如何将Airflow自定义算子序列打包以实现复用？

方法1：用Python函数封装任务序列 #

把创建三个算子、设置依赖的逻辑封装成可复用的函数，传入必要参数即可一键插入DAG。这种方式简单直接，适合快速复用。

示例代码：

def add_a_b_c_tasks(
    dag,
    project_id,
    region,
    body_str_data,
    task_prefix=""  # 可选，避免多个序列的task_id重复
):
    # 创建A算子
    a_data = AOperator(
        task_id=f"{task_prefix}a_data",
        project_id=project_id,
        region=region,
        body=body_str_data,
        dag=dag
    )
    # 创建B算子
    b_data = BOperator(
        task_id=f"{task_prefix}b_data",
        project_id=project_id,
        region=region,
        dag=dag
    )
    # 创建C算子
    c_data = COperator(
        task_id=f"{task_prefix}c_data",
        project_id=project_id,
        region=region,
        dag=dag
    )
    # 设置任务依赖
    a_data >> b_data >> c_data
    # 返回最后一个任务，方便和其他任务关联
    return c_data

调用方式：

from airflow import DAG
from datetime import datetime

PROJECT_ID = "your-project-id"
REGION = "your-region"
body_str_data = "your-body-content"

with DAG(
    dag_id="example_dag",
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False
) as dag:
    # 插入A->B->C任务序列
    abc_last_task = add_a_b_c_tasks(
        dag=dag,
        project_id=PROJECT_ID,
        region=REGION,
        body_str_data=body_str_data,
        task_prefix="batch1_"
    )
    # 可继续关联其他任务
    # another_task = SomeOperator(task_id="another_task", dag=dag)
    # abc_last_task >> another_task

方法2：用TaskGroup封装（Airflow 2.0+推荐） #

如果需要在Airflow UI上看到分组的任务结构，推荐使用TaskGroup。它会把三个任务打包成一个逻辑组，既复用代码，又提升可视化效果。

示例代码：

from airflow.utils.task_group import TaskGroup

def create_a_b_c_task_group(
    project_id,
    region,
    body_str_data,
    group_id="a_b_c_task_group"  # 可选，分组ID需唯一
):
    with TaskGroup(group_id=group_id) as task_group:
        a_data = AOperator(
            task_id="a_data",
            project_id=project_id,
            region=region,
            body=body_str_data
        )
        b_data = BOperator(
            task_id="b_data",
            project_id=project_id,
            region=region
        )
        c_data = COperator(
            task_id="c_data",
            project_id=project_id,
            region=region
        )
        # 设置依赖
        a_data >> b_data >> c_data
    return task_group

调用方式：

with DAG(
    dag_id="example_dag_with_group",
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False
) as dag:
    # 插入任务组
    abc_group = create_a_b_c_task_group(
        project_id=PROJECT_ID,
        region=REGION,
        body_str_data=body_str_data,
        group_id="batch_abc_group"
    )
    # 关联其他任务时直接用分组对象即可
    # another_task = SomeOperator(task_id="another_task")
    # abc_group >> another_task

注意事项 #

• 确保task_id全局唯一：使用函数时可通过task_prefix区分不同序列；使用TaskGroup时，分组内的task_id会自动带上分组前缀，无需额外处理。
• 灵活扩展参数：如果算子有更多可选配置，可以在封装函数中添加默认参数，或允许传入**kwargs来覆盖默认设置。
• 兼容性：TaskGroup仅在Airflow 2.0及以上版本支持，若使用旧版本，优先选择函数封装方式。

内容的提问来源于stack exchange，提问作者Brian Jones