使用Billiard多进程向MySQL插入数据时数据重复三次的问题求助

任务流程 #

• 查询数据库 → 将结果存入pandas DataFrame → 将结果集传递给工作进程 → 每个工作进程处理数据行并加载到另一个数据库

DAG文件简化版本 #

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

任务代码实现 #

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()

问题现象 #

任务运行无报错，但有时加载的数据会重复三次

原因分析 #

1. MySqlHook序列化异常：MySqlHook对象无法被billiard进程池正确序列化传递给子进程。init_worker中传递的destination_DB在跨进程复制时，可能触发重复初始化，导致子进程创建多个无效连接，进而引发重复插入。
2. 全局连接变量的进程共享风险：子进程通过global conn复用连接，但billiard的写时复制机制可能导致父进程连接状态被意外复制，多个子进程的连接操作出现异常，重复提交数据。
3. 字符串替换的SQL逻辑隐患：直接通过replace修改SQL字符串的方式，若进程调度出现重复触发（如进程池任务重发），容易导致同一行数据被多次处理插入。
4. 进程池任务分发的不确定性：billiard进程池在某些场景下可能存在任务重复分发的情况，尤其当父进程资源未正确隔离时，会导致同一数据行被多个进程重复处理。

解决措施 #

1. 避免传递不可序列化对象到子进程 #

不要直接传递MySqlHook给子进程，改为传递连接ID，在子进程内重新初始化Hook：

# 修改process函数
def process(source_DB, dest_conn_id):
    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=(dest_conn_id,)) as pool:
        items = [(idx, row) for idx, row in data.iterrows()]
        pool.map(load_data, items)

# 修改init_worker
def init_worker(dest_conn_id):
    global conn
    from airflow.providers.mysql.hooks.mysql import MySqlHook
    dest_hook = MySqlHook(mysql_conn_id=dest_conn_id)
    conn = dest_hook.get_conn()

同步更新DAG的op_kwargs：

LOADING = PythonOperator(
            task_id='LOADING',
            python_callable=process,
            op_kwargs={
                'source_DB': MySqlHook(mysql_conn_id='source_DB'),
                'dest_conn_id': 'destination_DB'
            },
            dag=dag,
        )

2. 改用参数化查询替代字符串替换 #

字符串替换不仅存在SQL注入风险，还容易引发逻辑错误，直接使用数据库参数化查询：

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']
    }
    # 用参数化执行替代字符串替换
    try:
        destination_DB_cur.execute(insert_sql, params)
    except Exception as e:
        print(e)
    destination_DB_cur.close()

注意：需要将insert.sql中的{{para1}}等占位符改为对应数据库的参数格式，如MySQL的%s或%(para1)s

3. 替换进程池为线程池（IO密集型场景） #

数据加载属于IO密集型任务，线程池无需处理进程序列化问题，效率更高：

from concurrent.futures import ThreadPoolExecutor
import os

def process(source_DB, dest_conn_id):
    get_data = open(f"{CUR_DIR}/path/to/get_data.sql").read()
    data = source_DB.get_pandas_df(sql=get_data, parameters={})

    def load_single_row(item):
        index, row = item
        dest_hook = MySqlHook(mysql_conn_id=dest_conn_id)
        insert_sql = open(f"{CUR_DIR}/path/to/insert.sql").read()
        with dest_hook.get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute(insert_sql, {'para1': row['para1'], 'para2': row['para2']})
            conn.commit()

    with ThreadPoolExecutor(max_workers=os.cpu_count()*2) as executor:
        executor.map(load_single_row, data.iterrows())

4. 数据库层面添加唯一性约束 #

在目标表添加唯一索引，从底层阻止重复数据写入：

ALTER TABLE target_table ADD UNIQUE KEY unique_para (para1, para2);

内容的提问来源于stack exchange，提问作者moe_