SQLAlchemy Autoflush 不适用于 Session.add

Question

我想在数据库中保存新对象并在不调用的情况下接收新实体的自动增量 ID session.flush()。我尝试了 autoflush=True (docs) 选项，但没有成功。

这是我的代码。

from sqlalchemy import create_engine, Column, Integer, String, Boolean
from sqlalchemy.orm import sessionmaker, as_declarative

engine = create_engine("postgresql://postgres:@localhost/postgres")
Session = sessionmaker(autocommit=False, autoflush=True, bind=engine)
session = Session()
print(session.autoflush)  # -> True

@as_declarative()
class Base:
    __name__: str

class UserDb(Base):
    __tablename__ = "user"
    id = Column(Integer, primary_key=True)
    username: str = Column(String, unique=True)
    hashed_password: str = Column(String)
    is_active: bool = Column(Boolean, default=True)
    is_superuser: bool = Column(Boolean, default=False)


user = UserDb(
    username="username",
    hashed_password="password",
)
session.add(user)
# Here I want to have user.id
print(user.id)  # -> None
session.flush()
print(user.id)  # -> 12

有没有办法在 SQLAlchemy 中实现这种行为？

环境：python 3.9.7、SQLAlchemy 1.4.27

Answer 1

我敢肯定，这是一种反模式，但我已经通过猴子修补 SQLAlchemy 实现了预期的行为

from sqlalchemy.orm import Session


def add_decorator(method):  # type: ignore
    """
    Flush session to Database for Create/Update queries
    Useful for instant retrieving Autoincrement IDs

    So instead
    session.add(obj)
    session.flush()

    We can write just
    session.add(obj)

    But this solution requires more DB calls.
    To optimize code with large amount of session.add calls use

    with session.no_autoflush:
        ...
    """

    def wrapper(self, instance, _warn=True):  # type: ignore
        method(self, instance, _warn)
        if self.autoflush:
            self.flush()

    return wrapper


def patch() -> None:
    Session.add = add_decorator(Session.add)  # type: ignore

然后在需要开启的地方调用patch()autoflush

SQLAlchemy Autoflush 不适用于 Session.add

SQLAlchemy Autoflush does not work on Session.add

python

postgresql

sqlalchemy