AsyncPG COPY 批量装载

示例说明

逐行 INSERT 的往返开销在大批量面前不可接受，COPY 把整批数据编码成一条协议消息传输，速度差一到两个数量级。异步生成器作为数据源让装载与生产并发流水线化——外部数据分页读出时不必先全部驻留内存。临时表暂存后再原子切换，业务表不被半成品污染。

技术要点

• COPY 单协议消息替代逐行往返
• 异步生成器边产边装零全量驻留
• 临时表隔离装载与业务可见性

代码解读

generate_rows 模拟两万五千行的异步供给；bulk_copy 在事务里先建临时表再 copy_to_table，生成器直接作为 source 传入；事务提交后临时表在会话内可见，取计数验证装载量，finally 里关闭连接释放服务端资源。

使用建议

临时表是会话级的，切正式表要在同一连接内完成。错误行处理可先 COPY 到带宽松约束的暂存再清洗。大装载注意目标库的 WAL 与磁盘压力。

代码示例

import asyncio

import asyncpg


async def generate_rows(n: int):
    for i in range(n):
        # 生产环境中这里通常是异步分页读取外部数据源
        yield (i, f'user-{i}', i % 2 == 0 and 'active' or 'inactive')


async def bulk_copy(dsn: str, rows, batch_size: int = 10_000) -> int:
    conn = await asyncpg.connect(dsn)
    loaded = 0
    try:
        async with conn.transaction():
            await conn.execute(
                'CREATE TEMP TABLE staging_users '
                '(id bigint PRIMARY KEY, name text, status text)'
            )


            await conn.copy_to_table(
                'staging_users',
                source=rows,
            )
            loaded = await conn.fetchval('SELECT count(*) FROM staging_users')
    finally:
        await conn.close()
    return loaded


async def main() -> None:
    rows = generate_rows(25_000)
    total = await bulk_copy(
        'postgresql://user:pass@localhost/app',
        rows,
    )
    print('rows staged:', total)


asyncio.run(main())