在 Airflow 中实现基于日期条件的任务流控制

from datetime import datetime, timedelta
import calendar

def is_last_tuesday_of_month(execution_date: datetime, **context) -> bool:
    """
    检查给定的 execution_date 是否是该月的最后一个周二。
    """
    # 获取当前月份的最后一天
    year = execution_date.year
    month = execution_date.month
    last_day_of_month = calendar.monthrange(year, month)[1]

    # 从最后一天开始向前查找第一个周二
    current_day = datetime(year, month, last_day_of_month)

    # 遍历直到找到周二 (weekday() 返回 0-6，周一到周日)
    # 周二的 weekday() 值为 1
    while current_day.weekday() != calendar.TUESDAY:
        current_day -= timedelta(days=1)

    # 检查找到的周二是否与 execution_date 的日期部分相同
    # 考虑到 execution_date 通常是 DAG 运行的开始时间，我们只关心日期部分
    if current_day.date() == execution_date.date():
        print(f"条件满足：{execution_date.date()} 是 {month} 月的最后一个周二。")
        return True
    else:
        print(f"条件不满足：{execution_date.date()} 不是 {month} 月的最后一个周二。")
        return False

from airflow import DAG
from airflow.sensors.python import PythonSensor
from airflow.operators.bash import BashOperator
from datetime import datetime, timedelta
import calendar

# 导入上面定义的条件函数
# from your_module import is_last_tuesday_of_month 
# 如果函数定义在同一个文件中，则不需要导入

# DAG 定义
with DAG(
    dag_id='conditional_last_tuesday_run',
    start_date=datetime(2023, 1, 1),
    schedule_interval='0 0 * * *', # 每天运行一次，以便检查条件
    catchup=False,
    tags=['example', 'sensor', 'condition'],
) as dag:

    # 定义 PythonSensor 任务
    check_last_tuesday = PythonSensor(
        task_id='check_if_last_tuesday_of_month',
        python_callable=is_last_tuesday_of_month,
        # op_kwargs 允许向 python_callable 传递额外的关键字参数
        # 在这里，我们将 execution_date 显式传递，以确保函数使用正确的日期
        op_kwargs={'execution_date': '{{ ds }}'}, # {{ ds }} 是 Airflow 模板，解析为执行日期字符串
        poke_interval=60 * 5, # 每 5 分钟检查一次条件
        timeout=60 * 60 * 24, # 最多等待 24 小时
        mode='poke', # 使用 'poke' 模式，周期性检查
    )

    # 原始任务 T1-T5
    # T1 = deletes all files from GCS
    t1 = BashOperator(
        task_id='delete_gcs_files',
        bash_command='echo "Deleting GCS files..."',
    )

    # T2 = Runs SQL query 1 and outputs to a table within BigQuery
    t2 = BashOperator(
        task_id='run_sql_query_1',
        bash_command='echo "Running SQL Query 1..."',
    )

    # T3 = Runs SQL query 2 and outputs to a table within BigQuery
    t3 = BashOperator(
        task_id='run_sql_query_2',
        bash_command='echo "Running SQL Query 2..."',
    )

    # T4 = Runs SQL query 3 and places a copy of this output as csv into the GCS that was emptied in T1
    t4 = BashOperator(
        task_id='run_sql_query_3_and_upload_to_gcs',
        bash_command='echo "Running SQL Query 3 and uploading to GCS..."',
    )

    # T5= Copies and Appends the reference numbers from the file in T4 to a history table in BigQuery.
    t5 = BashOperator(
        task_id='append_to_history_table',
        bash_command='echo "Appending to history table..."',
    )

    # 定义任务依赖关系
    # 只有当 check_last_tuesday 成功后，后续任务才会被执行
    check_last_tuesday >> t1 >> t2 >> t3 >> t4 >> t5