循环调用psycopg2 COPY_FROM向空表导入多数据集仅最后一条生效问题

问题原因 #

1. 代码层面的直接错误：你将conn.close()放在了客户端循环内部，每处理完一个客户端的数据就直接关闭数据库连接，后续循环的所有数据库操作都是基于已关闭的连接执行，自然全部失败。只有最后一次循环的插入操作是在连接关闭前完成，因此始终只有最后一份数据插入成功。
2. 多进程场景的潜在错误：psycopg2的连接对象不支持跨进程共享，如果你是在多进程中传递同一个连接对象使用，会出现事务混乱、操作丢失等未定义行为，也会触发数据插入不全的问题。

修复方案 #

1. 调整连接关闭逻辑（单进程循环场景） #

将连接关闭操作移到所有数据处理完成之后，同时不要覆盖循环内的df变量，避免逻辑异常：

conn = psycopg2.connect(//你的数据库连接参数)
for df, client_id in clients_data:
    with conn:
          with conn.cursor() as cursor:
              for chunk in df:
                  cleaned_df = pipeline.clean_df(chunk, cursor, client_id)
                  copy_from_df(cleaned_df, 'mytable', cursor, chunksize=50000)
# 所有数据处理完成后再关闭连接
conn.close()

2. 多进程场景适配 #

按业务要求需要在子进程独立处理数据的，每个子进程必须单独创建、销毁数据库连接，禁止传递父进程生成的连接对象：

def process_single_client(args):
    df, client_id = args
    # 子进程独立创建连接
    conn = psycopg2.connect(//你的数据库连接参数)
    try:
        with conn:
            with conn.cursor() as cursor:
                for chunk in df:
                    cleaned_df = pipeline.clean_df(chunk, cursor, client_id)
                    copy_from_df(cleaned_df, 'mytable', cursor, chunksize=50000)
    finally:
        conn.close()

3. 辅助排查优化 #

如果调整后仍有问题，可以做以下修改定位原因：

• 在copy_from_df的cursor.copy_from执行后增加日志打印，确认每次插入的行数：

  logger.info(f"本次COPY插入行数：{cursor.rowcount}")
  

• 将copy_from_df中每次执行的SET search_path语句移到连接创建后统一执行，减少重复交互。
• 检查pipeline.clean_df方法是否存在依赖表中现有数据的分支逻辑，避免表为空时返回空数据集导致无数据插入。

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