如何优化psycopg2跨库数据迁移速度？支持批量处理吗？

你的代码速度慢的核心原因是逐行处理+频繁提交：每处理1行就执行一次COPY和COMMIT，这两个操作本身涉及磁盘IO和网络交互，小批量重复执行会导致开销被放大上万倍。以下是具体优化方案，完全支持批量处理多行数据：

一、核心优化：批量读取+批量写入 #

将原逐行处理改为批量读取多行，一次性写入目标库，大幅减少COPY和COMMIT的次数。

优化后的代码示例 #

import io
import csv
import psycopg2

# 调整批量参数：根据内存和表结构设置，建议10000-50000
batch_size = 10000
transfer_table_name = 'my_table'
# 每处理N批提交一次，避免频繁磁盘IO
commit_interval = 10

# 假设已建立db1（源库）和db2（目标库）的连接，且db2连接autocommit=False
with db1.cursor(name='my_cursor') as src_cursor:
    # 设置服务器端游标每次拉取的行数，和batch_size一致
    src_cursor.itersize = batch_size
    src_cursor.execute('SELECT * FROM schema.table')
    
    dest_cursor = db2.cursor()
    total_processed = 0
    
    while True:
        # 批量读取源库数据
        rows = src_cursor.fetchmany(batch_size)
        if not rows:
            break
        
        # 批量写入StringIO
        csv_io = io.StringIO()
        writer = csv.writer(csv_io)
        writer.writerows(rows)  # 一次性写入多行，替代逐行writerow
        
        csv_io.seek(0)
        # 批量导入目标库
        dest_cursor.copy_expert(
            f'COPY {transfer_table_name} FROM STDIN WITH (NULL \'\' , DELIMITER \',\', FORMAT CSV)',
            csv_io
        )
        
        total_processed += len(rows)
        # 按间隔提交事务
        if total_processed % (batch_size * commit_interval) == 0:
            db2.commit()
            print(f'已提交 {total_processed} 行数据')
    
    # 提交最后一批剩余数据
    db2.commit()
    dest_cursor.close()
    print(f'迁移完成，共处理 {total_processed} 行')

二、额外优化手段，进一步提速 #

1. 临时禁用目标表的索引和约束 #

如果目标表有索引、外键或触发器，迁移时这些结构会导致每批数据插入都要更新索引、检查约束，大幅拖慢速度。可以先禁用，迁移完成后再恢复：

-- 迁移前执行（db2中）
ALTER TABLE {transfer_table_name} DISABLE TRIGGER ALL;
DROP INDEX IF EXISTS idx_your_index_name;
-- 迁移完成后执行
ALTER TABLE {transfer_table_name} ENABLE TRIGGER ALL;
CREATE INDEX idx_your_index_name ON {transfer_table_name}(your_column);

2. 使用PostgreSQL原生工具替代Python中转 #

如果允许直接操作数据库，用pg_dump+pg_restore或者dblink跨库拉取数据，效率远高于Python代码：

• pg_dump/pg_restore：适合全表或整库迁移，命令示例：

  pg_dump -h src_host -U src_user -d db1 -t schema.table > table_dump.sql
  psql -h dest_host -U dest_user -d db2 < table_dump.sql
  

• dblink跨库插入：在目标库直接拉取源库数据，无需客户端中转：

  -- 在db2中执行
  CREATE EXTENSION IF NOT EXISTS dblink;
  SELECT dblink_connect('src_db', 'dbname=db1 user=src_user password=src_pwd host=src_host');
  INSERT INTO {transfer_table_name} SELECT * FROM dblink('src_db', 'SELECT * FROM schema.table') AS t(col1 INT, col2 VARCHAR, ...);
  SELECT dblink_disconnect('src_db');
  

3. 调整连接参数 #

确保源库和目标库的连接使用最优配置，比如在psycopg2连接时添加：

conn = psycopg2.connect(
    dbname='xxx', user='xxx', password='xxx', host='xxx',
    keepalives=1, keepalives_idle=30, keepalives_interval=10, keepalives_count=5
)

内容的提问来源于stack exchange，提问作者Иван Иваныч