Query efficiency and performance - By retrieving only the queried fields from the database - Test Passed

Author: Arun S Kumar

graphene_mongo/converter.py

@@ -7,6 +7,7 @@
 from mongoengine.base import get_document
 from . import advanced_types
 from .utils import import_single_dispatch, get_field_description, get_query_fields
+from concurrent.futures import ThreadPoolExecutor, wait, as_completed
 
 singledispatch = import_single_dispatch()
 
@@ -104,6 +105,46 @@ def convert_file_to_field(field, registry=None):
 def convert_field_to_list(field, registry=None):
     base_type = convert_mongoengine_field(field.field, registry=registry)
     if isinstance(base_type, graphene.Field):
+        if isinstance(field.field, mongoengine.GenericReferenceField):
+            def get_reference_objects(*args, **kwargs):
+                if args[0][1]:
+                    document = get_document(args[0][0])
+                    document_field = mongoengine.ReferenceField(document)
+                    document_field = convert_mongoengine_field(document_field, registry)
+                    document_field_type = document_field.get_type().type._meta.name
+                    only_fields = [to_snake_case(i) for i in get_query_fields(args[0][3])[document_field_type].keys()]
+                    return document.objects().no_dereference().only(*only_fields).filter(pk__in=args[0][1])
+                else:
+                    return []
+
+            def reference_resolver(root, *args, **kwargs):
+                choice_to_resolve = dict()
+                to_resolve = getattr(root, field.name or field.db_name)
+                for each in to_resolve:
+                    if each['_cls'] not in choice_to_resolve:
+                        choice_to_resolve[each['_cls']] = list()
+                    choice_to_resolve[each['_cls']].append(each["_ref"].id)
+
+                pool = ThreadPoolExecutor(5)
+                futures = list()
+                for model, object_id_list in choice_to_resolve.items():
+                    futures.append(pool.submit(get_reference_objects, (model, object_id_list, registry, *args)))
+                result = list()
+                for x in as_completed(futures):
+                    result += x.result()
+                to_resolve_object_ids = [each["_ref"].id for each in to_resolve]
+                result_to_resolve_object_ids = [each.id for each in result]
+                ordered_result = list()
+                for each in to_resolve_object_ids:
+                    ordered_result.append(result[result_to_resolve_object_ids.index(each)])
+                return ordered_result
+
+            return graphene.List(
+                base_type._type,
+                description=get_field_description(field, registry),
+                required=field.required,
+                resolver=reference_resolver
+            )
         return graphene.List(
             base_type._type,
             description=get_field_description(field, registry),
@@ -128,7 +169,7 @@ def convert_field_to_list(field, registry=None):
     return graphene.List(
         base_type,
         description=get_field_description(field, registry),
-        required=field.required
+        required=field.required,
     )
 
 
@@ -172,7 +213,8 @@ def reference_resolver(root, *args, **kwargs):
         return document.objects().no_dereference().only(*only_fields).get(pk=dereferenced["_ref"].id)
 
     if isinstance(field, mongoengine.GenericReferenceField):
-        return graphene.Field(_union, resolver=reference_resolver)
+        return graphene.Field(_union, resolver=reference_resolver,
+                              description=get_field_description(field, registry))
 
     return graphene.Field(_union)
 

graphene_mongo/fields.py

@@ -5,6 +5,7 @@
 
 import graphene
 import mongoengine
+from bson import DBRef
 from graphene import Context
 from graphene.utils.str_converters import to_snake_case
 from promise import Promise
@@ -61,6 +62,12 @@ def model(self):
     def order_by(self):
         return self.node_type._meta.order_by
 
+    @property
+    def only_fields(self):
+        if isinstance(self.node_type._meta.only_fields, str):
+            return self.node_type._meta.only_fields.split(",")
+        return list()
+
     @property
     def registry(self):
         return getattr(self.node_type._meta, "registry", get_global_registry())
@@ -206,8 +213,10 @@ def get_queryset(self, model, info, only_fields=list(), **args):
     def default_resolver(self, _root, info, only_fields=list(), **args):
         args = args or {}
 
-        if _root is not None and getattr(_root, info.field_name, []) is not None:
-            args["pk__in"] = [r.id for r in getattr(_root, info.field_name, [])]
+        if _root is not None:
+            field_name = to_snake_case(info.field_name)
+            if getattr(_root, field_name, []) is not None:
+                args["pk__in"] = [r.id for r in getattr(_root, field_name, [])]
 
         connection_args = {
             "first": args.pop("first", None),
@@ -243,10 +252,13 @@ def default_resolver(self, _root, info, only_fields=list(), **args):
 
     def chained_resolver(self, resolver, is_partial, root, info, **args):
         only_fields = list()
+        for field in self.only_fields:
+            if field in self.model._fields_ordered:
+                only_fields.append(field)
         for field in get_query_fields(info):
             if to_snake_case(field) in self.model._fields_ordered:
                 only_fields.append(to_snake_case(field))
-        if root is None and (not bool(args) or not is_partial):
+        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()
@@ -259,7 +271,13 @@ def chained_resolver(self, resolver, is_partial, root, info, **args):
             # XXX: Filter nested args
             resolved = resolver(root, info, **args)
             if resolved is not None:
-                return resolved
+                if isinstance(resolved, list):
+                    if resolved == list():
+                        return resolved
+                    elif not isinstance(resolved[0], DBRef):
+                        return resolved
+                else:
+                    return resolved
         return self.default_resolver(root, info, only_fields, **args)
 
     @classmethod

graphene_mongo/tests/test_relay_query.py

@@ -17,7 +17,7 @@ class Query(graphene.ObjectType):
         reporter = graphene.Field(nodes.ReporterNode)
 
         def resolve_reporter(self, *args, **kwargs):
-            return models.Reporter.objects.first()
+            return models.Reporter.objects.no_dereference().first()
 
     query = """
         query ReporterQuery {