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)
 

example/views.py

@@ -12,7 +12,7 @@
 from rest_framework_json_api.filters import OrderingFilter, QueryParameterValidationFilter
 from rest_framework_json_api.pagination import JsonApiPageNumberPagination
 from rest_framework_json_api.utils import format_drf_errors
-from rest_framework_json_api.views import ModelViewSet, RelationshipView
+from rest_framework_json_api.views import ModelViewSet, RelationshipView, PreloadIncludesMixin
 
 from example.models import Author, Blog, Comment, Company, Entry, Project, ProjectType
 from example.serializers import (
@@ -200,7 +200,7 @@ def get_queryset(self, *args, **kwargs):
         return super(CommentViewSet, self).get_queryset()
 
 
-class CompanyViewset(ModelViewSet):
+class CompanyViewset(PreloadIncludesMixin, viewsets.ModelViewSet):
     queryset = Company.objects.all()
     serializer_class = CompanySerializer
     prefetch_for_includes = {

rest_framework_json_api/views.py

@@ -35,7 +35,7 @@ class PrefetchForIncludesHelperMixin(object):
 
     def __init__(self, *args, **kwargs):
         warnings.warn("PrefetchForIncludesHelperMixin is deprecated. "
-                      "Use AutoPreloadMixin instead",
+                      "Use PreloadIncludesMixin instead",
                       DeprecationWarning)
         super(PrefetchForIncludesHelperMixin, self).__init__(*args, **kwargs)
 
@@ -70,7 +70,7 @@ class MyViewSet(viewsets.ModelViewSet):
         return qs
 
 
-class AutoPreloadMixin(object):
+class PreloadIncludesMixin(object):
     """
     This mixin provides a helper attributes to select or prefetch related models
     based on the include specified in the URL.
@@ -99,22 +99,30 @@ def get_prefetch_related(self, include):
         return getattr(self, 'prefetch_for_includes', {}).get(include, None)
 
     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)
+        qs = super(PreloadIncludesMixin, 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
 
+        return qs
+
+
+class AutoPreloadMixin(object):
+
+    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__']:
             # If include was not defined, trying to resolve it automatically
             included_model = None
             levels = included.split('.')
@@ -252,6 +260,7 @@ def get_related_instance(self):
 
 
 class ModelViewSet(AutoPreloadMixin,
+                   PreloadIncludesMixin,
                    RelatedMixin,
                    viewsets.ModelViewSet):
     pass