本文详细介绍了如何在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。
from airflow import DAG
from airflow.sensors.python import PythonSensor
from airflow.operators.dummy import DummyOperator
from datetime import datetime, timedelta
# 导入上面定义的日期检查函数
# from your_module import is_last_tuesday_of_month
# 或者直接将函数定义在DAG文件顶部
with DAG(
dag_id='conditional_last_tuesday_dag',
start_date=datetime(2023, 1, 1),
schedule_interval='@daily', # 每天运行，传感器会进行判断
catchup=False,
tags=['sensor', 'condition', 'date'],
) as dag:
# 传感器任务：检查是否为当月最后一个周二
check_last_tuesday = PythonSensor(
task_id='wait_for_last_tuesday',
python_callable=is_last_tuesday_of_month,
poke_interval=60 * 60 * 24, # 每天检查一次 (24小时)
timeout=60 * 60 * 24 * 31, # 最大等待时间，约一个月，防止无限等待
mode='poke', # 默认模式，周期性调用callable
soft_fail=True, # 如果超时，将任务标记为skipped而不是failed
# op_kwargs={'some_arg': 'value'} # 如果你的callable需要额外的固定参数，可以通过这里传递
)
# 核心业务任务，只有当传感器成功时才运行
T1 = DummyOperator(task_id='delete_gcs_files')
T2 = DummyOperator(task_id='run_sql_query_1')
T3 = DummyOperator(task_id='run_sql_query_2')
T4 = DummyOperator(task_id='run_sql_query_3_to_gcs')
T5 = DummyOperator(task_id='copy_ref_to_history_table')
# 定义任务依赖关系
# 传感器成功后，才执行T1及后续任务
check_last_tuesday >> T1
T1 >> T2
T2 >> T3
T3 >> T4
T4 >> T5
代码说明：
- schedule_interval='@daily'：DAG被配置为每天触发一次。这意味着check_last_tuesday传感器将每天运行。
- PythonSensor：
- task_id：任务的唯一标识符。
- python_callable：指定要执行的Python函数。
- poke_interval：传感器两次检查条件之间的时间间隔（秒）。这里设置为每天检查一次。
- timeout：传感器在放弃并标记为失败（或跳过）之前等待条件的最大时间（秒）。这里设置为大约一个月，确保它有足够时间等待到当月最后一个周二。
- mode='poke'：传感器会周期性地调用python_callable。
- soft_fail=True：这是实现“停止并运行任何任务”的关键。如果check_last_tuesday传感器在timeout时间内未能返回True，它将不会失败整个DAG，而是将自身标记为skipped。由于依赖关系，所有下游任务（T1到T5）也将被skipped，从而有效地阻止了它们的执行。
注意事项与最佳实践
-
传感器超时与重试：
- poke_interval和timeout的设置至关重要。对于日期条件，poke_interval可以设置为24小时（即每天检查一次），timeout则应覆盖一个完整周期的最长时间（例如，一个月）。
- 如果soft_fail=False，传感器超时会导致任务失败，进而可能导致整个DAG失败。使用soft_fail=True可以更优雅地跳过不符合条件的执行。
-
python_callable的幂等性：
- python_callable函数应该没有副作用，并且每次调用都应返回相同的结果，直到条件满足。我们的is_last_tuesday_of_month函数是幂等的，因为它只基于输入日期进行计算。
-
日志记录：
- 在python_callable函数中添加清晰的print语句（如示例所示），这些信息会显示在Airflow任务日志中，有助于调试和理解传感器的工作状态。
-
替代方案（适用于不同场景）：
- BranchPythonOperator： 如果你需要在条件满足时走A路径，不满足时走B路径（而不是直接停止），BranchPythonOperator会是更好的选择。它允许你动态地选择下一个要执行的任务ID。
- ShortCircuitOperator： 如果条件不满足时，你希望立即跳过所有下游任务，但不需要等待，ShortCircuitOperator是一个更轻量级的选择。它的python_callable返回True继续，False则跳过。
-
资源消耗：
- 传感器会占用一个Worker槽位，直到其条件满足或超时。长时间运行或频繁轮询的传感器可能会消耗较多资源。合理设置poke_interval和timeout是关键。对于每天检查一次的日期条件，资源消耗相对较低。
总结
通过PythonSensor，Airflow为复杂的条件性工作流提供了强大的支持。结合自定义的Python逻辑，我们可以灵活地在DAG执行前进行各种检查，确保任务只在符合预设条件时才运行。这种机制极大地增强了Airflow DAG的鲁棒性和适应性，使其能够处理更加复杂的业务逻辑和调度需求。理解并熟练运用传感器，是构建高效、可靠Airflow工作流的关键一步。
