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

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

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


def push_data(**context):
    """
    向 xcom 推送数据的任务。

    **context：接收 Airflow 传入的上下文参数。
    其中包含 ti（TaskInstance），可以通过 ti 操作 xcom。
    """

    # 获取当前任务实例，并向 xcom 写入一条数据
    # key='my_key'：xcom 数据的键
    # value='my_value'：实际存储的数据
    context['ti'].xcom_push(
        key='my_key',
        value='my_value'
    )

    # 返回一个字符串
    # PythonOperator 默认会将函数返回值自动保存到 xcom
    # 默认的 xcom key 是 'return_value'
    return "默认的push"


def pull_data(**context):
    """
    从上游任务的 xcom 中读取数据。
    """

    # 获取当前任务实例（TaskInstance）
    ti = context['ti']

    # 从 push_task 任务中读取 key 为 my_key 的数据
    # 这里读取的是之前通过 xcom_push() 手动推送的数据
    value = ti.xcom_pull(
        task_ids='push_task',
        key='my_key'
    )

    # 从 push_task 任务中读取默认的返回值
    # PythonOperator 会将 push_data() 的返回值
    # 自动保存到 key='return_value' 对应的 xcom 中
    return_value = ti.xcom_pull(
        task_ids='push_task',
        key='return_value'
    )

    # 将获取到的数据输出到任务日志中，方便查看执行结果
    print(f'Pulled value: {value}')
    print(f'return_value: {return_value}')


# 定义一个 DAG，即 Airflow 工作流
with DAG(
    # DAG 的唯一标识符
    dag_id='xcom_example',

    # DAG 的开始日期
    # 注意：开始日期不代表 DAG 会在这一天自动执行
    start_date=datetime(2024, 10, 1),

    # 不设置定时调度，通常需要手动触发 DAG
    schedule_interval=None,

    # 不补执行历史上未运行的调度任务
    catchup=False,

    # 为 DAG 添加标签，方便在 Airflow 页面中筛选
    tags=["demo"]
) as dag:

    # 定义第一个任务：向 xcom 写入数据
    push_task = PythonOperator(
        # 任务的唯一标识符
        task_id='push_task',

        # 指定需要执行的 Python 函数
        python_callable=push_data
    )

    # 定义第二个任务：从 xcom 读取数据
    pull_task = PythonOperator(
        # 任务的唯一标识符
        task_id='pull_task',

        # 指定需要执行的 Python 函数
        python_callable=pull_data
    )

    # 设置任务依赖关系：
    # push_task 必须先成功完成，pull_task 才会执行
    push_task >> pull_task