Allow defining of select_related per include

CHANGELOG.md

@@ -16,6 +16,7 @@ any parts of the framework not mentioned in the documentation should generally b
 
 ### Fixed
 
+* Added `select_for_includes` handling. [PR](https://github.com/django-json-api/django-rest-framework-json-api/pull/600). [Docs](https://django-rest-framework-json-api.readthedocs.io/en/stable/usage.html#performance-improvements).
 * Avoid exception when trying to include skipped relationship
 * Don't swallow `filter[]` params when there are several
 * Fix DeprecationWarning regarding collections.abc import in Python 3.7

docs/usage.md

@@ -823,17 +823,22 @@ class QuestSerializer(serializers.ModelSerializer):
 
 Be aware that using included resources without any form of prefetching **WILL HURT PERFORMANCE** as it will introduce m\*(n+1) queries.
 
-A viewset helper was designed to allow for greater flexibility and it is automatically available when subclassing
+A viewset helper was designed to allow for greater flexibility and it is automatically available when subclassing.
+You can also define your custom queryset for `select` or `prefetch` related for each `include` that comes from the url.
+It has a priority over automatically added preloads.
 `rest_framework_json_api.views.ModelViewSet`:
 ```python
 from rest_framework_json_api import views
 
 # When MyViewSet is called with ?include=author it will dynamically prefetch author and author.bio
 class MyViewSet(views.ModelViewSet):
     queryset = Book.objects.all()
+    select_for_includes = {
+        'author': ['author__bio'],
+    }
     prefetch_for_includes = {
         '__all__': [],
-        'author': ['author', 'author__bio'],
+        'all_authors': [Prefetch('all_authors', queryset=Author.objects.select_related('bio'))],
         'category.section': ['category']
     }
 ```

example/tests/test_performance.py

@@ -53,6 +53,16 @@ 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)
+
+    def test_query_select_related_entry(self):
+        """ We expect a list view with an include have two queries:
+
+        1. Primary resource COUNT query
+        2. Primary resource SELECT + SELECT RELATED writer(author) and bio
+        """
+        with self.assertNumQueries(2):
+            response = self.client.get('/comments?include=writer&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 (
@@ -184,6 +184,9 @@ class AuthorViewSet(ModelViewSet):
 class CommentViewSet(ModelViewSet):
     queryset = Comment.objects.all()
     serializer_class = CommentSerializer
+    select_for_includes = {
+        'writer': ['author__bio']
+    }
     prefetch_for_includes = {
         '__all__': [],
         'author': ['author__bio', 'author__entries'],
@@ -197,9 +200,12 @@ 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 = {
+        'current_project': ['current_project'],
+    }
 
 
 class ProjectViewset(ModelViewSet):

rest_framework_json_api/views.py

@@ -1,3 +1,4 @@
+import warnings
 
 from django.core.exceptions import ImproperlyConfigured
 from django.db.models import Model
@@ -31,6 +32,13 @@
 
 
 class PrefetchForIncludesHelperMixin(object):
+
+    def __init__(self, *args, **kwargs):
+        warnings.warn("PrefetchForIncludesHelperMixin is deprecated. "
+                      "Use PreloadIncludesMixin instead",
+                      DeprecationWarning)
+        super(PrefetchForIncludesHelperMixin, self).__init__(*args, **kwargs)
+
     def get_queryset(self):
         """
         This viewset provides a helper attribute to prefetch related models
@@ -62,33 +70,86 @@ class MyViewSet(viewsets.ModelViewSet):
         return qs
 
 
-class AutoPrefetchMixin(object):
+class PreloadIncludesMixin(object):
+    """
+    This mixin provides a helper attributes to select or prefetch related models
+    based on the include specified in the URL.
+
+    __all__ can be used to specify a prefetch which should be done regardless of the include
+
+    .. code:: python
+
+    # When MyViewSet is called with ?include=author it will prefetch author and authorbio
+    class MyViewSet(viewsets.ModelViewSet):
+        queryset = Book.objects.all()
+        prefetch_for_includes = {
+            '__all__': [],
+            'category.section': ['category']
+        }
+        select_for_includes = {
+            '__all__': [],
+            'author': ['author', 'author__authorbio'],
+        }
+    """
+
+    def get_select_related(self, include):
+        return getattr(self, 'select_for_includes', {}).get(include, None)
+
+    def get_prefetch_related(self, include):
+        return getattr(self, 'prefetch_for_includes', {}).get(include, None)
+
+    def get_queryset(self, *args, **kwargs):
+        qs = super(PreloadIncludesMixin, self).get_queryset(*args, **kwargs)
+
+        included_resources = get_included_resources(self.request)
+        for included in included_resources + ['__all__']:
+
+            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
+
+
+class AutoPreloadMixin(object):
+
     def get_queryset(self, *args, **kwargs):
         """ This mixin adds automatic prefetching for OneToOne and ManyToMany fields. """
-        qs = super(AutoPrefetchMixin, self).get_queryset(*args, **kwargs)
+        qs = super(AutoPreloadMixin, self).get_queryset(*args, **kwargs)
         included_resources = get_included_resources(self.request)
 
-        for included in included_resources:
+        for included in included_resources + ['__all__']:
+            # If include was not defined, trying to resolve it automatically
             included_model = None
             levels = included.split('.')
             level_model = qs.model
+            # Suppose we can do select_related by default
+            can_select_related = True
             for level in levels:
                 if not hasattr(level_model, level):
                     break
                 field = getattr(level_model, level)
                 field_class = field.__class__
 
                 is_forward_relation = (
-                    issubclass(field_class, ForwardManyToOneDescriptor) or
-                    issubclass(field_class, ManyToManyDescriptor)
+                    issubclass(field_class, (ForwardManyToOneDescriptor, ManyToManyDescriptor))
                 )
                 is_reverse_relation = (
-                    issubclass(field_class, ReverseManyToOneDescriptor) or
-                    issubclass(field_class, ReverseOneToOneDescriptor)
+                    issubclass(field_class, (ReverseManyToOneDescriptor, ReverseOneToOneDescriptor))
                 )
                 if not (is_forward_relation or is_reverse_relation):
                     break
 
+                # Figuring out if relation should be select related rather than prefetch_related
+                # If at least one relation in the chain is not "selectable" then use "prefetch"
+                can_select_related &= (
+                    issubclass(field_class, (ForwardManyToOneDescriptor, ReverseOneToOneDescriptor))
+                )
+
                 if level == levels[-1]:
                     included_model = field
                 else:
@@ -104,11 +165,23 @@ def get_queryset(self, *args, **kwargs):
                         level_model = model_field.model
 
             if included_model is not None:
-                qs = qs.prefetch_related(included.replace('.', '__'))
+                if can_select_related:
+                    qs = qs.select_related(included.replace('.', '__'))
+                else:
+                    qs = qs.prefetch_related(included.replace('.', '__'))
 
         return qs
 
 
+class AutoPrefetchMixin(AutoPreloadMixin):
+
+    def __init__(self, *args, **kwargs):
+        warnings.warn("AutoPrefetchMixin is deprecated. "
+                      "Use AutoPreloadMixin instead",
+                      DeprecationWarning)
+        super(AutoPrefetchMixin, self).__init__(*args, **kwargs)
+
+
 class RelatedMixin(object):
     """
     This mixin handles all related entities, whose Serializers are declared in "related_serializers"
@@ -186,15 +259,14 @@ def get_related_instance(self):
                 raise NotFound
 
 
-class ModelViewSet(AutoPrefetchMixin,
-                   PrefetchForIncludesHelperMixin,
+class ModelViewSet(AutoPreloadMixin,
+                   PreloadIncludesMixin,
                    RelatedMixin,
                    viewsets.ModelViewSet):
     pass
 
 
-class ReadOnlyModelViewSet(AutoPrefetchMixin,
-                           PrefetchForIncludesHelperMixin,
+class ReadOnlyModelViewSet(AutoPreloadMixin,
                            RelatedMixin,
                            viewsets.ReadOnlyModelViewSet):
     pass