请教：如何在Airflow中编写可延迟（Deferrable）的SqlSensor？

针对你遇到的EOD结算等待场景，下面直接给出可延迟（Deferrable）SqlSensor的完整实现，以及关键逻辑的说明：

核心原理 #

Deferrable算子的核心是将等待逻辑转移到Airflow的Triggerer组件，而非占用Worker Slot。当条件不满足时，算子会暂停并释放Slot，Triggerer会异步轮询检查条件，满足后再唤醒算子继续执行。

完整代码实现 #

from airflow.sensors.base import BaseSensorOperator, DeferrableMixin
from airflow.triggers.base import BaseTrigger, TriggerEvent
from airflow.hooks.sql import SqlHook
import asyncio


class SqlTrigger(BaseTrigger):
    def __init__(self, sql: str, conn_id: str, poll_interval: float, success_condition: str = "return_value"):
        self.sql = sql
        self.conn_id = conn_id
        self.poll_interval = poll_interval
        self.success_condition = success_condition

    def serialize(self):
        # 序列化Trigger参数，用于Triggerer重启时恢复状态
        return {
            "sql": self.sql,
            "conn_id": self.conn_id,
            "poll_interval": self.poll_interval,
            "success_condition": self.success_condition
        }

    async def run(self, trigger_context):
        # 异步轮询检查SQL条件
        while True:
            # 使用同步Hook的异步执行方式（避免阻塞Triggerer）
            hook = SqlHook(sql_conn_id=self.conn_id)
            result = hook.get_first(self.sql)
            
            # 根据success_condition判断是否满足条件
            # 默认判断查询结果是否非空，你可以自定义逻辑（比如result[0] > 0）
            if eval(self.success_condition, {"return_value": result}):
                yield TriggerEvent(True)
                return
            
            # 等待指定间隔后再次检查
            await asyncio.sleep(self.poll_interval)


class DeferrableSqlSensor(BaseSensorOperator, DeferrableMixin):
    def __init__(self, sql: str, conn_id: str, poll_interval: float = 60, success_condition: str = "return_value is not None", **kwargs):
        super().__init__(**kwargs)
        self.sql = sql
        self.conn_id = conn_id
        self.poll_interval = poll_interval
        self.success_condition = success_condition

    def execute(self, context):
        # 先同步检查一次条件，避免直接进入延迟状态
        hook = SqlHook(sql_conn_id=self.conn_id)
        result = hook.get_first(self.sql)
        
        if eval(self.success_condition, {"return_value": result}):
            return
        
        # 条件不满足时，触发延迟并释放Worker Slot
        self.defer(
            trigger=SqlTrigger(
                sql=self.sql,
                conn_id=self.conn_id,
                poll_interval=self.poll_interval,
                success_condition=self.success_condition
            ),
            method_name="execute_complete"
        )

    def execute_complete(self, context, event=None):
        # 延迟结束后执行的回调方法
        if event is True:
            self.log.info("SQL条件满足，传感器执行完成")
            return
        raise ValueError("传感器执行失败或超时")

关键部分说明 #

1. SqlTrigger类

   • serialize方法：必须实现，用于Triggerer组件重启时序列化参数，恢复轮询状态
   • run方法：异步循环执行SQL查询，通过eval解析success_condition判断是否满足条件，满足则返回TriggerEvent唤醒算子

2. DeferrableSqlSensor类

   • 继承DeferrableMixin：标记该算子为可延迟类型
   • execute方法：先同步检查一次条件（优化性能），不满足则调用defer方法，将等待逻辑交给Triggerer
   • execute_complete方法：延迟结束后的回调，处理触发结果

使用示例 #

在你的DAG中直接调用这个传感器即可：

from airflow import DAG
from datetime import datetime
# 导入上面定义的DeferrableSqlSensor

with DAG(
    dag_id="eod_settlement_waiter",
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False
) as dag:
    wait_for_eod_settlement = DeferrableSqlSensor(
        task_id="wait_for_eod_settlement",
        sql="SELECT COUNT(*) FROM settlement_records WHERE settlement_date = CURRENT_DATE AND status = 'COMPLETED'",
        conn_id="your_sql_connection",  # 替换为你的数据库连接ID
        poll_interval=300,  # 每5分钟轮询一次
        success_condition="return_value > 0",  # 当结算记录数大于0时满足条件
        timeout=3600 * 4,  # 超时时间4小时
        mode="deferrable"  # 显式指定模式（Airflow 2.2+可省略，因为继承了DeferrableMixin）
    )

    # 后续任务
    process_settlement = ...

    wait_for_eod_settlement >> process_settlement

注意事项 #

• 确保你的Airflow版本支持Deferrable算子（Airflow 2.2+），并且Triggerer组件已启动
• success_condition支持自定义表达式，比如return_value[0] == 'SUCCESS'，请根据你的实际SQL返回结果调整
• 数据库连接ID需要提前在Airflow的Connections中配置好

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