Skip to content

22. FastAPI 使用 SQLAlchemy 2 Async 新增关系数据(单条与百万条) ​

小规模用对象图,大规模用异步批量语句 ​

单条业务创建强调关系清晰和事务一致;百万条导入强调控制内存、减少网络往返、按批提交。两者都使用第 21 篇的映射类,并通过第 20 篇的 AsyncSession 操作 MySQL。

准备:完成第 20、21 篇的配置和建表。以下脚本使用主库 SessionLocal,数据库驱动为 aiomysql,所有数据库 I/O 都需要 await。

情况一:创建一个用户,同时建立三种关系 ​

保存为 insert_one.py:

python
import asyncio
from uuid import uuid4

from sqlalchemy import select

from database import SessionLocal, master_engine
from models import Post, Profile, Role, User


async def create_one() -> dict:
    async with SessionLocal.begin() as db:
        role = await db.scalar(select(Role).where(Role.name == "reader"))
        if role is None:
            role = Role(name="reader")

        user = User(
            name="小明",
            email=f"{uuid4().hex}@example.com",
            profile=Profile(bio="开始学习 FastAPI"),
            posts=[
                Post(title="第一篇", content="你好"),
                Post(title="第二篇", content="学习笔记"),
            ],
            roles=[role],
        )

        db.add(user)       # 只登记对象,不加 await
        await db.flush()   # 执行 INSERT,取得各表自增主键

        # 这些关系在构造对象图时已经加载,不需要隐式 lazy load
        result = {
            "user_id": user.id,
            "profile_id": user.profile.id,
            "post_ids": [post.id for post in user.posts],
            "role_id": role.id,
        }

    # 离开 begin() 上下文时自动提交;异常则回滚
    return result


async def main() -> None:
    try:
        print(await create_one())
    finally:
        await master_engine.dispose()


if __name__ == "__main__":
    asyncio.run(main())

执行:

powershell
uv run python insert_one.py

db.add(user) 会沿保存级联登记关联对象;await db.flush() 后取得主键;退出 SessionLocal.begin() 异步事务块才真正提交。任何一张表写入失败,本次事务整体回滚。

reader 首次并发创建可能竞争唯一键。异步请求同样不能消除数据库唯一约束竞争,正式接口应捕获 IntegrityError、回滚后重试整个工作单元,或预先创建角色。

情况二:异步导入一百万用户及其关联数据 ​

下面每个用户带一份资料、两篇文章和一个角色关联:一百万用户会额外生成一百万资料、两百万文章、一百万关联行。总量远大于一百万行,应预先规划磁盘、索引和备份空间。

保存为 import_many.py,默认先导入 1000 条,确认后再显式指定百万条:

python
import argparse
import asyncio
from uuid import uuid4

from sqlalchemy import insert, select

from database import SessionLocal, master_engine
from models import Post, Profile, Role, User, user_roles


async def import_users(total: int, batch_size: int = 2000) -> str:
    if total < 1 or batch_size < 1:
        raise ValueError("总量和批大小必须为正")

    tag = uuid4().hex

    async with SessionLocal.begin() as db:
        role = Role(name=f"bulk-{tag}")
        db.add(role)
        await db.flush()
        role_id = role.id

    print(f"batch_tag={tag}", flush=True)

    for start in range(0, total, batch_size):
        stop = min(start + batch_size, total)
        rows = [
            {
                "name": f"用户{i}",
                "email": f"{tag}-{i}@example.com",
                "active": True,
                "batch_tag": tag,
            }
            for i in range(start, stop)
        ]
        emails = [row["email"] for row in rows]

        async with SessionLocal.begin() as db:
            await db.execute(insert(User), rows)

            # MySQL 不照搬 PostgreSQL 的批量 INSERT RETURNING。
            # 根据唯一 email 查回本批主键,不猜测自增 ID 连续性。
            result = await db.execute(
                select(User.email, User.id).where(User.email.in_(emails))
            )
            ids = dict(result.all())
            if len(ids) != len(rows):
                raise RuntimeError("主键映射不完整")

            profiles, posts, links = [], [], []
            for row in rows:
                user_id = ids[row["email"]]
                profiles.append({"user_id": user_id, "bio": "批量导入"})
                posts.extend(
                    [
                        {
                            "user_id": user_id,
                            "title": "文章 A",
                            "content": "内容 A",
                        },
                        {
                            "user_id": user_id,
                            "title": "文章 B",
                            "content": "内容 B",
                        },
                    ]
                )
                links.append({"user_id": user_id, "role_id": role_id})

            await db.execute(insert(Profile), profiles)
            await db.execute(insert(Post), posts)
            await db.execute(user_roles.insert(), links)

        print(f"已提交 {stop}/{total}", flush=True)

    return tag


async def main(total: int, batch_size: int) -> None:
    try:
        tag = await import_users(total, batch_size)
        print("完成,批次标识:", tag)
    finally:
        await master_engine.dispose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--total", type=int, default=1000)
    parser.add_argument("--batch-size", type=int, default=2000)
    args = parser.parse_args()
    asyncio.run(main(args.total, args.batch_size))

先执行小规模测试,再执行百万导入:

powershell
uv run python import_many.py --total 1000
uv run python import_many.py --total 1000000 --batch-size 2000

两次执行生成不同批次,不会覆盖彼此。保存终端打印的 batch_tag,后续查询、更新和删除都用它隔离数据。批大小 2000 是起点,不是最优值;根据 SQL 包大小、网络延迟、内存和数据库负载调整。

为什么批量语句仍需 await ​

insert(User) 只是构造 SQL 语句,不会访问数据库;await db.execute(...) 才通过 aiomysql 真正执行。字典批量执行不需要构造一百万个 User 对象,也不会自动遍历 relationship,所以必须显式先插父表、取回 ID,再插子表和关联表。

异步不等于无限并发。本例按顺序提交批次,避免同时占用大量 MySQL 连接。不要在同一个 AsyncSession 中用 asyncio.gather() 并发执行多个批次;需要并行时必须使用独立 Session,并限制并发数。

每批提交意味着一批内部原子,但整个百万条任务不是全有或全无。失败时前面已提交的批次会保留。本示例没有自动断点续传:重新运行会创建新批次。正式导入应增加持久任务表、稳定业务唯一键、事务内检查点、错误明细和幂等策略。

性能验证与落地 ​

先用 1000、10000 条测量耗时、内存、每批延迟及数据库慢日志,再逐步放大。异步驱动主要改善等待 I/O 时的事件循环调度,不会让 MySQL 的单条 INSERT 本身突然加快,也不应编造“固定几秒插完百万条”的结论。

百万级导入不要占用一个普通 HTTP 请求一直等待。由独立命令或可靠任务队列执行,API 返回任务 ID,客户端查询进度。后台任务需要自己的生命周期管理,不能依赖请求结束后即将关闭的 Session。

练习:导入 1000 用户后,按批次统计四类记录,预期用户 1000、资料 1000、文章 2000、角色关联 1000;故意在某批中制造重复邮箱,验证本批回滚而前批保留。