Query efficiency and performance - Retrieving only the queried fields from database

Author: Arun S Kumar

graphene_mongo/converter.py

@@ -6,7 +6,7 @@
 from mongoengine.base import get_document
 
 from . import advanced_types
-from .utils import import_single_dispatch, get_field_description
+from .utils import import_single_dispatch, get_field_description, get_query_fields
 
 singledispatch = import_single_dispatch()
 
@@ -186,7 +186,9 @@ def convert_lazy_field_to_dynamic(field, registry=None):
 
     def lazy_resolver(root, *args, **kwargs):
         if getattr(root, field.name or field.db_name):
-            return getattr(root, field.name or field.db_name).fetch()
+            only_fields = get_query_fields(args[0]).keys()
+            document = getattr(root, field.name or field.db_name)
+            return document.document_type.objects().only(*only_fields).get(pk=document.pk)
 
     def dynamic_type():
         _type = registry.get_type_for_model(model)

graphene_mongo/fields.py

@@ -5,6 +5,7 @@
 
 import graphene
 import mongoengine
+from graphql.utils.ast_to_dict import ast_to_dict
 from promise import Promise
 from graphql_relay import from_global_id
 from graphene.relay import ConnectionField
@@ -21,7 +22,7 @@
 )
 from .converter import convert_mongoengine_field, MongoEngineConversionError
 from .registry import get_global_registry
-from .utils import get_model_reference_fields, get_node_from_global_id
+from .utils import get_model_reference_fields, get_node_from_global_id, get_query_fields, camel_to_snake
 
 
 class MongoengineConnectionField(ConnectionField):
@@ -98,18 +99,18 @@ def is_filterable(k):
             if isinstance(converted, (ConnectionField, Dynamic)):
                 return False
             if callable(getattr(converted, "type", None)) and isinstance(
-                converted.type(),
-                (
-                    FileFieldType,
-                    PointFieldType,
-                    MultiPolygonFieldType,
-                    graphene.Union,
-                    PolygonFieldType,
-                ),
+                    converted.type(),
+                    (
+                            FileFieldType,
+                            PointFieldType,
+                            MultiPolygonFieldType,
+                            graphene.Union,
+                            PolygonFieldType,
+                    ),
             ):
                 return False
             if isinstance(converted, (graphene.List)) and issubclass(
-                getattr(converted, "_of_type", None), graphene.Union
+                    getattr(converted, "_of_type", None), graphene.Union
             ):
                 return False
 
@@ -160,16 +161,16 @@ def get_reference_field(r, kv):
             field = kv[1]
             mongo_field = getattr(self.model, kv[0], None)
             if isinstance(
-                mongo_field,
-                (mongoengine.LazyReferenceField, mongoengine.ReferenceField),
+                    mongo_field,
+                    (mongoengine.LazyReferenceField, mongoengine.ReferenceField),
             ):
                 field = convert_mongoengine_field(mongo_field, self.registry)
             if callable(getattr(field, "get_type", None)):
                 _type = field.get_type()
                 if _type:
                     node = _type._type._meta
                     if "id" in node.fields and not issubclass(
-                        node.model, (mongoengine.EmbeddedDocument,)
+                            node.model, (mongoengine.EmbeddedDocument,)
                     ):
                         r.update({kv[0]: node.fields["id"]._type.of_type()})
             return r
@@ -180,7 +181,7 @@ def get_reference_field(r, kv):
     def fields(self):
         return self._type._meta.fields
 
-    def get_queryset(self, model, info, **args):
+    def get_queryset(self, model, info, only_fields=list(), **args):
         if args:
             reference_fields = get_model_reference_fields(self.model)
             hydrated_references = {}
@@ -198,12 +199,13 @@ def get_queryset(self, model, info, **args):
                 return queryset_or_filters
             else:
                 args.update(queryset_or_filters)
+
         return model.objects(**args).order_by(self.order_by)
 
-    def default_resolver(self, _root, info, **args):
+    def default_resolver(self, _root, info, only_fields=list(), **args):
         args = args or {}
 
-        if _root is not None:
+        if _root is not None and getattr(_root, info.field_name, []) is not None:
             args["pk__in"] = [r.pk for r in getattr(_root, info.field_name, [])]
 
         connection_args = {
@@ -219,7 +221,7 @@ def default_resolver(self, _root, info, **args):
             args['pk'] = from_global_id(_id)[-1]
 
         if callable(getattr(self.model, "objects", None)):
-            iterables = self.get_queryset(self.model, info, **args)
+            iterables = self.get_queryset(self.model, info, only_fields, **args)
             list_length = iterables.count()
         else:
             iterables = []
@@ -239,23 +241,31 @@ def default_resolver(self, _root, info, **args):
         return connection
 
     def chained_resolver(self, resolver, is_partial, root, info, **args):
+        only_fields = list()
+        for field in get_query_fields(info):
+            if camel_to_snake(field) in self.model._fields_ordered:
+                only_fields.append(camel_to_snake(field))
         if not bool(args) or not is_partial:
+            if isinstance(self.model, mongoengine.Document) or isinstance(self.model,
+                                                                          mongoengine.base.metaclasses.TopLevelDocumentMetaclass):
+                args_copy = args.copy()
+                for arg_name, arg in args.copy().items():
+                    if arg_name not in self.model._fields_ordered:
+                        args_copy.pop(arg_name)
             # XXX: Filter nested args
             resolved = resolver(root, info, **args)
             if resolved is not None:
                 return resolved
-        return self.default_resolver(root, info, **args)
+        return self.default_resolver(root, info, only_fields, **args)
 
     @classmethod
     def connection_resolver(cls, resolver, connection_type, root, info, **args):
         iterable = resolver(root, info, **args)
         if isinstance(connection_type, graphene.NonNull):
             connection_type = connection_type.of_type
-
         on_resolve = partial(cls.resolve_connection, connection_type, args)
         if Promise.is_thenable(iterable):
             return Promise.resolve(iterable).then(on_resolve)
-
         return on_resolve(iterable)
 
     def get_resolver(self, parent_resolver):

graphene_mongo/utils.py

@@ -6,6 +6,7 @@
 import mongoengine
 from graphene import Node
 from graphene.utils.trim_docstring import trim_docstring
+from graphql.utils.ast_to_dict import ast_to_dict
 
 
 def get_model_fields(model, excluding=None):
@@ -23,8 +24,8 @@ def get_model_reference_fields(model, excluding=None):
     attributes = dict()
     for attr_name, attr in model._fields.items():
         if attr_name in excluding or not isinstance(
-            attr,
-            (mongoengine.fields.ReferenceField, mongoengine.fields.LazyReferenceField),
+                attr,
+                (mongoengine.fields.ReferenceField, mongoengine.fields.LazyReferenceField),
         ):
             continue
         attributes[attr_name] = attr
@@ -33,8 +34,8 @@ def get_model_reference_fields(model, excluding=None):
 
 def is_valid_mongoengine_model(model):
     return inspect.isclass(model) and (
-        issubclass(model, mongoengine.Document)
-        or issubclass(model, mongoengine.EmbeddedDocument)
+            issubclass(model, mongoengine.Document)
+            or issubclass(model, mongoengine.EmbeddedDocument)
     )
 
 
@@ -101,3 +102,61 @@ def get_node_from_global_id(node, info, global_id):
                 return interface.get_node_from_global_id(info, global_id)
     except AttributeError:
         return Node.get_node_from_global_id(info, global_id)
+
+
+def collect_query_fields(node, fragments):
+    """Recursively collects fields from the AST
+
+    Args:
+        node (dict): A node in the AST
+        fragments (dict): Fragment definitions
+
+    Returns:
+        A dict mapping each field found, along with their sub fields.
+
+        {'name': {},
+         'image': {'id': {},
+                                   'name': {},
+                                   'description': {}},
+         'slug': {}}
+    """
+
+    field = {}
+
+    if node.get('selection_set'):
+        for leaf in node['selection_set']['selections']:
+            if leaf['kind'] == 'Field':
+                field.update({
+                    leaf['name']['value']: collect_query_fields(leaf, fragments)
+                })
+            elif leaf['kind'] == 'FragmentSpread':
+                field.update(collect_query_fields(fragments[leaf['name']['value']],
+                                                  fragments))
+
+    return field
+
+
+def get_query_fields(info):
+    """A convenience function to call collect_query_fields with info
+
+    Args:
+        info (ResolveInfo)
+
+    Returns:
+        dict: Returned from collect_query_fields
+    """
+
+    fragments = {}
+    node = ast_to_dict(info.field_asts[0])
+
+    for name, value in info.fragments.items():
+        fragments[name] = ast_to_dict(value)
+
+    query = collect_query_fields(node, fragments)
+    if "edges" in query:
+        return query["edges"]["node"].keys()
+    return query
+
+
+def camel_to_snake(field):
+    return ''.join(['_' + c.lower() if c.isupper() else c for c in field]).lstrip('_')