我正在尝试使用Airflow将数据导入Redshift。我的源数据库中有许多表格,我编写了Python模块,可实现端到端的ETL进入Redshift。我打算使用Airflow协调它们的执行,其中每个表格的ETL将是DAG中的任务。重要的是,我希望避免为每个表格建立Redshift连接。相反,我想在开始时建立一次连接,并希望重复使用该连接来处理每个下游任务。
我尝试使用xcom,但似乎连接对象不可pickle-able,因此无法使用(而且pickle化“connection”的概念并不太合理)。我还尝试了以下方法使用全局变量,但似乎并不起作用,因为它显然为每个任务实例化一个新对象:
redcon = None
red_iam_role = None
dag = DAG('redshift_etl', description='An example redshift etl prototype', schedule_interval='0 12 * * *', start_date=datetime(2017, 3, 20), catchup=False)
def connect_to_redshift():
global redcon, red_iam_role
if redcon is None or red_iam_role is None:
(redcon, red_iam_role) = redshift_utilities.get_connection('redshift_cluster', 15) #a module that returns a tuple consisting of redshift connection object (using psycopg2) and redshift_iam_role which I use for a few tasks.
return {"connection": redcon, "iam_role": red_iam_role}
else:
return {"connection": redcon, "iam_role": red_iam_role}
def first_task(redshift_params,**context):
print(redshift_params["connection"])
def second_task(redshift_params,**context):
print(redshift_params["connection"])
first_task_dag = PythonOperator(task_id='first_task',provide_context=True,op_kwargs={"redshift_params":connect_to_redshift()}, python_callable=first_task, dag=dag)
second_task_dag = PythonOperator(task_id='second_task',provide_context=True,op_kwargs={"redshift_params":connect_to_redshift()}, python_callable=second_task, dag=dag)
first_task_dag >> second_task_dag
有没有办法实现这个?我不想在Airflow中设置连接(例如JdbcHook),并在此上下文中使用它。