优化SQLAlchemy多表关联批量插入至MySQL Aurora时的ID映射与关系保留问题

看起来你已经在优化批量插入性能的路上做了不少尝试，但卡在了ID映射维护和关联关系保留的关键点上——我来帮你拆解现有代码的问题，再给出几个针对性的优化方案，既能减少数据库往返次数，又能确保关联关系正确。

先聊聊现有代码的核心问题 #

1. ORM对象内存开销大：用add_all创建大量ModelCalculation、Dashboard等ORM对象，不仅占用内存，还需要SQLAlchemy追踪每个对象的状态，拖慢速度。
2. 关系更新效率低：公式的parent_id需要先设为None插入，再逐个修改对象后flush，生成大量零散的UPDATE语句，数据库压力大。
3. Graph部分逐行插入：graph_domain.create逐个创建Graph及关联数据，没有利用批量操作的优势。
4. Flush次数还是偏多：虽然分了阶段，但每个add_all后都flush，本质上还是多次数据库往返，没有最大化批量操作的收益。

优化方案：用批量操作替代ORM对象逐个处理，精准维护ID映射 #

SQLAlchemy 2.0的bulk_insert_mappings和bulk_update_mappings是批量插入的利器，配合return_defaults=True可以直接获取自增ID，完美解决你的ID映射问题，同时大幅减少数据库交互次数。

步骤1：精简ActiveTemplate插入（仅这步需要flush） #

因为后续所有对象都依赖active_template.id，所以这步必须先插入并获取ID：

active_template = ActiveTemplate(company_id=company_id, template_id=template_id,
                                 start_date=payload.start_date, branch_id=branch_id)
db.add(active_template)
db.flush()  # 必须flush拿到active_template的自增ID

步骤2：批量插入ModelCalculation，直接建立新旧ID映射 #

用bulk_insert_mappings批量插入，避免创建大量ORM对象，同时直接返回新生成的ID：

old_id_to_new_model_id = {}
model_dicts = []

for model_attr in data.models:
    old_id = model_attr.id
    # 整理插入字典，覆盖必要字段
    model_dict = model_attr.model_dump()
    model_dict.update({
        'active_template_id': active_template.id,
        'company_id': company_id,
        'branch_id': branch_id,
        'id': None  # 让数据库生成自增ID
    })
    model_dicts.append(model_dict)
    # 记录旧ID在列表中的索引，后面用来匹配新ID
    old_id_to_new_model_id[old_id] = len(model_dicts) - 1

# 批量插入并返回新ID（return_defaults=True仅适用于支持返回自增ID的数据库，Aurora MySQL完全支持）
result = db.bulk_insert_mappings(ModelCalculation, model_dicts, return_defaults=True)

# 更新映射：把旧ID关联到新生成的数据库ID
for old_id, idx in old_id_to_new_model_id.items():
    old_id_to_new_model_id[old_id] = result[idx]['id']

步骤3：批量插入Dashboard，复用Model的ID映射 #

和Model的逻辑一致，先替换依赖的model_id，再批量插入：

old_id_to_new_dashboard_id = {}
dashboard_dicts = []

for dash_attr in data.dashboards:
    old_id = dash_attr.id
    new_model_id = old_id_to_new_model_id.get(dash_attr.model_id)
    if not new_model_id:
        logging.error(f"关联模型ID {dash_attr.model_id} 在仪表板 {old_id} 中不存在")
        continue
    
    dash_dict = dash_attr.model_dump()
    dash_dict.update({
        'model_id': new_model_id,
        'branch_id': branch_id,
        'company_id': company_id,
        'id': None
    })
    dashboard_dicts.append(dash_dict)
    old_id_to_new_dashboard_id[old_id] = len(dashboard_dicts) - 1

# 批量插入并获取新ID
result = db.bulk_insert_mappings(Dashboard, dashboard_dicts, return_defaults=True)
for old_id, idx in old_id_to_new_dashboard_id.items():
    old_id_to_new_dashboard_id[old_id] = result[idx]['id']

步骤4：优化Formula的插入与parent_id更新 #

先批量插入所有Formula（parent_id暂设为None），再用批量更新替换parent_id，避免逐个修改ORM对象：

old_id_to_new_formula_id = {}
formula_dicts = []
old_id_to_old_parent_id = {}  # 保存原始parent_id，用于后续映射

for formula_attr in data.formulas:
    old_id = formula_attr.id
    new_model_id = old_id_to_new_model_id.get(formula_attr.model_id)
    if not new_model_id:
        logging.error(f"关联模型ID {formula_attr.model_id} 在公式 {old_id} 中不存在")
        continue
    
    # 保存原始parent_id
    old_id_to_old_parent_id[old_id] = formula_attr.parent_id
    
    formula_dict = formula_attr.model_dump()
    formula_dict.update({
        'model_id': new_model_id,
        'branch_id': branch_id,
        'id': None,
        'parent_id': None  # 暂设为None，后续批量更新
    })
    formula_dicts.append(formula_dict)
    old_id_to_new_formula_id[old_id] = len(formula_dicts) - 1

# 批量插入公式并获取新ID
result = db.bulk_insert_mappings(Formula, formula_dicts, return_defaults=True)
for old_id, idx in old_id_to_new_formula_id.items():
    old_id_to_new_formula_id[old_id] = result[idx]['id']

# 整理批量更新的操作
update_ops = []
for old_id, new_formula_id in old_id_to_new_formula_id.items():
    old_parent_id = old_id_to_old_parent_id[old_id]
    if old_parent_id:
        new_parent_id = old_id_to_new_formula_id.get(old_parent_id)
        if new_parent_id:
            update_ops.append({
                'id': new_formula_id,
                'parent_id': new_parent_id
            })

# 批量更新parent_id，仅生成一条（或少量）UPDATE语句
if update_ops:
    db.bulk_update_mappings(Formula, update_ops)

步骤5：批量处理Graph及关联数据 #

把Graph、DataGroup、DataPoint的批量插入整合起来，避免逐个创建：

graph_objects = []

for graph_attr in data.graphs:
    old_dashboard_id = graph_attr.graph_data.dashboard_id
    new_dashboard_id = old_id_to_new_dashboard_id.get(old_dashboard_id)
    if not new_dashboard_id:
        logging.error(f"关联仪表板ID {old_dashboard_id} 在图表中不存在")
        continue
    
    # 更新Graph的dashboard_id和branch_id
    graph_data_dict = graph_attr.graph_data.model_dump()
    graph_data_dict.update({
        'dashboard_id': new_dashboard_id,
        'branch_id': branch_id
    })
    graph = Graph(**graph_data_dict)
    
    # 处理DataGroup和DataPoint的formula_id映射
    for dg_attr in graph_attr.data_groups:
        dg = DataGroup(**dg_attr.model_dump())
        graph.data_groups.append(dg)
        
        for dp_attr in dg_attr.data_points:
            old_formula_id = dp_attr.formula_id
            dp_attr.formula_id = old_id_to_new_formula_id.get(old_formula_id, old_formula_id)
            dp = DataPoint(**dp_attr.model_dump())
            dg.data_points.append(dp)
    
    graph_objects.append(graph)

# 批量保存Graph及其关联对象，SQLAlchemy会自动处理关联关系的插入
db.bulk_save_objects(graph_objects, return_defaults=True)

最后统一提交 #

db.commit()
return active_template

关键注意事项 #

1. 模型字段配置：确保所有自增ID字段在SQLAlchemy模型中正确配置autoincrement=True，比如：

   id = Column(Integer, primary_key=True, autoincrement=True)
   

2. 批量操作的限制：bulk_insert_mappings不会触发ORM的before_insert/after_insert钩子，如果你的模型依赖这些钩子逻辑，需要手动处理。
3. 事务一致性：整个流程放在一个事务中，任何步骤出错都会回滚，确保数据一致性。
4. 内存占用优化：用字典列表替代ORM对象，大幅减少内存开销，适合插入大量数据的场景。

这样优化后，你会发现flush次数大幅减少（仅ActiveTemplate需要一次flush，其余批量操作直接执行SQL），数据库往返次数降到最低，同时ID映射和关联关系也能精准维护。

备注：内容来源于stack exchange，提问作者Tomer Lankri