如何将psycopg2 ConnectionPool传入ProcessPoolExecutor？CSV入PG报错求解

问题描述 #

我有一批CSV文件需要导入PostgreSQL表，尝试使用ProcessPoolExecutor并将psycopg2的ConnectionPool对象传入load_csv_file函数时，提交函数出现TypeError，等待过程中同样触发该错误且程序挂起。求问是否有其他开发者遇到过相同问题及解决方法，同时希望了解更高效的CSV数据加载方案。

相关代码 #

def load_csv_file(file_path, connection, semaphore):
    with open(file_path, "r") as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            try:
                semaphore.acquire()
                connection.cursor().execute(
                    """
                    INSERT INTO table (column1, column2, column3)
                    VALUES (%s, %s, %s)
                    """,
                    row,
                )
            finally:
                semaphore.release()
def main():
    # Create a connection pool
    connection_pool = ThreadedConnectionPool(
        max_conn=10,
        host="localhost",
        port=5432,
        dbname="my_database",
        user="my_user",
        password="my_password",
    )

with ProcessPoolExecutor(max_workers=10) as executor:
        # Load the CSV files to the table
        for file_path in ["file1.csv", "file2.csv", "file3.csv"]:
            executor.submit(
                load_csv_file,
                file_path,
                connection_pool.getconn(),
                semaphore=semaphore,
            )

问题原因与解决方法 #

1. TypeError与程序挂起的根源 #

ProcessPoolExecutor基于多进程实现，而ThreadedConnectionPool的数据库连接是线程绑定的，无法跨进程序列化传递。你直接在主进程获取连接后传给子进程，会因为连接对象无法被pickle序列化触发TypeError，同时进程间的资源竞争会导致程序挂起。

2. 可行修复方案 #

方案一：替换为ThreadPoolExecutor（推荐，IO密集场景更合适）

CSV导入属于IO密集型任务，线程池比进程池更高效，且ThreadedConnectionPool本身就是为线程场景设计的，不会有跨进程序列化问题。修正后的代码：

import csv
from psycopg2.pool import ThreadedConnectionPool
from concurrent.futures import ThreadPoolExecutor
import threading

def load_csv_file(file_path, connection_pool, semaphore):
    conn = None
    cursor = None
    try:
        semaphore.acquire()
        # 每个线程单独从连接池获取连接
        conn = connection_pool.getconn()
        cursor = conn.cursor()
        with open(file_path, "r") as csvfile:
            reader = csv.reader(csvfile)
            # 跳过表头（如果CSV有表头的话）
            next(reader)
            for row in reader:
                cursor.execute(
                    """
                    INSERT INTO table (column1, column2, column3)
                    VALUES (%s, %s, %s)
                    """,
                    row,
                )
        # 批量提交事务，提升效率
        conn.commit()
    except Exception as e:
        if conn:
            conn.rollback()
        print(f"处理文件{file_path}出错: {str(e)}")
    finally:
        # 关闭游标，归还连接到池
        if cursor:
            cursor.close()
        if conn:
            connection_pool.putconn(conn)
        semaphore.release()

def main():
    # 注意ThreadedConnectionPool需要指定minconn和maxconn参数
    connection_pool = ThreadedConnectionPool(
        minconn=2,
        maxconn=10,
        host="localhost",
        port=5432,
        dbname="my_database",
        user="my_user",
        password="my_password",
    )
    # 用信号量控制并发数，避免连接池耗尽
    semaphore = threading.Semaphore(10)
    csv_files = ["file1.csv", "file2.csv", "file3.csv"]
    
    with ThreadPoolExecutor(max_workers=10) as executor:
        for file in csv_files:
            executor.submit(load_csv_file, file, connection_pool, semaphore)
    
    # 所有任务完成后关闭连接池
    connection_pool.closeall()

if __name__ == "__main__":
    main()

方案二：多进程场景下独立创建连接

如果一定要用ProcessPoolExecutor，需要在每个子进程内部单独创建数据库连接（不要共享主进程的连接池），因为多进程之间内存隔离，无法共享连接对象。但这种方式开销较大，不推荐IO密集场景使用。

更高效的CSV加载方案 #

逐行INSERT效率极低，推荐使用PostgreSQL原生的COPY命令，这是批量导入CSV的最优方案，速度比逐行INSERT快10~100倍。结合psycopg2的copy_from方法实现：

基于COPY命令的批量导入代码 #

def load_csv_with_copy(file_path, connection_pool, semaphore):
    conn = None
    cursor = None
    try:
        semaphore.acquire()
        conn = connection_pool.getconn()
        cursor = conn.cursor()
        with open(file_path, "r") as csvfile:
            # 跳过CSV表头
            next(csvfile)
            # 使用copy_from批量导入
            cursor.copy_from(
                csvfile,
                "table",  # 目标表名
                columns=("column1", "column2", "column3"),  # 对应CSV列的表字段
                sep=","  # CSV分隔符，根据实际情况调整
            )
        conn.commit()
    except Exception as e:
        if conn:
            conn.rollback()
        print(f"处理文件{file_path}出错: {str(e)}")
    finally:
        if cursor:
            cursor.close()
        if conn:
            connection_pool.putconn(conn)
        semaphore.release()

额外优化建议 #

• 导入前临时禁用表的索引和外键约束，导入完成后再重建，能大幅提升导入速度
• 对于超大CSV文件，可以拆分后并行导入，但注意不要超过数据库的最大连接数限制
• 调整连接池的minconn和maxconn参数，匹配线程池的并发数，避免连接等待

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