166 lines
4.8 KiB
Python
166 lines
4.8 KiB
Python
# 导入 DAG 类，用于定义工作流
|
||
from airflow import DAG
|
||
|
||
# 导入 PythonOperator：执行普通的 Python 函数
|
||
# 导入 BranchPythonOperator：根据 Python 函数的返回值选择执行分支
|
||
from airflow.operators.python import PythonOperator, BranchPythonOperator
|
||
|
||
# 导入 DummyOperator：空操作任务，用于组织任务依赖关系
|
||
from airflow.operators.dummy import DummyOperator
|
||
|
||
# 导入 TriggerRule：用于设置任务的触发规则
|
||
from airflow.utils.trigger_rule import TriggerRule
|
||
|
||
# 导入 datetime，用于设置 DAG 的开始日期
|
||
from datetime import datetime
|
||
|
||
# 导入 random，用于生成随机数
|
||
import random
|
||
|
||
|
||
# ==================== 1. 定义分支选择函数 ====================
|
||
|
||
def choose_branch(**kwargs):
|
||
"""
|
||
分支选择函数。
|
||
|
||
BranchPythonOperator 执行该函数，
|
||
根据函数的返回值决定执行哪个下游任务。
|
||
|
||
**kwargs：接收 Airflow 自动传入的任务上下文参数。
|
||
"""
|
||
|
||
# 随机生成 1～3 之间的整数，包含 1 和 3
|
||
random_number = random.randint(1, 3)
|
||
|
||
# 将生成的随机数输出到任务日志中，方便调试
|
||
print(f'Random number generated: {random_number}')
|
||
|
||
# 根据随机数选择执行的分支
|
||
# 注意：返回值必须是下游任务的 task_id
|
||
|
||
if random_number == 1:
|
||
return 'task_a' # 随机数为 1，选择 A 分支
|
||
|
||
elif random_number == 2:
|
||
return 'task_b' # 随机数为 2，选择 B 分支
|
||
|
||
else:
|
||
return 'task_c' # 随机数为 3，选择 C 分支
|
||
|
||
|
||
# ==================== 2. 定义普通任务函数 ====================
|
||
|
||
def print_task(task_name):
|
||
"""
|
||
普通任务函数。
|
||
|
||
task_name：任务名称，用于区分不同的分支任务。
|
||
"""
|
||
|
||
# 输出当前执行的任务名称
|
||
print(f'Executing {task_name}')
|
||
|
||
|
||
# ==================== 3. 创建 DAG 工作流 ====================
|
||
|
||
with DAG(
|
||
# DAG 的唯一标识符
|
||
dag_id='branch_example',
|
||
|
||
# 设置调度周期为每天执行一次
|
||
schedule_interval='@daily',
|
||
|
||
# 设置 DAG 开始日期
|
||
start_date=datetime(2026, 10, 1),
|
||
|
||
# 开启历史补跑
|
||
# Airflow 会根据调度规则处理开始日期之后符合条件的历史调度周期
|
||
catchup=True,
|
||
|
||
# 设置 DAG 标签，方便在 Airflow 页面中查找
|
||
tags=['demo']
|
||
|
||
) as dag:
|
||
|
||
# ==================== 4. 定义分支选择任务 ====================
|
||
|
||
branching_task = BranchPythonOperator(
|
||
# 任务 ID
|
||
task_id='branching_task',
|
||
|
||
# 指定执行的 Python 函数
|
||
# 函数的返回值决定执行哪个下游分支
|
||
python_callable=choose_branch,
|
||
)
|
||
|
||
# ==================== 5. 定义 A 分支任务 ====================
|
||
|
||
task_a = PythonOperator(
|
||
# 任务 ID，必须与分支函数返回的名称一致
|
||
task_id='task_a',
|
||
|
||
# 使用 lambda 创建一个无参数函数
|
||
# 执行任务时，调用 print_task('Task A')
|
||
python_callable=lambda: print_task('Task A'),
|
||
)
|
||
|
||
# ==================== 6. 定义 B 分支任务 ====================
|
||
|
||
task_b = PythonOperator(
|
||
# B 分支的任务 ID
|
||
task_id='task_b',
|
||
|
||
# 执行时打印 Task B
|
||
python_callable=lambda: print_task('Task B'),
|
||
)
|
||
|
||
# ==================== 7. 定义 C 分支任务 ====================
|
||
|
||
task_c = PythonOperator(
|
||
# C 分支的任务 ID
|
||
task_id='task_c',
|
||
|
||
# 执行时打印 Task C
|
||
python_callable=lambda: print_task('Task C'),
|
||
)
|
||
|
||
# ==================== 8. 定义结束任务 ====================
|
||
|
||
end_task = DummyOperator(
|
||
# 结束任务的 ID
|
||
task_id='end_task',
|
||
|
||
# 设置触发规则：
|
||
# 所有上游任务中没有失败的任务，并且至少有一个上游任务成功，
|
||
# 当前任务就可以执行。
|
||
#
|
||
# 因为分支选择会把未选中的任务标记为 skipped，
|
||
# 所以不能简单使用默认的 all_success，
|
||
# 否则 end_task 可能因为存在 skipped 上游任务而被跳过。
|
||
trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
|
||
)
|
||
|
||
# ==================== 9. 设置任务依赖关系 ====================
|
||
|
||
# 分支任务执行完成后，根据返回值选择一个分支：
|
||
#
|
||
# 返回 'task_a' -> 执行 task_a
|
||
# 返回 'task_b' -> 执行 task_b
|
||
# 返回 'task_c' -> 执行 task_c
|
||
#
|
||
# 未被选中的分支会被标记为 skipped。
|
||
#
|
||
# 被选中的分支执行完成后，end_task 在满足触发规则时执行。
|
||
|
||
branching_task >> [task_a, task_b, task_c] >> end_task
|
||
|
||
# 上面的依赖关系等价于：
|
||
#
|
||
# branching_task >> task_a
|
||
# branching_task >> task_b
|
||
# branching_task >> task_c
|
||
#
|
||
# task_a >> end_task
|
||
# task_b >> end_task
|
||
# task_c >> end_task |