Merge pull request #36 from vintasoftware/feat/support-nested-prefetch

Add support to nested prefetch lookups

django_virtual_models/fields.py

@@ -319,10 +319,11 @@ def hydrate_queryset(
         # always include the "back reference" field name in the Prefetch's lookup list
         # to avoid N+1s in internal Django prefetch code
         field_to_prefetch = self.lookup if self.lookup else self.field_name
-        field_descriptor = getattr(self.parent.model_cls, field_to_prefetch)
-        if type(field_descriptor) == ReverseManyToOneDescriptor:  # don't use isinstance
-            back_reference = field_descriptor.rel.field.name
-            new_lookup_list.append(back_reference)
+        if "__" not in field_to_prefetch:
+            field_descriptor = getattr(self.parent.model_cls, field_to_prefetch)
+            if type(field_descriptor) == ReverseManyToOneDescriptor:  # don't use isinstance
+                back_reference = field_descriptor.rel.field.name
+                new_lookup_list.append(back_reference)
 
         # defer fields on prefetch_queryset
         prefetch_queryset = _defer_fields(

tests/optimization/test_exceptions.py

@@ -41,7 +41,7 @@ class Meta:
     def get_lesson_titles_from_method(self, course: Annotated[Course, hints.Virtual("lessons")]):
         ...
 
-    @v.hints.from_types_of(get_lesson_title_list, "course")
+    @hints.from_types_of(get_lesson_title_list, "course")
     def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
         ...
 
@@ -166,7 +166,7 @@ def get_lesson_title_list(course):  # no type annotation here
             ...
 
         class BrokenCourseSerializer(BaseCourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
                 ...
 
@@ -189,7 +189,7 @@ def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
 
     def test_function_with_wrong_param_name_raises_exception(self):
         class BrokenCourseSerializer(BaseCourseSerializer):
-            @v.hints.from_types_of(
+            @hints.from_types_of(
                 get_lesson_title_list, obj_param_name="course_blabla"  # wrong param name
             )
             def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
@@ -219,7 +219,7 @@ def get_lesson_title_list(course: Course):  # no Annotated here
             ...
 
         class BrokenCourseSerializer(BaseCourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
                 ...
 
@@ -244,7 +244,7 @@ def get_lesson_title_list(course: Annotated[Course, ()]):  # wrong Annotated her
             ...
 
         class BrokenCourseSerializer(BaseCourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
                 ...
 
@@ -276,7 +276,7 @@ def get_lesson_title_list(
             ...
 
         class BrokenCourseSerializer(BaseCourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles_from_function(self, course, get_lesson_title_list_helper):
                 ...
 

tests/optimization/test_lookup_finder.py

@@ -63,6 +63,18 @@ class Meta:
         deferred_fields = ["description"]
 
 
+class NestedFacilitatorUsers(v.VirtualModel):
+    class Meta:
+        model = User
+
+
+class VirtualLesson(v.VirtualModel):
+    facilitator_users = NestedFacilitatorUsers(lookup="course__facilitators")
+
+    class Meta:
+        model = Lesson
+
+
 class NestedAssignmentSerializer(serializers.ModelSerializer):
     email = serializers.EmailField()
     lessons_total = serializers.IntegerField()
@@ -149,11 +161,33 @@ def get_user_assignment(self, obj, serializer_cls):
             return serializer_cls(obj.user_assignment[0]).data
         return None
 
-    @v.hints.from_types_of(get_lesson_title_list, "course")
+    @hints.from_types_of(get_lesson_title_list, "course")
     def get_lesson_titles(self, course, get_lesson_title_list_helper):
         return get_lesson_title_list_helper(course)
 
 
+class LessonSerializer(serializers.ModelSerializer):
+    facilitator_users = serializers.SerializerMethodField()
+
+    class Meta:
+        model = Lesson
+        virtual_model = VirtualLesson
+        fields = [
+            "title",
+            "content",
+            "facilitator_users",
+        ]
+
+    @hints.defined_on_virtual_model()
+    def get_facilitator_users(self, lesson):
+        if hasattr(lesson, "facilitator_users"):
+            return list({u.email for u in lesson.facilitator_users})
+
+        # this won't run because it's defined on virtual model,
+        # but one could add fallback code here:
+        return None
+
+
 class LookupFinderTests(TestCase):
     def setUp(self):
         super().setUp()
@@ -264,6 +298,21 @@ class Meta:
             ]
         )
 
+    def test_prefetch_with_nested_lookup(self):
+        qs = Lesson.objects.all()
+        serializer_instance = LessonSerializer(instance=qs, many=True)
+        virtual_model = VirtualLesson()
+
+        lookup_list = LookupFinder(
+            serializer_instance=serializer_instance,
+            virtual_model=virtual_model,
+        ).recursively_find_lookup_list()
+
+        optimized_qs = virtual_model.get_optimized_queryset(qs=qs, lookup_list=lookup_list)
+        with self.assertNumQueries(3):
+            lesson_list = list(optimized_qs)
+            assert len(lesson_list) == 9
+
     def test_ignored_nested_serializer_with_noop(self):
         """
         Sometimes one needs a nested serializer generated dynamically.
@@ -375,7 +424,7 @@ class Meta:
     @override_settings(DEBUG=True)
     def test_prefetch_hints_block_queries_on_serializer_evaluation(self):
         class BrokenCourseSerializer(CourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles(self, course, get_lesson_title_list_helper):
                 list(course.lessons.order_by("title"))  # new query
 
@@ -398,7 +447,7 @@ def get_lesson_titles(self, course, get_lesson_title_list_helper):
     @override_settings(DEBUG=True)
     def test_prefetch_hints_does_not_block_queries_if_false(self):
         class BrokenCourseSerializer(CourseSerializer):
-            @v.hints.from_types_of(get_lesson_title_list, "course")
+            @hints.from_types_of(get_lesson_title_list, "course")
             def get_lesson_titles(self, course, get_lesson_title_list_helper):
                 list(course.lessons.order_by("title"))  # new query