如何基于唯一非主键用SQLAlchemy Model实现MySQL Upsert？

环境 #

• Python 3.10
• SQLAlchemy 1.4.44
• 数据库：MySQL 8

问题描述 #

需要对继承自Base的ORM模型实例执行Upsert操作，要求基于ORM Model结合Session（支持提交、回滚）完成，而非直接使用核心API的insert()。数据库表foobar结构如下，且email和company为组合唯一约束：

mysql> desc foobar;
+---------+--------------+------+-----+---------+----------------+
| Field   | Type         | Null | Key | Default | Extra          |
+---------+--------------+------+-----+---------+----------------+
| id      | int          | NO   | PRI | NULL    | auto_increment |
| email   | varchar(255) | NO   |     | NULL    |                |
| company | varchar(255) | NO   |     | NULL    |                |
| memo    | varchar(255) | YES  |     | NULL    |                |
+---------+--------------+------+-----+---------+----------------+
4 rows in set (0.01 sec)

已编写基础代码，需在insert_to_db函数中实现Upsert逻辑。

解决方案 #

步骤1：完善ORM模型（添加唯一约束） #

首先在Foobar模型中显式声明email和company的组合唯一约束，让SQLAlchemy明确Upsert的判断依据：

from sqlalchemy import UniqueConstraint

class Foobar(Base):
    __tablename__ = 'foobar'

    id = Column(Integer, primary_key=True)
    email = Column(String(255), nullable=False)
    company = Column(String(255), nullable=False)
    memo = Column(String(255), nullable=True)
    
    # 添加组合唯一约束
    __table_args__ = (
        UniqueConstraint('email', 'company', name='uq_foobar_email_company'),
    )

    def __repr__(self):
        columns = ', '.join([
            '{0}={1}'.format(k, repr(self.__dict__[k]))
            for k in self.__dict__.keys() if k[0] != '_'
        ])
        return '<{0}({1})>'.format(
            self.__class__.__name__, columns
        )

步骤2：实现Upsert逻辑 #

提供两种实现方式，可根据场景选择：

方式1：ORM风格（查询+更新/插入）

先根据唯一约束查询记录，存在则更新字段，不存在则添加新实例，直观符合ORM使用习惯：

def insert_to_db(session, memo):
    # 根据唯一约束查询现有记录
    existing = session.query(Foobar).filter(
        Foobar.email == 'example@example.com',
        Foobar.company == 'example'
    ).first()

    if existing:
        # 存在则更新memo字段
        existing.memo = memo
        print(f"更新现有记录: {existing}")
    else:
        # 不存在则创建新实例并添加
        foobar = Foobar(
            email='example@example.com',
            company='example',
            memo=memo,
        )
        session.add(foobar)
        print(f"插入新记录: {foobar}")

方式2：高效Upsert（核心API+Session执行）

如果需要更高性能（减少一次查询），可使用SQLAlchemy核心API结合MySQL的ON DUPLICATE KEY UPDATE，通过Session执行，一次数据库操作完成Upsert：

from sqlalchemy import insert

def insert_to_db(session, memo):
    # 构造Upsert语句
    stmt = insert(Foobar).values(
        email='example@example.com',
        company='example',
        memo=memo
    ).on_duplicate_key_update(
        # 唯一键冲突时更新memo字段
        memo=memo
    )
    # 通过Session执行语句
    session.execute(stmt)
    # 查询验证结果
    updated = session.query(Foobar).filter(
        Foobar.email == 'example@example.com',
        Foobar.company == 'example'
    ).first()
    print(f"Upsert结果: {updated}")

完整测试 #

整合逻辑后，原main函数无需修改，执行后第一次会插入记录，第二次会将memo更新为memo2。

注意事项 #

• 确保数据库表中已存在email和company的组合唯一索引，否则ON DUPLICATE KEY UPDATE不会生效。
• 方式1适合单条记录的简单场景，方式2适合批量操作或对性能要求较高的场景。
• 使用Session时，需调用session.commit()提交事务，异常时调用session.rollback()回滚。

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