SQLAlchemy Autoflush does not work on Session.add

user_8153147 0
Viewed 195

I want to save new objects in DB and receive autoincrement ID of new entity without call session.flush(). I tried autoflush=True (docs) option but had no luck.

Here my code.

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

Is there a way to achieve this behavior in SQLAlchemy?

Env: python 3.9.7, SQLAlchemy 1.4.27

sqlalchemy postgresql python
1 Answers

I'm sure, that it's an anti-pattern, but I've achieved desired behavior with monkey patching 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

Then you call patch() wherever you need to turn on autoflush

user_81531470
Related