使用一个Python任务的输出并将其用作Airflow上另一个Python任务的输入

时间:2019-09-09 21:09:54

标签: python mongodb airflow pipeline

因此,我正在使用Apache Airflow创建一个数据流,以获取存储在Pandas Dataframe中的某些数据,然后将其存储到MongoDB中。因此,我有两种python方法,一种用于获取数据并返回数据帧,另一种用于将其存储到相关数据库中。如何获取一个任务的输出并将其作为输入传递给另一任务?这就是我到目前为止(概述和精简版)

我研究了xcom pull and push的概念,这就是我在下面实现的概念,我还看到有用于气流的MongoHook,但不确定如何使用它。

import pandas as pd
import pymongo
import airflow
from datetime import datetime, timedelta
from airflow.models import DAG
from airflow.operators.python_operator import PythonOperator


def get_data(name, **context):
    data = pd.read_csv('dataset.csv')
    df = data.loc[data.name == name]
    context['ti'].xcom_push(task_ids=['get-data'], value=data)

def push_to_db(df, dbname, collection):
    client = pymongo.MongoClient(-insert creds here-)
    db = client[dbname][collection]
    data = df.to_dict(orient='records')
    db.insert_many(data)

args = {
    'owner': 'Airflow',
    'start_date': airflow.utils.dates.days_ago(2),
}

dag = DAG(
  dag_id='simple_xcom',
  default_args=args,
  start_date=datetime(2019, 09, 02),
  schedule_interval="@daily",
  retries=2
)

task1 = PythonOperator(task_id='get-data', params=['name': 'John'], 
        python_callable=get_data, 
        provide_context=True, dag=dag)

task2 = PythonOperator(task_id='load-db', params=['df': context['ti'].xcom_pull(task_ids=['get-data'], key='data'), 
    'dbname': 'person', 'table': 'salary'), 
    python_callable=push_to_db, provide_context=True, dag=dag) 

task1 >> task2 


每次我尝试运行它时,它都会显示该上下文不存在。那么也许在将一项任务的输出作为另一项任务的输入方面做错了吗?

1 个答案:

答案 0 :(得分:1)

看看示例xcom DAG。

https://github.com/apache/airflow/blob/master/airflow/example_dags/example_xcom.py