psycopg2上传数据时连接中断，如何重连并断点续传？

你的原代码直接用executemany一次性提交所有10000行数据，一旦插入过程中连接中断，整个批次会因为未提交而全部回滚，而且没有记录已插入的位置，所以无法直接从5001行继续。但通过分片批量插入+断点记录的方式，就能实现重连后从断点处续传。

改造思路 #

• 分片插入：把数据分成小批次（比如每1000行一批），每批插入后立即提交，确保成功的批次不会因为后续中断丢失
• 断点记录：实时记录已成功插入的行数，中断后重连时从该断点开始处理剩余数据
• 重连逻辑：捕获连接异常时，自动重新建立数据库连接，然后从断点处继续执行

修改后的完整代码 #

import psycopg2
from psycopg2 import Error
import os
import json

def insert_hp_csv_data_to_pgdb(df, table, batch_size=1000):
    # 提前完成表检查、列映射和SQL模板生成（避免重连后重复执行）
    try:
        # 先做一次表存在检查
        temp_conn = psycopg2.connect(
            host="IP",
            database="database name",
            user="user",
            password="admin",
            port=5432
        )
        temp_cur = temp_conn.cursor()
        query = "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = %s)"
        temp_cur.execute(query, (table,))
        table_exists = temp_cur.fetchone()[0]
        temp_cur.close()
        temp_conn.close()

        if not table_exists:
            print(f"表 '{table}' 不存在，请先在PostgreSQL中创建对应表，确保列名映射匹配")
            raise Exception("目标表不存在")

        # 生成PostgreSQL列名映射
        postgresql_columns = []
        column_map_path = os.path.join(os.getcwd(), f'column_map/{table}.json') 
        with open(column_map_path, 'r') as file:
            col_map = json.load(file)
        for col in df.columns:
            postgresql_column_name = col_map.get(col, col)
            postgresql_columns.append(postgresql_column_name)
        
        columns = ', '.join(postgresql_columns)
        placeholders = ', '.join(['%s'] * len(df.columns))
        insert_query = f'''INSERT INTO analytical_db.{table}({columns}) VALUES ({placeholders})'''

        # 初始化断点
        inserted_rows = 0
        total_rows = len(df)

        while inserted_rows < total_rows:
            connection = None
            cur = None
            try:
                # 建立连接
                connection = psycopg2.connect(
                    host="IP",
                    database="database name",
                    user="user",
                    password="admin",
                    port=5432
                )
                cur = connection.cursor()

                # 计算当前批次的起始和结束索引
                start_idx = inserted_rows
                end_idx = min(inserted_rows + batch_size, total_rows)
                batch_df = df.iloc[start_idx:end_idx]
                data_values = [tuple(row) for _, row in batch_df.iterrows()]

                # 插入当前批次
                cur.executemany(insert_query, data_values)
                connection.commit()

                # 更新断点
                inserted_rows = end_idx
                print(f"已成功插入 {inserted_rows}/{total_rows} 行数据")

            except (Error, Exception) as error:
                print(f"插入数据时出错: {error}")
                print(f"当前断点位置：已插入 {inserted_rows} 行，将尝试重连后从该位置继续")
                # 回滚当前批次（如果未提交）
                if connection:
                    connection.rollback()
                # 等待1秒后重试（可选，避免频繁重试）
                import time
                time.sleep(1)
            finally:
                # 关闭当前连接
                if cur:
                    cur.close()
                if connection:
                    connection.close()
                    print("PostgreSQL连接已关闭")

        print(f"所有数据从 {table} 上传成功，共插入 {total_rows} 行")

    except (Exception, Error) as error:
        print("初始化过程出错", error)

# 使用示例
# insert_hp_csv_data_to_pgdb(your_dataframe, "your_table_name", batch_size=1000)

关键修改说明 #

• 提前初始化：表存在检查、列映射和SQL模板只执行一次，避免重连后重复操作，提升效率
• 分片处理：按batch_size拆分数据，每批插入后立即提交，确保成功的批次不会丢失
• 断点维护：用inserted_rows记录已成功插入的行数，中断后从该位置继续
• 异常重试：捕获异常后回滚当前批次，等待片刻后重新建立连接，继续执行剩余数据
• 资源清理：每次循环结束后关闭连接，避免资源泄漏

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