From 25d680fa007c2ab7113067080f7a12e8395b5890 Mon Sep 17 00:00:00 2001
From: Syrus Akbary <me@syrusakbary.com>
Date: Tue, 19 Jan 2016 20:49:43 -0800
Subject: [PATCH] Initial optimize in connections

---
 graphene/contrib/django/fields.py             |  7 ++++-
 .../django/tests/test_optimizequery.py        | 26 +++++++++++++++++--
 graphene/contrib/django/utils.py              |  4 ++-
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/graphene/contrib/django/fields.py b/graphene/contrib/django/fields.py
index 8eee0736..8e2282a7 100644
--- a/graphene/contrib/django/fields.py
+++ b/graphene/contrib/django/fields.py
@@ -4,6 +4,7 @@ from ...core.types.base import FieldType
 from ...core.types.definitions import List
 from ...relay import ConnectionField
 from ...relay.utils import is_node
+from .fetcher import get_fields
 from .utils import DJANGO_FILTER_INSTALLED, get_type_for_model, maybe_queryset
 
 
@@ -37,10 +38,14 @@ class DjangoConnectionField(DjangoField, ConnectionField):
     def get_queryset(self, resolved_qs, args, info):
         return resolved_qs
 
+    def get_fetch_fields(self, info):
+        return get_fields(info)
+
     def from_list(self, connection_type, resolved, args, info):
         if not resolved:
             resolved = self.get_manager()
-        resolved_qs = maybe_queryset(resolved)
+        only = self.get_fetch_fields(info)
+        resolved_qs = maybe_queryset(resolved, only)
         qs = self.get_queryset(resolved_qs, args, info)
         return super(DjangoConnectionField, self).from_list(connection_type, qs, args, info)
 
diff --git a/graphene/contrib/django/tests/test_optimizequery.py b/graphene/contrib/django/tests/test_optimizequery.py
index 93fe7f18..3676460a 100644
--- a/graphene/contrib/django/tests/test_optimizequery.py
+++ b/graphene/contrib/django/tests/test_optimizequery.py
@@ -1,11 +1,12 @@
 import pytest
 
 import graphene
-from graphene.contrib.django import DjangoObjectType
+from graphene.contrib.django import DjangoNode
 
 from ..tests.models import Reporter
 from ..debug.plugin import DjangoDebugPlugin
 from ..fetcher import fetch_only_required, get_fields
+from ..fields import DjangoConnectionField
 
 # from examples.starwars_django.models import Character
 
@@ -18,7 +19,7 @@ def test_should_query_well():
     r2 = Reporter(last_name='Griffin')
     r2.save()
 
-    class ReporterType(DjangoObjectType):
+    class ReporterType(DjangoNode):
 
         class Meta:
             model = Reporter
@@ -26,6 +27,7 @@ def test_should_query_well():
     class Query(graphene.ObjectType):
         reporter = graphene.Field(ReporterType)
         all_reporters = ReporterType.List()
+        all_reporters_connection = DjangoConnectionField(ReporterType)
 
         @fetch_only_required
         def resolve_all_reporters(self, args, info):
@@ -40,6 +42,14 @@ def test_should_query_well():
             lastName
             email
           }
+          allReportersConnection(first:1) {
+            edges {
+              node {
+                lastName
+                email
+              }
+            }
+          }
           reporter {
             email
           }
@@ -58,12 +68,24 @@ def test_should_query_well():
             'lastName': 'Griffin',
             'email': '',
         }],
+        'allReportersConnection': {
+            'edges': [{
+                'node': {
+                    'lastName': 'ABA',
+                    'email': '',
+                }
+            }]
+        },
         'reporter': {
             'email': ''
         },
         '__debug': {
             'sql': [{
                 'rawSql': str(Reporter.objects.all().only('last_name', 'email').query)
+            }, {
+                'rawSql': 'SELECT COUNT(*) AS "__count" FROM "tests_reporter"'
+            }, {
+                'rawSql': str(Reporter.objects.all()[:1].query)
             }, {
                 'rawSql': str(Reporter.objects.only('email').order_by('pk')[:1].query)
             }]
diff --git a/graphene/contrib/django/utils.py b/graphene/contrib/django/utils.py
index 4a16702b..4a9a8396 100644
--- a/graphene/contrib/django/utils.py
+++ b/graphene/contrib/django/utils.py
@@ -49,10 +49,12 @@ class WrappedQueryset(LazyList):
         return self._origin.count()
 
 
-def maybe_queryset(value):
+def maybe_queryset(value, only=[]):
     if isinstance(value, Manager):
         value = value.get_queryset()
     if isinstance(value, QuerySet):
+        if only:
+            value = value.only(*only)
         return WrappedQueryset(value)
     return value