100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
# 导入 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 |