在Airflow中实现条件性任务跳过：short_circuit操作符指南

from __future__ import annotations

import pendulum

from airflow.decorators import dag, task

@dag(
    dag_id="conditional_tasks_example",
    schedule=None,
    start_date=pendulum.datetime(2023, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example", "conditional"],
)
def dag_runner():
    @task(task_id="get_data_src_a")
    def get_data_src_a() -> list:
        """模拟从数据源A获取用户列表"""
        print("Fetching data from source A...")
        return ["user1", "user2", "user3"]

    @task(task_id="get_data_src_b")
    def get_data_src_b() -> list:
        """模拟从数据源B获取用户列表"""
        print("Fetching data from source B...")
        # 尝试返回空列表以测试条件跳过
        return ["user1", "user4"] 
        # return [] # 改变此行以测试跳过场景

    @task(task_id="find_uniq_users")
    def find_uniq_users(users_from_a: list, users_from_b: list) -> list:
        """找出在源A中但不在源B中的唯一用户"""
        print(f"Users from A: {users_from_a}")
        print(f"Users from B: {users_from_b}")
        uniq_users = [u for u in users_from_a if u not in users_from_b]
        print(f"Unique users: {uniq_users}")
        return uniq_users

    @task(task_id="do_something_with_users")
    def do_something_with_users(uniq_users: list):
        """对唯一用户执行一些操作"""
        print(f"Performing action for unique users: {uniq_users}")
        # 实际操作，例如发送通知、更新数据库等
        if not uniq_users:
            print("No unique users to process, this message should ideally not appear if task is skipped.")

    users_from_a_data = get_data_src_a()
    users_from_b_data = get_data_src_b()
    unique_users_result = find_uniq_users(users_from_a_data, users_from_b_data)
    do_something_with_users(unique_users_result)

dag_runner()

from __future__ import annotations

import pendulum

from airflow.decorators import dag, task

@dag(
    dag_id="conditional_tasks_with_short_circuit",
    schedule=None,
    start_date=pendulum.datetime(2023, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example", "conditional", "short_circuit"],
)
def dag_runner_with_conditional():
    @task(task_id="get_data_src_a")
    def get_data_src_a() -> list:
        print("Fetching data from source A...")
        return ["user1", "user2", "user3"]

    @task(task_id="get_data_src_b")
    def get_data_src_b() -> list:
        print("Fetching data from source B...")
        # 场景一：有唯一用户，do_something_with_users 会执行
        return ["user1", "user4"] 
        # 场景二：没有唯一用户，do_something_with_users 会被跳过
        # return ["user1", "user2", "user3"] 

    @task(task_id="find_uniq_users")
    def find_uniq_users(users_from_a: list, users_from_b: list) -> list:
        print(f"Users from A: {users_from_a}")
        print(f"Users from B: {users_from_b}")
        uniq_users = [u for u in users_from_a if u not in users_from_b]
        print(f"Unique users found: {uniq_users}")
        return uniq_users

    # 引入 short_circuit 任务
    @task.short_circuit(task_id="should_process_unique_users")
    def should_process_unique_users(uniq_users: list) -> bool:
        """
        根据唯一用户列表是否为空来决定是否继续执行下游任务。
        如果列表非空，返回 True；否则返回 False。
        """
        if uniq_users:
            print(f"Unique users found: {len(uniq_users)}. Proceeding to process.")
            return True
        else:
            print("No unique users found. Skipping downstream processing task.")
            return False

    @task(task_id="do_something_with_users")
    def do_something_with_users(uniq_users: list):
        print(f"Performing action for unique users: {uniq_users}")
        # 实际操作，例如发送通知、更新数据库等
        # 如果此任务被跳过，此处的 print 不会执行

    users_from_a_data = get_data_src_a()
    users_from_b_data = get_data_src_b()
    unique_users_result = find_uniq_users(users_from_a_data, users_from_b_data)

    # 任务依赖现在包括 short_circuit 任务
    # unique_users_result 传递给 short_circuit 任务
    # short_circuit 任务的输出（True/False）决定了 do_something_with_users 是否执行
    should_process_unique_users_result = should_process_unique_users(unique_users_result)
    should_process_unique_users_result >> do_something_with_users(unique_users_result)

dag_runner_with_conditional()