From 2d2c5ee88939a929c9f4d1c15beb39ba76215060 Mon Sep 17 00:00:00 2001
From: Jacob Foster <jacob@britecore.com>
Date: Tue, 18 Jul 2017 16:30:00 -0500
Subject: [PATCH] Adding query optimization

---
 graphene_django/filter/fields.py |  2 +
 graphene_django/optimization.py  | 82 ++++++++++++++++++++++++++++++++
 graphene_django/types.py         |  6 ++-
 3 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 graphene_django/optimization.py

diff --git a/graphene_django/filter/fields.py b/graphene_django/filter/fields.py
index fc414bf..fb0d9f9 100644
--- a/graphene_django/filter/fields.py
+++ b/graphene_django/filter/fields.py
@@ -4,6 +4,7 @@ from functools import partial
 # from graphene.relay import is_node
 from graphene.types.argument import to_arguments
 from ..fields import DjangoConnectionField
+from ..optimization import optimize_queryset
 from .utils import get_filtering_args_from_filterset, get_filterset_class
 
 
@@ -75,6 +76,7 @@ class DjangoFilterConnectionField(DjangoConnectionField):
             data=filter_kwargs,
             queryset=default_manager.get_queryset()
         ).qs
+        qs = optimize_queryset(default_manager.model, qs, info.field_asts[0])
 
         return super(DjangoFilterConnectionField, cls).connection_resolver(
             resolver,
diff --git a/graphene_django/optimization.py b/graphene_django/optimization.py
new file mode 100644
index 0000000..f66d387
--- /dev/null
+++ b/graphene_django/optimization.py
@@ -0,0 +1,82 @@
+from collections import namedtuple
+
+from django.db.models import ForeignKey
+from django.db.models.fields.reverse_related import ForeignObjectRel
+from graphene.utils.str_converters import to_snake_case
+
+from .registry import get_global_registry
+from .utils import get_related_model
+
+REGISTRY = get_global_registry()
+SELECT = 'select'
+PREFETCH = 'prefetch'
+RelatedSelection = namedtuple('RelatedSelection', ['name', 'fetch_type'])
+
+
+def model_fields_as_dict(model):
+    return dict((f.name, f) for f in model._meta.get_fields())
+
+
+def get_related_fetches_for_model(model, graphql_ast):
+    model_fields = model_fields_as_dict(model)
+    selections = graphql_ast.selection_set.selections
+
+    graphene_obj_type = REGISTRY.get_type_for_model(model)
+    optimizations = {}
+    if graphene_obj_type and graphene_obj_type._meta.optimizations:
+        optimizations = graphene_obj_type._meta.optimizations
+
+    relateds = []
+
+    for selection in selections:
+        selection_name = to_snake_case(selection.name.value)
+        selection_field = model_fields.get(selection_name, None)
+
+        try:
+            related_model = get_related_model(selection_field)
+        except:
+            # This is not a ForeignKey or Relation, check manual optimizations
+            manual_optimizations = optimizations.get(selection_name)
+            if manual_optimizations:
+                for manual_select in manual_optimizations.get(SELECT, []):
+                    relateds.append(RelatedSelection(manual_select, SELECT))
+                for manual_prefetch in manual_optimizations.get(PREFETCH, []):
+                    relateds.append(RelatedSelection(manual_prefetch, PREFETCH))
+
+            continue
+
+        query_name = selection_field.name
+        if isinstance(selection_field, ForeignObjectRel):
+            query_name = selection_field.field.related_query_name()
+
+        nested_relateds = get_related_fetches_for_model(related_model, selection)
+
+        related_type = PREFETCH  # default to prefetch, it's safer
+        if isinstance(selection_field, ForeignKey):
+            related_type = SELECT  # we can only select for ForeignKeys
+
+        if nested_relateds:
+            for related in nested_relateds:
+                full_name = '{0}__{1}'.format(query_name, related.name)
+
+                nested_related_type = PREFETCH
+                if related_type == SELECT and related.fetch_type == SELECT:
+                    nested_related_type = related_type
+
+                relateds.append(RelatedSelection(full_name, nested_related_type))
+        else:
+            relateds.append(RelatedSelection(query_name, related_type))
+
+    return relateds
+
+
+def optimize_queryset(model, queryset, graphql_ast):
+    relateds = get_related_fetches_for_model(model, graphql_ast)
+
+    for related in relateds:
+        if related.fetch_type == SELECT:
+            queryset = queryset.select_related(related.name)
+        else:
+            queryset = queryset.prefetch_related(related.name)
+
+    return queryset
diff --git a/graphene_django/types.py b/graphene_django/types.py
index bb0a2f1..1d9b2a6 100644
--- a/graphene_django/types.py
+++ b/graphene_django/types.py
@@ -10,6 +10,7 @@ from graphene.types.utils import merge, yank_fields_from_attrs
 from graphene.utils.is_base_type import is_base_type
 
 from .converter import convert_django_field_with_choices
+from .optimization import optimize_queryset
 from .registry import Registry, get_global_registry
 from .utils import (DJANGO_FILTER_INSTALLED, get_model_fields,
                     is_valid_django_model)
@@ -55,6 +56,7 @@ class DjangoObjectTypeMeta(ObjectTypeMeta):
             only_fields=(),
             exclude_fields=(),
             interfaces=(),
+            optimizations=None,
             skip_registry=False,
             registry=None
         )
@@ -118,7 +120,9 @@ class DjangoObjectType(six.with_metaclass(DjangoObjectTypeMeta, ObjectType)):
 
     @classmethod
     def get_node(cls, id, context, info):
+        query = cls._meta.model._meta.default_manager
+        query = optimize_queryset(cls._meta.model, query, info.field_asts[0])
         try:
-            return cls._meta.model.objects.get(pk=id)
+            return query.get(pk=id)
         except cls._meta.model.DoesNotExist:
             return None