Airflow - 跨任务的数据库连接重用

Question

我正在尝试使用 Airflow 将数据摄取到 Redshift。我的源数据库中有许多表，我已经用 Python 编写了模块，这些模块在 Redshift 中执行端到端 ETL。我打算使用 Airflow 来协调它们的执行，其中每个表的 ETL 将是 DAG 中的一项任务。他们主要担心的是我想避免为我的每个表建立与 Redshift 的连接。相反，我想在开始时建立一次连接，并希望为我的每个下游任务重新使用该连接。 我尝试使用 xcom，但似乎连接对象不是可腌制的，因此不起作用（“连接”被腌制的概念也没有多大意义）。我也尝试如下使用全局变量，但似乎它不起作用，因为它显然为每个任务实例化了一个新对象：

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）并在这种情况下使用它。

Answer 1

如果连接是指数据库连接，那么不是。每个气流任务实例都在自己的进程中执行，因此您将无法重用相同的连接。

如果您想为多个操作重复使用相同的连接，则必须将它们组合成一个任务（例如，在 execute 中，遍历每个表并完成您的工作）。

应该不需要 jdbc。您应该可以使用PostgresHook 进行红移。

如果您的意思是气流连接（即Connection 模型，这是气流管理信用的方式），那么，是的。按照managing connections 中的说明定义一次连接，然后在每个任务定义中引用它。