Python SQLAlchemy 批量插入时忽略重复键

Posted 2023-02-02 gcygeeker

SQLAlchemy 批量插入的方法如 bulk_save_objects() 默认无法插入重复 key，如果插入会报错；而有时候我们需要忽略而非报错，经过研究发现下面代码可以解决这个问题：

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import Insert

"""
When imported, automatically make all insert not fail on duplicate keys
"""


@compiles(Insert, "mysql")
def mysql_insert_ignore(insert, compiler, **kw):
    return compiler.visit_insert(insert.prefix_with("IGNORE"), **kw)


@compiles(Insert, "postgresql")
def postgresql_on_conflict_do_nothing(insert, compiler, **kw):
    statement = compiler.visit_insert(insert, **kw)
    # IF we have a "RETURNING" clause, we must insert before it
    returning_position = statement.find("RETURNING")
    if returning_position >= 0:
        return (
            statement[:returning_position]
            + "ON CONFLICT DO NOTHING "
            + statement[returning_position:]
        )
    else:
        return statement + " ON CONFLICT DO NOTHING"

在需要忽略重复 key 的代码块中 import 该 package 就可以。比如保存上述代码为insert_ignore.py，然后在需要使用的地方：

import insert_ignore
...
session.bulk_save_objects(...)

注意控制使用的范围，在 import 的范围内所有插入遇到的重复 key 都会忽略

参考资料1