SQLAlchemy 自动刷新在 Session 上不起作用。

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

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