How to eager load both forward and backward relationships in SQLAchemy one-to-many relation

32 Views Asked by At

I have a classical one-to-many relationship with the SQLAchemy 2.0 ORM:

from __future__ import annotations

import sqlalchemy as sa
from sqlalchemy.orm import (
    sessionmaker,
    DeclarativeBase,
    MappedAsDataclass,
    Mapped,
    mapped_column,
    relationship,
    immediateload,
)

# tables definition

class Base(MappedAsDataclass, DeclarativeBase):
    pass


class Parent(Base):
    __tablename__ = "parent"
    id: Mapped[int] = mapped_column(primary_key=True, init=False)
    children: Mapped[list[Child]] = relationship("Child", back_populates="parent")


class Child(Base):
    __tablename__ = "child"
    id: Mapped[int] = mapped_column(primary_key=True, init=False)
    parent_id: Mapped[int] = mapped_column(sa.ForeignKey(Parent.id), init=False)
    parent: Mapped[Parent] = relationship(Parent, back_populates="children", init=False, repr=False)

# main

engine = sa.create_engine("sqlite:///tmp.db", echo=True)
session_factory = sessionmaker(bind=engine, autoflush=True, expire_on_commit=False)

Base.metadata.drop_all(engine, checkfirst=True)
Base.metadata.create_all(engine)

# insert mock data
with session_factory.begin() as session:
    session.add_all(
        [
            Parent(children=[Child()]),  # p1 with 1 child
            Parent(children=[Child(), Child()]),  # p2 with 2 children
        ]
    )

I would like to make a select query on Parent table, that would eagerly load both forward and backward relationships.

I manage to eagerly load the Parent.children relationship thanks to immediateload, but I can't manage to eagerly back populate the Child.parent relationships as well, whereas it should come for free:

# read mock data
with session_factory.begin() as session:
    stmt = sa.select(Parent).options(
        immediateload(Parent.children),
        # immediateload(Child.parent),  # does not work
    )
    p1, p2 = session.scalars(stmt).all()
    assert len(p1.children) == 1

    # this DOES work, but only within session `Child.parent`
    assert p1.children[0].parent is p1

# test objects properties outside session
assert len(p2.children) == 2
# work because, we populated manually `Child.parent` relationship within session (and we set `expire_on_commit=True`)
assert p1.children[0].parent is p1
# does NOT work, because `Child.parent` relationship was NOT populated eagerly
assert p2.children[0].parent is p2
assert p2.children[1].parent is p2

I understand that simplify calling [child.parent for child in p1.children + p2.children] within session would do the trick, but I can't understand how to somehow make this happen automatically.

0

There are 0 best solutions below