AsyncPG PostgreSQL COPY Binary Load

Bulk load millions of rows into PostgreSQL with asyncpg COPY, streaming records from an async generator with a single protocol message.

PythonData Fetching

Code

snippet
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())