from sqlalchemy import and_
from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import ForeignKey
from sqlalchemy import func
from sqlalchemy import Integer
from sqlalchemy import select
from sqlalchemy.orm import aliased
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import relationship
from sqlalchemy.orm import selectinload
from sqlalchemy.orm import Session
Base = declarative_base()

class A(Base):
    __tablename__ = "a"

    id = Column(Integer, primary_key=True)

class B(Base):
    __tablename__ = "b"
    id = Column(Integer, primary_key=True)
    a_id = Column(ForeignKey("a.id"))
    cs = relationship("C")

class C(Base):
    __tablename__ = "c"
    id = Column(Integer, primary_key=True)
    b_id = Column(ForeignKey("b.id"))

partition = select(
    B,
    func.row_number()
    .over(order_by=B.id, partition_by=B.a_id)
    .label("index"),
).alias()

partitioned_b = aliased(B, alias=partition)

A.partitioned_bs = relationship(
    partitioned_b,
    primaryjoin=and_(
        partitioned_b.a_id == A.id, partition.c.index < 10
    ),
)

e = create_engine("sqlite://", echo=True)
Base.metadata.create_all(e)

s = Session(e)
s.add_all([A(id=i) for i in range(1, 4)])
s.commit()
s.add_all(
    [
        B(a_id=i, cs=[C(), C()])
        for i in range(1, 4)
        for j in range(1, 21)
    ]
)
s.commit()

s.close()
# works
for a1 in s.query(A).options(
    selectinload(A.partitioned_bs).joinedload("cs")
):
    for b in a1.partitioned_bs:
        assert len(b.cs) == 2

s.close()

# fails on the first query due to the way the path comes out
for a1 in s.query(A).options(
    selectinload(A.partitioned_bs).joinedload(partitioned_b.cs)
):
    for b in a1.partitioned_bs:
        assert len(b.cs) == 2

# works with:
selectinload(A.partitioned_bs.of_type(partitioned_b)).joinedload(partitioned_b.cs)

# this should raise, silently fails
selectinload(A.partitioned_bs).joinedload(B.cs)