From 65a36c8da8bc06202d3bf2d326a6e9c869895184 Mon Sep 17 00:00:00 2001 From: Alexander Kozlovsky Date: Sat, 19 Jan 2019 16:57:15 +0300 Subject: [PATCH] Fixes 414: prefetching Optional relationships fails on 0.7.7 --- pony/orm/core.py | 2 +- pony/orm/tests/test_prefetching.py | 43 +++++++++++++++++++++++++++--- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/pony/orm/core.py b/pony/orm/core.py index 068776a0f..69d7a50cf 100644 --- a/pony/orm/core.py +++ b/pony/orm/core.py @@ -5843,7 +5843,7 @@ def _do_prefetch(query, query_result): collection_prefetch_dict[attr].add(obj) else: obj2 = attr.get(obj) - if obj2 not in all_objects: + if obj2 is not None and obj2 not in all_objects: all_objects.add(obj2) objects_to_prefetch.add(obj2) diff --git a/pony/orm/tests/test_prefetching.py b/pony/orm/tests/test_prefetching.py index d7afb54b4..68509a72f 100644 --- a/pony/orm/tests/test_prefetching.py +++ b/pony/orm/tests/test_prefetching.py @@ -14,6 +14,7 @@ class Student(db.Entity): dob = Optional(date) group = Required('Group') courses = Set('Course') + mentor = Optional('Teacher') biography = Optional(LongStr) class Group(db.Entity): @@ -25,6 +26,10 @@ class Course(db.Entity): name = Required(str, unique=True) students = Set(Student) +class Teacher(db.Entity): + name = Required(str) + students = Set(Student) + db.generate_mapping(create_tables=True) with db_session: @@ -33,10 +38,12 @@ class Course(db.Entity): c1 = Course(name='Math') c2 = Course(name='Physics') c3 = Course(name='Computer Science') - Student(id=1, name='S1', group=g1, gpa=3.1, courses=[c1, c2], biography='S1 bio') + t1 = Teacher(name='T1') + t2 = Teacher(name='T2') + Student(id=1, name='S1', group=g1, gpa=3.1, courses=[c1, c2], biography='S1 bio', mentor=t1) Student(id=2, name='S2', group=g1, gpa=4.2, scholarship=100, dob=date(2000, 1, 1), biography='S2 bio') Student(id=3, name='S3', group=g1, gpa=4.7, scholarship=200, dob=date(2001, 1, 2), courses=[c2, c3]) - Student(id=4, name='S4', group=g2, gpa=3.2, biography='S4 bio', courses=[c1, c3]) + Student(id=4, name='S4', group=g2, gpa=3.2, biography='S4 bio', courses=[c1, c3], mentor=t2) Student(id=5, name='S5', group=g2, gpa=4.5, biography='S5 bio', courses=[c1, c3]) class TestPrefetching(unittest.TestCase): @@ -108,7 +115,8 @@ def test_12(self): with db_session: s1 = Student.select().prefetch(Student.biography).first() self.assertEqual(s1.biography, 'S1 bio') - self.assertEqual(db.last_sql, '''SELECT "s"."id", "s"."name", "s"."scholarship", "s"."gpa", "s"."dob", "s"."group", "s"."biography" + self.assertEqual(db.last_sql, +'''SELECT "s"."id", "s"."name", "s"."scholarship", "s"."gpa", "s"."dob", "s"."group", "s"."mentor", "s"."biography" FROM "Student" "s" ORDER BY 1 LIMIT 1''') @@ -167,6 +175,35 @@ def test_17(self): query_count = db.local_stats[None].db_count self.assertEqual(query_count, 4) + def test_18(self): + db.merge_local_stats() + with db_session: + q = Group.select().prefetch(Group.students, Student.biography) + for g in q: # 2 queries + for s in g.students: + m = s.mentor # 0 queries + b = s.biography # 0 queries + query_count = db.local_stats[None].db_count + self.assertEqual(query_count, 2) + + def test_19(self): + db.merge_local_stats() + with db_session: + q = Group.select().prefetch(Group.students, Student.biography, Student.mentor) + mentors = set() + for g in q: # 3 queries + for s in g.students: + m = s.mentor # 0 queries + if m is not None: + mentors.add(m) + b = s.biography # 0 queries + query_count = db.local_stats[None].db_count + self.assertEqual(query_count, 3) + + for m in mentors: + n = m.name # 0 queries + query_count = db.local_stats[None].db_count + self.assertEqual(query_count, 3) if __name__ == '__main__':