防止 SQLAlchemy 自动设置 IDENTITY_INSERT

Question

我有一个带有自动递增主键的 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()

如果 tub 已经存在于数据库中，我希望更新它的权重。如果没有，我希望插入失败，因为它包含显式主键。

Answer 1

根据SQLAlchemy 1.3 Documentation：

对于不需要默认生成IDENTITY 的情况，请在第一个整数主键列上为Column.autoincrement 标志指定False：

m = MetaData()
t = Table('t', m,
        Column('id', Integer, primary_key=True, autoincrement=False),
        Column('x', Integer))
m.create_all(engine)

Answer 2

为什么不直接重写构造函数？

class Tub(Base):
    __tablename__ = 'Tub'

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

    def __init__(self, id=None, **kwargs):
        if id is not None:
            raise ValueError("Hey, don't pass in `id`!")
            # or you could even just ignore that it was passed in and issue a warning.
        super().__init__(**kwargs)

如果你想将同样的方法应用到很多类，你可以实现一个 mixin：

import warnings

class AutoIncrMixin:
    id = Column('ID', Integer, primary_key=True, autoincrement=True)

    def __init__(self, id=None, **kwargs):
        if id is not None:
            warnings.warn("supplied `id` value was ignored")
        super().__init__(**kwargs)

class Tub(Base, AutoIncrMixin):
    tare_weight = ...

Answer 3

从长远来看，您可以进行完整的猴子修补。我不认为它可以解决你所有的问题，但它确实阻止了SET IDENTITY_INSERT 的使用。你可以修改MSExecutionContexts pre_exec和post_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

我使用猴子的粗略示例修补了利用_enable_identity_insert 功能的功能，并删除了这些部分。现在每次尝试专门INSERTing 一个 ID 都会给你以下错误：

sqlalchemy.exc.IntegrityError: (pyodbc.IntegrityError) ('23000', "[23000] [Microsoft][SQL Server Native Client 11.0][SQL Server]Cannot insert explicit value for identity column in table 'Tub' when IDENTITY_INSERT 设置为 OFF。(544) (SQLExecDirectW)")

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

Answer 4

我正在添加我自己的答案，以详细说明我的 cmets 到 @SuperShoot 的答案。我认为，如果您打算首先阻止主键获取值的解决方法，您应该通过 property 而不是构造函数：

class Tub(Base):
    __tablename__ = 'Tub'

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

    @property
    def id(self):
        return self.__id

这会阻止你做这两件事，而不仅仅是第一件事：

tub = Tub(id=5, tare_weight=100)

or

tub = Tub(tare_weight=100)
tub.id = 5

AttributeError: 无法设置属性

这样做将停止显式主键插入，因为id 永远不会得到值。缺点是您不能创建对象然后调用session.merge() 来发出UPDATE 查询。

...除了 Python 允许您在确实需要时作弊：

tub = Tub(_Tub__id=5, tare_weight=100)

or

tub = Tub(tare_weight=100)
tub._Tub__id = 5