Airflow - Database connection reuse across tasks

user_1924991 0
Viewed 1445

I am trying to use Airflow to ingest data into Redshift. There are a number of tables in my source database and I have written modules in Python which perform end-to-end ETL into Redshift. I was intending to use Airflow to coordinate their execution where ETL for each table will be a task in DAG. They key concern is that I wanted to avoid establishing connection to Redshift for each of my table. Rather, I want to establish connection once at the start and wanted to re-use that connection for each of my downstream tasks. I tried using xcom but seems like the connection objects are not pickle-able and thus didn't work (also the notion of a "connection" getting pickled does not make much sense). I also tried using global variables as follows but seems like it does not work as it is apparently instantiating a new object for each task:

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

is there any way I can achieve this? I don't want to setup a connection in Airflow (e.g. JdbcHook) and use that in this context.

airflow python database
1 Answers

If by connection you mean database connection, then no. Each airflow task instance is executed in its own process, so you will not be able to reuse the same connection.

If you want to reuse same connection for multiple operations, you'll have to combine them into a single task (e.g. in execute, loop through each table and do your work).

There should be no need for jdbc. You should be able to use PostgresHook for redshift.

If you mean airflow connection (i.e. the Connection model, which is airflow's way of managing creds), then, yes. Define your connection once as described in managing connections, and reference it in each task definition.

user_27616820
Related