如何基于SqlSensor的成功/跳过状态在Airflow DAG中分支任务？

你的核心需求是：SqlSensor检查到数据则执行转换任务，无数据则DAG直接成功结束。之前用BranchPythonOperator结合XCom失效的原因是默认SqlSensor无论成功还是soft_fail跳过，都不会推送XCom数据，所以无法通过XCom判断分支。

下面提供两种直接可行的实现方案：

方案一：用ShortCircuitOperator替代（推荐，更简洁） #

ShortCircuitOperator的逻辑是：若执行结果为True，则继续执行下游任务；若为False，则跳过所有下游任务，DAG直接标记为成功结束，完美匹配你的需求。

完整代码示例 #

from airflow import DAG
from airflow.operators.python import ShortCircuitOperator
from airflow.providers.oracle.hooks.oracle import OracleHook
from airflow.operators.dummy import DummyOperator
from datetime import datetime

def check_for_records():
    # 用OracleHook执行SQL查询
    hook = OracleHook(conn_id='oracle_db')
    sql = "select count(*) from your_table where load_date = current_date"  # 替换为你的实际SQL
    count = hook.get_first(sql)[0]
    # 返回True（有数据）则继续下游，False（无数据）则短路结束
    return count > 0

with DAG(
    dag_id='oracle_data_processing',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False
) as dag:
    check_data = ShortCircuitOperator(
        task_id='check_for_records_with_load_date',
        python_callable=check_for_records
    )

    transform_task = DummyOperator(task_id='run_transformations')  # 替换为你的实际转换任务

    check_data >> transform_task

方案二：自定义SqlSensor推送XCom（若坚持用Sensor） #

如果你一定要保留SqlSensor，可以自定义一个继承自SqlSensor的类，在任务结束时推送XCom记录状态，再用BranchPythonOperator判断分支。

自定义Sensor代码 #

from airflow.providers.oracle.sensors.oracle import OracleSensor
from airflow.utils.state import State

class CustomSqlSensor(OracleSensor):
    def execute(self, context):
        try:
            # 执行原传感器逻辑
            super().execute(context)
            # 成功时推送XCom：有数据
            context['ti'].xcom_push(key='has_data', value=True)
        except Exception as e:
            # 因soft_fail被跳过的情况，推送XCom：无数据
            if self.soft_fail:
                context['ti'].xcom_push(key='has_data', value=False)
            raise e

# 后续使用这个自定义传感器替代原SqlSensor
oracle_sensor = CustomSqlSensor(
    task_id='check_for_records_with_load_date',
    conn_id='oracle_db',
    sql="select count(*) etc",
    mode='poke', 
    timeout=60, 
    soft_fail=True 
)

# 分支判断函数
def decide_next_task(**context):
    has_data = context['ti'].xcom_pull(task_ids='check_for_records_with_load_date', key='has_data')
    if has_data:
        return 'run_transformations'
    else:
        return 'end_dag_success'

# 分支Operator和结束任务
branch_op = BranchPythonOperator(
    task_id='branch_decision',
    python_callable=decide_next_task,
    provide_context=True
)

transform_task = DummyOperator(task_id='run_transformations')
end_task = DummyOperator(task_id='end_dag_success', trigger_rule='none_failed_min_one_success')

oracle_sensor >> branch_op >> [transform_task, end_task]

注意点 #

• 方案二中的end_task需要设置trigger_rule='none_failed_min_one_success'，确保即使分支走end_dag_success，整个DAG状态为成功。
• 自定义传感器需要处理soft_fail的异常捕获，确保跳过状态下能正确推送XCom。

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