Airflow - 跨任務的數據庫連接重用

Question

我正在嘗試使用 Airflow 將數據攝取到 Redshift。 我的源數據庫中有許多表，我在 Python 中編寫了模塊，這些模塊在 Redshift 中執行端到端 ETL。 我打算使用 Airflow 來協調它們的執行，其中每個表的 ETL 將是 DAG 中的一個任務。 他們主要擔心的是我想避免為我的每個表建立與 Redshift 的連接。 相反，我想在開始時建立一次連接，並希望為我的每個下游任務重新使用該連接。 我嘗試使用 xcom，但似乎連接對象不是可腌制的，因此不起作用（“連接”被腌制的概念也沒有多大意義）。 我也嘗試使用如下全局變量，但似乎它不起作用，因為它顯然為每個任務實例化了一個新的 object：

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

如果連接是指數據庫連接，那么不是。 每個 airflow 任務實例都在自己的進程中執行，因此您將無法重用相同的連接。

如果您想為多個操作重用相同的連接，則必須將它們組合成一個任務（例如在execute中，循環遍歷每個表並完成您的工作）。

應該不需要 jdbc。 您應該能夠使用PostgresHook進行紅移。

如果您的意思是airflow 連接（即Connection model，這是氣流管理信用的方式），那么，是的。 按照管理連接中的描述定義一次連接，並在每個任務定義中引用它。