SQLAlchemy+MariaDB复合主键批量Upsert实现问题求助

问题根源 #

你之前用and_组合单个主键字段的in查询，逻辑上是让每个主键字段各自匹配对应的值列表，而不是匹配主键字段的组合，导致数据库无法找到对应记录，自然不会执行更新逻辑。

两种可行解决方案 #

方案1：先查询过滤再分批次插入/更新 #

核心是用tuple_把复合主键字段打包，配合in_查询主键组合的集合，实现精准匹配：

1. 定义复合主键模型

from sqlalchemy import Column, Integer, String, tuple_, case, insert
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine

Base = declarative_base()

class CompositeModel(Base):
    __tablename__ = 'composite_table'
    # 复合主键：id + type_id
    id = Column(Integer, primary_key=True)
    type_id = Column(Integer, primary_key=True)
    name = Column(String(50))
    value = Column(Integer)

# 初始化数据库连接
engine = create_engine('mysql+mysqlconnector://用户名:密码@localhost/数据库名')
Session = sessionmaker(bind=engine)
session = Session()

2. 批量Upsert实现函数

def bulk_upsert_composite(session, model, data_list):
    # 提取模型的复合主键字段
    pk_cols = [col for col in model.__table__.columns if col.primary_key]
    # 提取待处理数据的主键组合（元组形式）
    pk_tuples = [tuple(d[col.name] for col in pk_cols) for d in data_list]
    
    # 精准查询已存在的主键组合
    existing_pks = set(session.query(*pk_cols).filter(
        tuple_(*pk_cols).in_(pk_tuples)
    ).all())
    
    # 拆分新增/更新数据集
    to_insert = []
    to_update = []
    for d in data_list:
        current_pk = tuple(d[col.name] for col in pk_cols)
        if current_pk not in existing_pks:
            to_insert.append(d)
        else:
            to_update.append(d)
    
    # 批量插入新数据
    if to_insert:
        session.bulk_insert_mappings(model, to_insert)
        print(f"插入 {len(to_insert)} 条新记录")
    
    # 批量更新现有数据
    if to_update:
        update_stmt = model.__table__.update()
        # 为每个非主键字段构造匹配更新逻辑
        non_pk_cols = [col for col in model.__table__.columns if not col.primary_key]
        for col in non_pk_cols:
            when_clauses = []
            for d in to_update:
                when_clauses.append(
                    (tuple_(*pk_cols) == tuple(d[pk_col.name] for pk_col in pk_cols), d[col.name])
                )
            update_stmt = update_stmt.values({col.name: case(*when_clauses, else_=col)})
        session.execute(update_stmt)
        print(f"更新 {len(to_update)} 条记录")
    
    session.commit()

3. 使用示例

test_data = [
    {"id": 1, "type_id": 1, "name": "A", "value": 10},
    {"id": 1, "type_id": 2, "name": "B", "value": 20},
    {"id": 2, "type_id": 1, "name": "C", "value": 30},
    # 这条已存在，会被更新
    {"id": 1, "type_id": 1, "name": "A_Updated", "value": 15}
]

bulk_upsert_composite(session, CompositeModel, test_data)

方案2：直接用MariaDB的ON DUPLICATE KEY UPDATE（更高效） #

MariaDB原生支持基于主键（包括复合主键）的重复判断，SQLAlchemy可以直接调用该语法，无需提前查询过滤，性能更优：

def bulk_upsert_on_duplicate(session, model, data_list):
    # 构造插入语句
    insert_stmt = insert(model.__table__).values(data_list)
    # 定义重复时的更新字段（所有非主键字段）
    update_fields = {
        col.name: insert_stmt.inserted[col.name]
        for col in model.__table__.columns if not col.primary_key
    }
    # 拼接ON DUPLICATE KEY UPDATE逻辑
    upsert_stmt = insert_stmt.on_duplicate_key_update(**update_fields)
    
    session.execute(upsert_stmt)
    session.commit()
    # 打印影响行数（插入+更新的总条数）
    print(f"批量Upsert完成，共影响 {session.execute('SELECT ROW_COUNT()').scalar()} 条记录")

关键提示 #

• 方案1适合需要在代码层对数据做额外处理（比如过滤、修改）的场景；
• 方案2直接委托数据库处理重复判断，代码更简洁，大数据量下性能更好，优先推荐；
• 确保MariaDB表的复合主键定义正确，否则ON DUPLICATE KEY UPDATE无法触发更新逻辑。

内容的提问来源于stack exchange，提问作者Shan-x