from sqlalchemy import Column, ForeignKey, Integer, create_engine, insert, select
from sqlalchemy.orm import Session, declarative_base, noload, relationship, subqueryload

Base = declarative_base()


# two tables, with back_populates
class A(Base):
    __tablename__ = "a"

    a_id = Column(
        Integer,
        nullable=False,
        primary_key=True,
        autoincrement=True,
    )
    b_connections = relationship("B", back_populates="a_connections", primaryjoin="foreign(B.a_id) == A.a_id", uselist=True)


class B(Base):
    __tablename__ = "b"

    b_id = Column(
        Integer,
        nullable=False,
        primary_key=True,
        autoincrement=True,
    )
    a_id = Column(
        ForeignKey("a.a_id"),
        nullable=False,
    )
    a_connections = relationship("A", back_populates="b_connections", primaryjoin="A.a_id == foreign(B.a_id)", uselist=True)


engine = create_engine("sqlite://", echo=False)
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)

# initial setup block
with Session(engine) as session:
    a_insert = session.execute(insert(A).values())
    a_id = a_insert.lastrowid
    b_insert = session.execute(insert(B).values(a_id=a_id))
    b_id = b_insert.lastrowid

    session.commit()


a_stmt = (
    select(A)
    .options(
        subqueryload(A.b_connections).options(noload(B.a_connections)),
    )
    .where(A.a_id == a_id)
)
b_stmt = (
    select(B)
    .options(
        subqueryload(B.a_connections).options(noload(A.b_connections)),
    )
    .where(B.b_id == b_id)
)

with Session(engine) as session:
    # these blocks can be run in either order, the second one always fails
    (a, *_), *_ = session.execute(a_stmt).fetchall()
    assert a_id == a.a_id
    assert 1 == len(a.b_connections)
    a_b = a.b_connections[0]
    assert b_id == a_b.b_id

    (b, *_), *_ = session.execute(b_stmt).fetchall()
    assert b_id == b.b_id
    assert 1 == len(b.a_connections)  # this assertion fails
    b_a = b.a_connections[0]
    assert a_id == b_a.a_id