Refactored AutoPrefetchRelated mixin

Author: Anton-Shutik

example/tests/test_performance.py

@@ -53,7 +53,7 @@ def test_query_count_include_author(self):
         4. Author types prefetched
         5. Entries prefetched
         """
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.get('/comments?include=author&page[size]=25')
             self.assertEqual(len(response.data['results']), 25)
 

rest_framework_json_api/views.py

@@ -98,22 +98,25 @@ def get_select_related(self, include):
     def get_prefetch_related(self, include):
         return getattr(self, 'prefetch_for_includes', {}).get(include, None)
 
+    def get_queryset_with_related(self, qs, included):
+        select_related = self.get_select_related(included)
+        if select_related is not None:
+            qs = qs.select_related(*select_related)
+
+        prefetch_related = self.get_prefetch_related(included)
+        if prefetch_related is not None:
+            qs = qs.prefetch_related(*prefetch_related)
+
+        return qs
+
     def get_queryset(self, *args, **kwargs):
         """ This mixin adds automatic prefetching for OneToOne and ManyToMany fields. """
         qs = super(AutoPreloadMixin, self).get_queryset(*args, **kwargs)
         included_resources = get_included_resources(self.request)
 
         for included in included_resources + ['__all__']:
             # Custom defined "select_related" and "prefetch_related" is a priority
-            select_related = self.get_select_related(included)
-            if select_related is not None:
-                qs = qs.select_related(*select_related)
-                continue
-
-            prefetch_related = self.get_prefetch_related(included)
-            if prefetch_related is not None:
-                qs = qs.prefetch_related(*prefetch_related)
-                continue
+            qs = self.get_queryset_with_related(qs, included)
 
             # If include was not defined, trying to resolve it automatically
             included_model = None