在Apache Airflow中实现基于日期的条件性DAG执行

本文详细介绍了如何在Apache Airflow中利用PythonSensor实现复杂的日期条件判断，从而精确控制DAG的任务执行流程。通过一个“每月最后一个周二”的实际案例，文章演示了如何编写自定义Python函数来检查特定日期条件，并将其集成到PythonSensor中，以决定是否继续执行下游任务，从而实现灵活且健壮的自动化工作流。

在自动化工作流管理平台Apache Airflow中，经常需要根据特定条件来决定是否执行一系列任务。这些条件可能包括外部文件是否存在、数据库状态、API响应，或者像本文将探讨的——基于日期的复杂逻辑。当我们需要在DAG运行前进行条件判断，并且如果条件不满足就停止后续任务时，Airflow的传感器（Sensor）机制提供了一个优雅且强大的解决方案。

Airflow中的条件执行与传感器

Airflow提供了多种实现条件逻辑的Operator，例如BranchPythonOperator用于根据条件选择不同的分支路径，或者ShortCircuitOperator用于在条件不满足时跳过下游任务。然而，当我们需要在任务开始前“等待”某个条件满足，或者当条件不满足时直接“阻止”下游任务执行时，传感器（Sensor）是更合适的选择。

传感器是一种特殊的Operator，它会周期性地检查某个条件，直到条件满足才会标记为成功并触发下游任务。如果条件在设定的超时时间内仍未满足，传感器可以选择失败或跳过（soft_fail=True）下游任务。对于“如果条件不满足就停止所有任务”的需求，PythonSensor结合适当的配置可以完美实现。

使用PythonSensor实现日期条件判断

本教程将以一个具体的场景为例：一个DAG只有在“每月最后一个周二”才需要运行其核心业务逻辑。如果不是，则不执行任何后续任务。

1. 定义日期条件检查函数

首先，我们需要一个Python函数来判断给定的日期是否是当月的最后一个周二。这个函数将作为PythonSensor的python_callable参数。

from datetime import datetime, timedelta
import calendar

def is_last_tuesday_of_month(**kwargs):
    """
    检查Airflow的执行日期是否是当月的最后一个周二。
    如果满足条件，返回True；否则返回False。
    """
    # 从Airflow上下文中获取执行日期
    # 'ds' 格式为 'YYYY-MM-DD'
    ds = kwargs.get('ds')
    if not ds:
        # 如果在Airflow上下文之外测试，可以使用当前日期
        print("警告: 'ds' 未在kwargs中找到。使用当前日期进行检查。")
        execution_date = datetime.now().date()
    else:
        execution_date = datetime.strptime(ds, '%Y-%m-%d').date()

    year = execution_date.year
    month = execution_date.month

    # 获取当前月的最后一天
    _, num_days = calendar.monthrange(year, month)
    last_day_of_month = datetime(year, month, num_days).date()

    # 从当月最后一天开始向前迭代，查找最后一个周二
    current_check_date = last_day_of_month
    while current_check_date.month == month:
        if current_check_date.weekday() == calendar.TUESDAY:
            # 找到了当月的最后一个周二
            is_match = (current_check_date == execution_date)
            print(f"执行日期: {execution_date}, 当月最后一个周二: {current_check_date}. 匹配结果: {is_match}")
            return is_match
        current_check_date -= timedelta(days=1)

    # 对于有效月份，通常不会执行到这里
    print(f"在 {year} 年 {month} 月中未找到周二。")
    return False

登录后复制

函数说明：

• kwargs.get('ds')：这是Airflow传递给python_callable的一个上下文变量，代表DAG的执行日期（YYYY-MM-DD格式）。
• calendar.monthrange(year, month)：返回指定月份的天数。
• current_check_date.weekday() == calendar.TUESDAY：weekday()返回0-6的整数（周一为0，周日为6），calendar.TUESDAY常量值为1。
• 函数通过从当月最后一天向前迭代，找到第一个（即最后一个）周二，然后与执行日期进行比较。

2. 将函数集成到PythonSensor

接下来，我们将上述函数集成到Airflow DAG中，使用PythonSensor。

标签： python apache ai python函数 yy