import datetime
import pendulum

from airflow import DAG
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.operators.bash import BashOperator
from airflow.models import DagRun


# 自定义函数：获取上游 DAG 最近一次查询结果中的 execution_date
# execution_date_fn 会调用这个函数，用于确定下游需要等待哪个上游执行日期
def getExecutionDate(kwargs=None):
    # DagRun.find() 根据 dag_id 查询 dag_dependency1 的 DAG 运行记录
    # pop() 取出查询结果列表中的最后一个元素，并获取该次运行的 execution_date
    # 注意：这里取的是查询结果列表的最后一个元素，不一定能保证它就是时间上最新的一次运行
    return DagRun.find(dag_id="dag_dependency1").pop().execution_date


# 创建下游 DAG：dag_dependency2
# 适用于需要等待上游 DAG 执行完成后，再执行下游任务的场景
dag = DAG(
    # 当前 DAG 的唯一标识
    dag_id='dag_dependency2',

    # DAG 描述信息，可在 Airflow Web UI 中查看
    description='this is a hello world workflow',

    # 不设置定时调度，None 表示该 DAG 不会按照固定时间自动调度
    # 可以通过 Airflow Web UI、命令行或其他 DAG 手动触发
    schedule_interval=None,

    # DAG 的开始时间，使用 UTC 时区
    # 它是调度相关的起始时间，不代表 DAG 一定会在这个时间自动运行
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),

    # 关闭历史任务补跑
    catchup=False,

    # 单次 DAG Run 的最长运行时间为 60 分钟
    dagrun_timeout=datetime.timedelta(minutes=60),

    # 添加标签，方便在 Airflow Web UI 中分类和筛选 DAG
    tags=['demo']
)


# 创建外部任务传感器
# ExternalTaskSensor 用于等待另一个 DAG 或其中的某个任务达到指定状态
# 当前配置监控上游 DAG：dag_dependency1
external_task = ExternalTaskSensor(
    # 当前传感器任务的名称
    task_id='external_task',

    # 指定需要监控的上游 DAG 的 dag_id
    external_dag_id="dag_dependency1",

    # 设置为 None，表示监控整个上游 DAG 的状态
    # 即等待上游 DAG Run 达到允许的完成状态，而不是只等待某一个指定任务
    external_task_id=None,

    # 指定需要检查的上游 DAG 执行日期
    # 默认情况下，可通过 execution_delta 指定相对当前执行日期的偏移量
    #
    # 例如，execution_delta=datetime.timedelta(minutes=1)
    # 表示检查当前执行日期之前 1 分钟对应的上游 DAG Run
    #
    # 注意：execution_delta 通常用于存在固定时间间隔的上、下游调度场景
    # 当前代码注释掉该参数，因此不会使用这个时间偏移量
    # execution_delta=datetime.timedelta(minutes=1),

    # 自定义执行日期计算函数
    # 由 getExecutionDate() 返回需要监控的上游 DAG 执行日期
    # 当前函数直接查询 dag_dependency1 的运行记录，并取查询结果的最后一条
    # 因此这里不是根据当前下游执行日期计算偏移量
    execution_date_fn=getExecutionDate,

    # 每隔 5 秒检查一次上游 DAG 是否达到要求的状态
    poke_interval=5,

    # 最多等待 180 秒
    # 如果超过该时间上游 DAG 仍未达到要求的状态，传感器任务将超时失败
    timeout=180,

    # 将当前传感器任务关联到 dag_dependency2
    dag=dag
)


# 创建下游 Bash 任务
# 只有当 external_task 传感器任务成功后，该任务才会开始执行
next_task = BashOperator(
    # 当前任务的唯一标识
    task_id='next_task',

    # 执行一条 Shell 命令，在终端输出提示信息
    bash_command='echo "last dag is completed"',

    # 将当前任务关联到 dag_dependency2
    dag=dag
)


# 设置任务依赖关系
# external_task 成功完成后，才会执行 next_task
#
# 执行流程：
# 等待 dag_dependency1 完成
#             ↓
#       external_task
#             ↓
#         next_task
external_task >> next_task