# 导入 pendulum，用于创建带时区的日期时间对象
import pendulum

# 导入 DAG，用于定义工作流
from airflow import DAG

# 导入 BashOperator，用于执行 Bash 命令
from airflow.operators.bash import BashOperator

# 导入 TriggerDagRunOperator，用于触发其他 DAG 运行
from airflow.operators.trigger_dagrun import TriggerDagRunOperator

# 导入 Python 的 datetime 模块
import datetime

# 导入 DagRun 模型，用于查询 DAG 的运行记录
from airflow.models.dagrun import DagRun


# 获取指定 DAG 最近一次运行的执行时间
def getExecutionDate(kwargs=None):
    # 查询 dag_id 为 dag_dependency1 的所有 DAG 运行记录
    # find() 返回 DagRun 对象列表
    # pop() 取出列表中的最后一个元素
    # execution_date 获取该次运行对应的逻辑执行时间
    return DagRun.find(dag_id="dag_dependency1").pop().execution_date


# 定义一个 DAG，即一个工作流
with DAG(
    # DAG 的唯一标识，其他 DAG 可以通过这个 ID 引用它
    dag_id='dag_dependency1',

    # DAG 的描述信息
    description="this is a hello world workflow",

    # 不设置定时调度，需要手动触发或由其他 DAG 触发
    schedule_interval=None,

    # 设置 DAG 的开始时间，并指定 UTC 时区
    # 注意：start_date 是调度的起始时间，不代表 DAG 会在这个时间自动运行
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),

    # 关闭历史任务补跑
    # 对于配置了定时调度的 DAG，catchup=False 表示不会补跑所有历史调度周期
    catchup=False,

    # 添加标签，便于在 Airflow Web UI 中筛选 DAG
    tags=['demo'],

) as dag:

    # -------------------- 第一个任务 --------------------
    # 执行一条 Bash 命令，在日志中输出提示信息
    first_task = BashOperator(
        # 任务的唯一标识
        task_id="first_task",

        # 需要执行的 Bash 命令
        bash_command="echo 'this is first task'",

        # 指定该任务所属的 DAG
        dag=dag
    )

    # -------------------- 触发第一个子 DAG --------------------
    # 当该任务执行时，触发 args_example 这个 DAG
    trigger_task = TriggerDagRunOperator(
        # 当前触发任务的 ID
        task_id="triger_task",

        # 要触发的目标 DAG 的 dag_id
        trigger_dag_id="args_example",

        # 传递给子 DAG 的配置参数
        # 子 DAG 可以通过 dag_run.conf 获取这些参数
        # 例如：dag_run.conf.get("city") 将得到 "beijing"
        conf={"city": "beijing"},

        # 是否等待子 DAG 执行结束
        # True 表示等待子 DAG 完成后，当前触发任务才结束
        wait_for_completion=True,

        # 每隔 5 秒检查一次子 DAG 的运行状态
        # 该参数在等待子 DAG 完成时生效
        poke_interval=5
    )

    # -------------------- 触发第二个子 DAG --------------------
    # 当该任务执行时，触发 branch_example 这个 DAG
    trigger_task2 = TriggerDagRunOperator(
        # 当前触发任务的 ID
        task_id="triger_task2",

        # 要触发的目标 DAG 的 dag_id
        trigger_dag_id="branch_example",

        # 传递给子 DAG 的配置参数
        # 子 DAG 可以通过 dag_run.conf 获取 city 参数
        conf={"city": "beijing"},

        # 等待子 DAG 执行结束后，再将当前任务标记为完成
        wait_for_completion=True,

        # 每隔 5 秒检查一次子 DAG 的执行状态
        poke_interval=5
    )

    # -------------------- 最后一个任务 --------------------
    # 当前面两个触发任务都成功完成后，执行该任务
    last_task = BashOperator(
        # 任务的唯一标识
        task_id="last_task",

        # 执行 Bash 命令，输出提示信息
        bash_command="echo 'this is last task'",

        # 指定该任务所属的 DAG
        dag=dag
    )

    # -------------------- 设置任务依赖关系 --------------------

    # 并行触发模式
    # 1. first_task 先执行
    # 2. first_task 成功后，trigger_task 和 trigger_task2 可以并行执行
    # 3. 两个触发任务都成功完成后，last_task 才执行
    first_task >> [trigger_task, trigger_task2] >> last_task

    # -------------------- 串行触发模式 --------------------
    # 如果希望两个子 DAG 按顺序触发，可以使用下面的依赖关系
    # 执行顺序：first_task
    #          -> trigger_task
    #          -> trigger_task2
    #          -> last_task
    #
    # 因为 wait_for_completion=True，
    # trigger_task 会等待 args_example 执行完成，
    # 然后才会执行 trigger_task2。
    #
    # first_task >> trigger_task >> trigger_task2 >> last_task