防止SQLAlChemy自动设置IDENTITY_INSERT

2022-02-22 00:00:00 python sqlalchemy sql-server

我有一个SQL Server表,其中包含一个自动递增的主键,我正在使用SQLAlChemy操作该主键。如果我创建包含主键表模型的实例,而不是关闭它并调用session.add(instance),则SQLAlChemy会自动向数据库发出SET IDENTITY_INSERT [table] ON,并且插入成功。

有没有办法自己控制呢?如果我尝试插入特定的主键而不是让它自动递增,我宁愿收到错误。如果确实需要设置主键,我宁愿显式说明。

模型如下:

Base = declarative_base()

class Tub(Base):
    __tablename__ = 'Tub'

    id = Column('ID', Integer, primary_key=True, autoincrement=True)
    tare_weight = Column('TareWeight', Float(53), nullable=False)

这里有一个插入示例:

t = Tub(id=20, tare_weight=200)
session.add(t)
session.commit()

以下是生成的SQL:

BEGIN
SET IDENTITY_INSERT [Tub] ON
INSERT INTO [Tub] ([ID], [TareWeight]) VALUES (20, 200)
SET IDENTITY_INSERT [Tub] OFF
COMMIT

更新:

我意识到的另一个相关情况是:

t = Tub(id=20, tare_weight=200)
session.merge(t)
session.commit()
如果数据库中已存在浴缸,我希望更新其权重。如果没有,我希望插入失败,因为它包含显式主键。


解决方案

作为一种可能,您可以完全修补猴子。我不认为它能解决您所有的问题,但它确实阻止了SET IDENTITY_INSERT的使用。您可以使用猴子补丁MSExecutionContextspre_execpost_exec。例如:

from sqlalchemy import engine
from sqlalchemy.dialects.mssql.base import MSExecutionContext

def pre_exec(self):
    if self.isinsert:
        tbl = self.compiled.statement.table
        seq_column = tbl._autoincrement_column
        insert_has_sequence = seq_column is not None

        self._select_lastrowid = (
            not self.compiled.inline
            and insert_has_sequence
            and not self.compiled.returning
            and not self._enable_identity_insert
            and not self.executemany
        )


def post_exec(self):
    conn = self.root_connection
    if self._select_lastrowid:
        if self.dialect.use_scope_identity:
            conn._cursor_execute(
                self.cursor,
                "SELECT scope_identity() AS lastrowid",
                (),
                self,
            )
        else:
            conn._cursor_execute(
                self.cursor, "SELECT @@identity AS lastrowid", (), self
            )
        # fetchall() ensures the cursor is consumed without closing it
        row = self.cursor.fetchall()[0]
        self._lastrowid = int(row[0])

    if (
        self.isinsert or self.isupdate or self.isdelete
    ) and self.compiled.returning:
        self._result_proxy = engine.FullyBufferedResultProxy(self)


MSExecutionContext.pre_exec = pre_exec
MSExecutionContext.post_exec = post_exec

我在使用的粗略示例Money会修补利用_enable_identity_insert功能的函数,并删除这些部分。现在,每次尝试明确INSERT查找ID时都会出现以下错误:

sqlalChemy.exc.IntegrityError:(pyodbc.IntegrityError)(‘23000’,"[23000][Microsoft][SQL Server Native Client 11.0][SQL Server]当IDENTITY_INSERT设置为OFF时,[SQL Server]无法为表‘Tub’中的IDENTITY列插入显式值。(23000)(SQLExecDirectW)")

您现在的问题可能是您完全依赖自动递增功能,但它可能对某些人有用,或者您找到了其他方法来绕过此功能。

相关文章