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

# 导入 Variable 类，用于获取 Airflow 中保存的全局变量
from airflow.models import Variable

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

# 导入 PythonOperator，用于执行 Python 函数
from airflow.operators.python import PythonOperator

# 导入 datetime，用于设置 DAG 的开始日期
from datetime import datetime


# ==================================================
# 第一部分：通过运行上下文获取城市参数
# ==================================================

# **context 表示接收多个关键字参数
# Airflow 执行任务时，会把运行上下文传递给这个函数
def print_city(**context):

    # 从上下文中获取当前 DAG Run（一次工作流运行）的配置信息
    # dag_run.conf 通常用于获取手动触发 DAG 时传入的 JSON 参数
    # 例如：{"city": "Guangzhou"}

    # 如果没有传入配置，就使用一个空字典，避免 conf 为 None
    conf = context["dag_run"].conf or {}

    # 从配置字典中获取 city 对应的值
    # 如果没有传入 city，就使用默认值 default_city
    city = conf.get("city", "default_city")

    # 将获取到的城市打印到任务日志中
    print(f"The city is: {city}")


# ==================================================
# 第二部分：获取 Airflow Variable 全局变量
# ==================================================

def use_variable():

    # 从 Airflow 的 Variables 中读取名为 foo 的变量
    # 如果 Airflow 中不存在 foo，就使用 default_value
    foo = Variable.get("foo", default_var="default_value")

    # 打印变量的值
    print(f"The value of foo is: {foo}")


# ==================================================
# 第三部分：创建 DAG，也就是工作流
# ==================================================

with DAG(

    # DAG 的唯一标识名称
    dag_id="args_example",

    # 每天运行一个调度周期
    schedule_interval="@daily",

    # 设置调度开始日期
    start_date=datetime(2023, 1, 1),

    # 不补跑过去未执行的调度周期
    catchup=False,

    # 设置标签，方便在 Airflow 网页中筛选
    tags=["demo"]

) as dag:

    # ==================================================
    # 第四部分：任务一——通过 context 获取城市参数
    # ==================================================

    task = PythonOperator(

        # 任务的唯一标识
        task_id="print_city_task",

        # 指定任务要调用的 Python 函数
        # 注意：传入函数名，不要写 print_city()
        python_callable=print_city

    )

    # ==================================================
    # 第五部分：任务二——通过 Jinja 模板获取城市参数
    # ==================================================

    task_bash = BashOperator(

        # 任务的唯一标识
        task_id="print_city_bash",

        # BashOperator 会在执行前解析 Jinja 模板
        # dag_run.conf 用来获取本次 DAG Run 的配置
        # 如果没有传入 city，就使用 default_city
        bash_command=(
            'echo "city is '
            '{{ dag_run.conf.get("city", "default_city") '
            'if dag_run and dag_run.conf else "default_city" }}"'
        )

    )

    # ==================================================
    # 第六部分：任务三——获取 Airflow 全局变量
    # ==================================================

    # 全局变量（Variable）
    task_var = PythonOperator(

        # 任务的唯一标识
        task_id="use_variable_task",

        # 调用前面定义的 use_variable 函数
        python_callable=use_variable

    )

    # ==================================================
    # 第七部分：任务四——使用宏变量
    # ==================================================

    # 官方文档参考：
    # https://airflow.apache.org/docs/apache-airflow/2.2.5/templates-ref.html

    # 宏变量可以在支持模板的 Operator 参数中使用
    # 这里使用了 ds 和 dag.dag_id 两个模板变量
    task_ds = BashOperator(

        # 任务的唯一标识
        task_id="print_execution_date",

        # {{ ds }}：本次调度运行对应的日期，格式通常为 YYYY-MM-DD
        # {{ dag.dag_id }}：当前 DAG 的唯一标识名称
        bash_command=(
            'echo "Execution date is {{ ds }}, '
            'dag_id is {{ dag.dag_id }}"'
        )

    )

    # ==================================================
    # 第八部分：定义任务之间的依赖关系
    # ==================================================

    # task 和 task_bash 是两个并行的上游任务
    # 两者都成功后，task_var 才可以执行
    # task_var 成功后，task_ds 才可以执行
    [task, task_bash] >> task_var >> task_ds