如何在 flask_sqlalchemy 中使用 PostgreSQL 的“INSERT...ON CONFLICT”（UPSERT）功能？

Posted 2023-02-23

【中文标题】如何在 flask_sqlalchemy 中使用 PostgreSQL 的“INSERT...ON CONFLICT”（UPSERT）功能？

【英文标题】：How to use PostgreSQL's "INSERT...ON CONFLICT" (UPSERT) feature with flask_sqlalchemy?

【发布时间】：2017-05-28 22:01:10

【问题描述】：

INSERT 语句中的 PostgreSQL ON CONFLICT 子句提供“更新插入”功能（即更新现有记录，或在不存在此类记录时插入新记录）。 SQLAlchemy 通过 PostgreSQL 方言的 Insert 对象上的 on_conflict_do_nothing 和 on_conflict_do_update 方法支持此功能（如 here 所述）：

from sqlalchemy.dialects.postgresql import insert

insert_stmt = insert(my_table).values(
    id='some_existing_id',
    data='inserted value'
)

do_nothing_stmt = insert_stmt.on_conflict_do_nothing(
    index_elements=['id']
)

conn.execute(do_nothing_stmt)

do_update_stmt = insert_stmt.on_conflict_do_update(
    constraint='pk_my_table',
    set_=dict(data='updated value')
)

conn.execute(do_update_stmt)

我正在使用flask_sqlalchemy，它为您管理 SQLAlchemy 的引擎、会话和连接。为了向数据库中添加元素，我创建了一个模型实例，将其添加到数据库会话中，然后调用 commit，如下所示：

from flask import Flask
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)
db = SQLAlchemy(app)

class MyTable(db.Model):
    id = db.Column(UUID, primary_key=True)
    data = db.Column(db.String)

relation = MyTable(id=1, data='foo')
db.session.add(relation)
db.session.commit()

所以Insert 对象完全被flask_sqlalchemy 包裹和遮蔽。

如何访问特定于 PostgreSQL 的方言方法来执行 upsert？我是否需要绕过flask_sqlalchemy 并创建自己的会话？如果这样做，如何确保没有冲突？

【参考方案1】：

事实证明，您可以在db.session 上执行较低级别的语句。所以一个解决方案看起来像这样：

from flask import Flask
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.dialects.postgresql import insert as pg_insert

app = Flask(__name__)
db = SQLAlchemy(app)

class MyTable(db.Model):
    id = db.Column(UUID, primary_key=True)
    data = db.Column(db.String)

    def __init__(self, _id, *args, **kwargs):
        self.id = _id
        self.data = kwargs['data']

    def as_dict(self):
        return 'id': self.id, 'data': self.data

    def props_dict(self):
        d = self.as_dict()
        d.pop('id')
        return d

relation = MyTable(id=1, data='foo')
statement = pg_insert(MyTable)\.
    values(**relation.as_dict()).\
    on_conflict_do_update(constraint='id',
                          set_=relation.props_dict())

db.session.execute(statement)
db.session.commit()

我的模型类中的as_dict() 和props_dict() 方法允许我使用构造函数从传入的HTTP 请求中过滤掉不需要的属性。

【讨论】：

即使id 被创建为db.Column(db.Integer, primary_key=True) 字段，这对我来说也失败了constraint "id" for table "foobar" does not exist 的错误。

这对我帮助很大。很好地满足了我运行 on_conflict_do_nothing() 的需要，并对您的代码示例进行了一些小改动。

@deed02392 我遇到了同样的错误，你找到解决办法了吗？

嗨@avocado，我最后一定完成了，但我想这是显而易见的，因此我没有再次在这里更新。如果你想不通，我们可以聊天。

【参考方案2】：

使用编译扩展的替代方法 (https://docs.sqlalchemy.org/en/13/core/compiler.html)：

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

@compiles(Insert)
def compile_upsert(insert_stmt, compiler, **kwargs):
    """
    converts every SQL insert to an upsert  i.e;
    INSERT INTO test (foo, bar) VALUES (1, 'a')
    becomes:
    INSERT INTO test (foo, bar) VALUES (1, 'a') ON CONFLICT(foo) DO UPDATE SET (bar = EXCLUDED.bar)
    (assuming foo is a primary key)
    :param insert_stmt: Original insert statement
    :param compiler: SQL Compiler
    :param kwargs: optional arguments
    :return: upsert statement
    """
    pk = insert_stmt.table.primary_key
    insert = compiler.visit_insert(insert_stmt, **kwargs)
    ondup = f'ON CONFLICT (",".join(c.name for c in pk)) DO UPDATE SET'
    updates = ', '.join(f"c.name=EXCLUDED.c.name" for c in insert_stmt.table.columns)
    upsert = ' '.join((insert, ondup, updates))
    return upsert

这应该确保所有插入语句都表现为 upsert。