优化SQLAlchemy Flush：拦截INSERT实现MySQL批量多行插入

需求背景 #

使用SQLAlchemy 1.4，希望拦截session.flush()生成的所有INSERT语句，不直接发送到数据库，而是按表合并为单条多行INSERT语句执行，以此减少MySQL批量插入的网络调用次数，提升性能。当前代码能捕获单条INSERT，但无法自动合并为多行语句。

可行方案：捕获ORM待插入对象而非SQL #

直接解析捕获的SQL语句进行合并容易出错，更可靠的方式是从Session中收集待插入的ORM对象，按表分组后生成多行INSERT，同时维护ORM对象的状态（如自增ID、关联关系）。

实现步骤 #

1. 监听before_flush事件收集待插入对象：拦截Session的flush过程，收集所有待插入的ORM对象，阻止默认的单条INSERT生成。
2. 按表依赖排序插入：处理表之间的外键依赖（如先插入User表，再插入依赖User的Profile表）。
3. 生成多行INSERT执行：使用SQLAlchemy的insert().values()生成批量插入语句，并将数据库返回的自增ID赋值给ORM对象，保持Session状态一致。

完整代码示例 #

from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, insert
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship, event
from pymysql.constants import CLIENT

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    id   = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(50), nullable=False)
    profiles = relationship(
        'Profile', back_populates='user', cascade='all, delete-orphan'
    )

class Profile(Base):
    __tablename__ = 'profiles'
    id      = Column(Integer, primary_key=True, autoincrement=True)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    bio     = Column(String(200), nullable=True)
    user    = relationship('User', back_populates='profiles')

engine = create_engine(
    'mysql+pymysql://user:pass@host/dbname',
    connect_args={'client_flag': CLIENT.MULTI_STATEMENTS}
)
Session = sessionmaker(bind=engine, autoflush=False)

# 存储待插入的ORM对象，key为模型类，value为对象列表
pending_inserts = {}

def before_flush(session, flush_context, instances):
    # 收集所有新创建的ORM对象
    for instance in session.new:
        cls = type(instance)
        if cls not in pending_inserts:
            pending_inserts[cls] = []
        pending_inserts[cls].append(instance)
    # 清空session.new，阻止默认的单条INSERT逻辑
    session.new.clear()

# 注册flush前的监听事件
event.listen(Session, 'before_flush', before_flush)

def get_table_dependency_order(classes):
    """根据外键依赖对表进行排序，先插入无依赖的父表"""
    order = []
    processed = set()
    remaining_classes = list(classes)
    
    while remaining_classes:
        for cls in list(remaining_classes):
            table = cls.__table__
            has_unprocessed_fk = False
            # 检查当前表的外键是否指向未处理的表
            for fk in table.foreign_keys:
                referenced_table = fk.column.table
                referenced_cls = next(c for c in remaining_classes if c.__table__ == referenced_table)
                if referenced_cls not in processed:
                    has_unprocessed_fk = True
                    break
            if not has_unprocessed_fk:
                order.append(cls)
                processed.add(cls)
                remaining_classes.remove(cls)
                break
    return order

def bulk_insert_from_pending(session):
    """批量插入收集到的ORM对象，并维护对象状态"""
    if not pending_inserts:
        return
    
    # 按依赖顺序处理表
    sorted_classes = get_table_dependency_order(list(pending_inserts.keys()))
    
    for cls in sorted_classes:
        instances = pending_inserts[cls]
        table = cls.__table__
        
        # 提取对象的字段值，排除自增主键（由数据库生成）
        values = []
        for instance in instances:
            obj_data = {}
            for col in table.columns:
                if not col.autoincrement:
                    obj_data[col.name] = getattr(instance, col.name)
            values.append(obj_data)
        
        if not values:
            continue
        
        # 执行多行INSERT
        result = session.execute(insert(table).values(values))
        
        # 处理自增主键，将生成的ID赋值给ORM对象
        if table.c.id.autoincrement:
            first_id = result.lastrowid
            for idx, instance in enumerate(instances):
                setattr(instance, 'id', first_id + idx)
                # 维护关联关系（比如Profile的user属性）
                if hasattr(instance, 'user') and instance.user_id is not None:
                    instance.user = session.query(User).get(instance.user_id)
        
        # 将对象标记为已持久化，同步Session状态
        for instance in instances:
            session.add(instance)
            session.expunge(instance)
            session.add(instance)
    
    # 清空待插入容器
    pending_inserts.clear()

# 使用示例
if __name__ == '__main__':
    session = Session()
    
    # 创建带关联的ORM对象
    alice = User(name='alice')
    alice.profiles.extend([
        Profile(bio='Bio 1'),
        Profile(bio='Bio 2'),
        Profile(bio='Bio 3')
    ])
    session.add(alice)
    
    # 触发flush（此时仅收集对象，不执行默认插入）
    session.flush()
    # 执行批量插入
    bulk_insert_from_pending(session)
    # 提交事务
    session.commit()

方案优势 #

• 保留ORM的关联关系管理，无需手动编写原生SQL
• 自动按表依赖排序，避免外键约束错误
• 生成标准的多行INSERT语句，减少网络调用次数
• 同步ORM对象状态，保证Session的一致性

替代方案：使用bulk_save_objects #

如果不需要自动拦截flush，也可以手动使用bulk_save_objects，但需手动处理关联对象的外键：

session = Session()
alice = User(name='alice')
# 先批量插入User
session.bulk_save_objects([alice])
session.flush()  # 获取自增ID
# 给Profile设置user_id
for p in alice.profiles:
    p.user_id = alice.id
# 批量插入Profile
session.bulk_save_objects(alice.profiles)
session.commit()

内容的提问来源于stack exchange，提问作者Meenti