I have a task in Airflow where I am loading data into a table. The process goes like this:
- query a database -> retrieve results as pandas data frame -> send the result set to multiple worker processes -> each worker process processes the rows and loads data into a different database.
Below is a simplified version of the DAG file:
import process
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.operators.python import PythonOperator
LOADING = PythonOperator(
task_id='LOADING',
python_callable=process,
op_kwargs={
'source_DB': MySqlHook(mysql_conn_id='source_DB'),
'destination_DB': MySqlHook(mysql_conn_id='destination_DB')
},
dag=dag,
)
start >> LOADING >> end
This is the code for the task:
import os
import logging
import billiard as mp
CUR_DIR = os.path.abspath(os.path.dirname(__file__))
def process(source_DB, destination_DB):
get_data = open(f"{CUR_DIR}/path/to/get_data.sql").read()
data = source_DB.get_pandas_df(
sql=get_data,
parameters={}
)
with mp.Pool(processes=mp.cpu_count(), initializer=init_worker, initargs=(destination_DB,)) as pool:
items = [(idx, row) for idx, row in data.iterrows()]
pool.map(load_data, items)
def init_worker(destination_DB):
global conn
conn = destination_DB.get_conn()
def load_data(args):
index, data = args
insert_sql = open(
f"{CUR_DIR}/path/to/insert.sql").read()
conn.autocommit(True)
destination_DB_cur = conn.cursor()
params = {
'para1': data['para1'],
'para2': data['para2']
}
for word, replacement in params.items():
insert_sql = insert_sql.replace(
'{{' + str(word) + '}}', str(replacement))
try:
destination_DB_cur.execute(insert_sql)
except Exception as e:
print(e)
destination_DB_cur.close()
The job runs without any errors, but sometimes the loaded data gets duplicated three times. After doing some research, I found conflicting advice - some attribute it to the billiard library while others suggest using connection pooling for synchronization and coordination. Can anyone provide insight into this issue and suggest how to prevent the duplication from happening?