from sqlalchemy import and_ from sqlalchemy import Column from sqlalchemy import ForeignKey from sqlalchemy import func from sqlalchemy import Integer from sqlalchemy import join from sqlalchemy import select from sqlalchemy import testing from sqlalchemy.orm import aliased from sqlalchemy.orm import joinedload from sqlalchemy.orm import noload from sqlalchemy.orm import relationship from sqlalchemy.orm import selectinload from sqlalchemy.orm import Session from sqlalchemy.testing import eq_ from sqlalchemy.testing import fixtures from sqlalchemy.testing.assertsql import CompiledSQL from sqlalchemy.testing.fixtures import ComparableEntity class PartitionByFixture(fixtures.DeclarativeMappedTest): @classmethod def setup_classes(cls): Base = cls.DeclarativeBasic 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 ), ) @classmethod def insert_data(cls, connection): A, B, C = cls.classes("A", "B", "C") s = Session(connection) s.add_all([A(id=i) for i in range(1, 4)]) s.flush() s.add_all( [ B(a_id=i, cs=[C(), C()]) for i in range(1, 4) for j in range(1, 21) ] ) s.commit() class AliasedClassRelationshipTest( PartitionByFixture, testing.AssertsCompiledSQL ): # TODO: maybe make this more backend agnostic __requires__ = ("window_functions",) __dialect__ = "default" def test_lazyload(self): A, B, C = self.classes("A", "B", "C") s = Session(testing.db) def go(): for a1 in s.query(A): # 1 query eq_(len(a1.partitioned_bs), 9) # 3 queries for b in a1.partitioned_bs: eq_(len(b.cs), 2) # 9 * 3 = 27 queries self.assert_sql_count(testing.db, go, 31) def test_join_one(self): A, B, C = self.classes("A", "B", "C") s = Session(testing.db) q = s.query(A).join(A.partitioned_bs) self.assert_compile( q, "SELECT a.id AS a_id FROM a JOIN " "(SELECT b.id AS id, b.a_id AS a_id, row_number() " "OVER (PARTITION BY b.a_id ORDER BY b.id) " "AS index FROM b) AS anon_1 " "ON anon_1.a_id = a.id AND anon_1.index < :index_1", ) def test_join_two(self): A, B, C = self.classes("A", "B", "C") s = Session(testing.db) q = s.query(A, A.partitioned_bs.entity).join(A.partitioned_bs) self.assert_compile( q, "SELECT a.id AS a_id, anon_1.id AS anon_1_id, " "anon_1.a_id AS anon_1_a_id " "FROM a JOIN " "(SELECT b.id AS id, b.a_id AS a_id, row_number() " "OVER (PARTITION BY b.a_id ORDER BY b.id) " "AS index FROM b) AS anon_1 " "ON anon_1.a_id = a.id AND anon_1.index < :index_1", ) def test_selectinload_w_noload_after(self): A, B, C = self.classes("A", "B", "C") s = Session(testing.db) def go(): for a1 in s.query(A).options( noload("*"), selectinload(A.partitioned_bs) ): for b in a1.partitioned_bs: eq_(b.cs, []) self.assert_sql_count(testing.db, go, 2) def test_selectinload_w_joinedload_after(self): A, B, C = self.classes("A", "B", "C") s = Session(testing.db) def go(): for a1 in s.query(A).options( selectinload(A.partitioned_bs).joinedload("cs") ): for b in a1.partitioned_bs: eq_(len(b.cs), 2) self.assert_sql_count(testing.db, go, 2) class AltSelectableTest( fixtures.DeclarativeMappedTest, testing.AssertsCompiledSQL ): __dialect__ = "default" @classmethod def setup_classes(cls): Base = cls.DeclarativeBasic class A(ComparableEntity, Base): __tablename__ = "a" id = Column(Integer, primary_key=True) b_id = Column(ForeignKey("b.id")) class B(ComparableEntity, Base): __tablename__ = "b" id = Column(Integer, primary_key=True) class C(ComparableEntity, Base): __tablename__ = "c" id = Column(Integer, primary_key=True) a_id = Column(ForeignKey("a.id")) class D(ComparableEntity, Base): __tablename__ = "d" id = Column(Integer, primary_key=True) c_id = Column(ForeignKey("c.id")) b_id = Column(ForeignKey("b.id")) # 1. set up the join() as a variable, so we can refer # to it in the mapping multiple times. j = join(B, D, D.b_id == B.id).join(C, C.id == D.c_id) # 2. Create an AliasedClass to B B_viacd = aliased(B, j, flat=True) A.b = relationship(B_viacd, primaryjoin=A.b_id == j.c.b_id) @classmethod def insert_data(cls, connection): A, B, C, D = cls.classes("A", "B", "C", "D") sess = Session(connection) for obj in [ B(id=1), A(id=1, b_id=1), C(id=1, a_id=1), D(id=1, c_id=1, b_id=1), ]: sess.add(obj) sess.flush() sess.commit() def test_lazyload(self): A, B = self.classes("A", "B") sess = Session() a1 = sess.query(A).first() with self.sql_execution_asserter() as asserter: # note this is many-to-one. use_get is unconditionally turned # off for relationship to aliased class for now. eq_(a1.b, B(id=1)) asserter.assert_( CompiledSQL( "SELECT b.id AS b_id FROM b JOIN d ON d.b_id = b.id " "JOIN c ON c.id = d.c_id WHERE :param_1 = b.id", [{"param_1": 1}], ) ) def test_joinedload(self): A, B = self.classes("A", "B") sess = Session() with self.sql_execution_asserter() as asserter: # note this is many-to-one. use_get is unconditionally turned # off for relationship to aliased class for now. a1 = sess.query(A).options(joinedload(A.b)).first() eq_(a1.b, B(id=1)) asserter.assert_( CompiledSQL( "SELECT a.id AS a_id, a.b_id AS a_b_id, b_1.id AS b_1_id " "FROM a LEFT OUTER JOIN (b AS b_1 " "JOIN d AS d_1 ON d_1.b_id = b_1.id " "JOIN c AS c_1 ON c_1.id = d_1.c_id) ON a.b_id = b_1.id " "LIMIT :param_1", [{"param_1": 1}], ) ) def test_selectinload(self): A, B = self.classes("A", "B") sess = Session() with self.sql_execution_asserter() as asserter: # note this is many-to-one. use_get is unconditionally turned # off for relationship to aliased class for now. a1 = sess.query(A).options(selectinload(A.b)).first() eq_(a1.b, B(id=1)) asserter.assert_( CompiledSQL( "SELECT a.id AS a_id, a.b_id AS a_b_id " "FROM a LIMIT :param_1", [{"param_1": 1}], ), CompiledSQL( "SELECT a_1.id AS a_1_id, b.id AS b_id FROM a AS a_1 " "JOIN (b JOIN d ON d.b_id = b.id JOIN c ON c.id = d.c_id) " "ON a_1.b_id = b.id WHERE a_1.id " "IN ([POSTCOMPILE_primary_keys])", [{"primary_keys": [1]}], ), ) def test_join(self): A, B = self.classes("A", "B") sess = Session() self.assert_compile( sess.query(A).join(A.b), "SELECT a.id AS a_id, a.b_id AS a_b_id " "FROM a JOIN (b JOIN d ON d.b_id = b.id " "JOIN c ON c.id = d.c_id) ON a.b_id = b.id", )