Merge pull request #165 from arunsureshkumar/bug-fix-feat-retrieving-queried-fields-only

Bug fix feat retrieving queried fields only

Author: abawchen

graphene_mongo/converter.py

@@ -4,7 +4,7 @@
 
 from graphene.types.json import JSONString
 from graphene.utils.str_converters import to_snake_case
-from mongoengine.base import get_document
+from mongoengine.base import get_document, LazyReference
 from . import advanced_types
 from .utils import import_single_dispatch, get_field_description, get_query_fields
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -138,10 +138,19 @@ def reference_resolver(root, *args, **kwargs):
                     to_resolve_models = list()
                     for each in querying_union_types:
                         to_resolve_models.append(registry._registry_string_map[each])
+                    to_resolve_object_ids = list()
                     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)
+                        if isinstance(each, LazyReference):
+                            to_resolve_object_ids.append(each.pk)
+                            model = each.document_type._class_name
+                            if model not in choice_to_resolve:
+                                choice_to_resolve[model] = list()
+                            choice_to_resolve[model].append(each.pk)
+                        else:
+                            to_resolve_object_ids.append(each["_ref"].id)
+                            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():
@@ -153,7 +162,6 @@ def reference_resolver(root, *args, **kwargs):
                     result = list()
                     for x in as_completed(futures):
                         result += x.result()
-                    to_resolve_object_ids = [each["_ref"].id for each in to_resolve]
                     result_object_ids = list()
                     for each in result:
                         result_object_ids.append(each.id)
@@ -249,7 +257,7 @@ def reference_resolver(root, *args, **kwargs):
                 return document.objects().no_dereference().only(*list(
                     set(list(_type._meta.required_fields) + queried_fields))).get(
                     pk=de_referenced["_ref"].id)
-            return document
+            return document()
         return None
 
     def lazy_reference_resolver(root, *args, **kwargs):

graphene_mongo/fields.py

@@ -232,7 +232,6 @@ def get_queryset(self, model, info, required_fields=list(), skip=None, limit=Non
                 args.update(queryset_or_filters)
         if limit is not None:
             if reversed:
-                order_by = ""
                 if self.order_by:
                     order_by = self.order_by + ",-pk"
                 else:
@@ -244,7 +243,6 @@ def get_queryset(self, model, info, required_fields=list(), skip=None, limit=Non
                     skip if skip else 0).limit(limit)
         elif skip is not None:
             if reversed:
-                order_by = ""
                 if self.order_by:
                     order_by = self.order_by + ",-pk"
                 else:
@@ -260,10 +258,13 @@ def default_resolver(self, _root, info, required_fields=list(), **args):
         args = args or {}
         if _root is not None:
             field_name = to_snake_case(info.field_name)
-            if field_name in _root._fields_ordered and not (isinstance(_root._fields[field_name].field,
-                                                                       mongoengine.EmbeddedDocumentField) or
-                                                            isinstance(_root._fields[field_name].field,
-                                                                       mongoengine.GenericEmbeddedDocumentField)):
+            if not hasattr(_root, "_fields_ordered"):
+                if getattr(_root, field_name, []) is not None:
+                    args["pk__in"] = [r.id for r in getattr(_root, field_name, [])]
+            elif field_name in _root._fields_ordered and not (isinstance(_root._fields[field_name].field,
+                                                                         mongoengine.EmbeddedDocumentField) or
+                                                              isinstance(_root._fields[field_name].field,
+                                                                         mongoengine.GenericEmbeddedDocumentField)):
                 if getattr(_root, field_name, []) is not None:
                     args["pk__in"] = [r.id for r in getattr(_root, field_name, [])]
 
@@ -351,10 +352,10 @@ def chained_resolver(self, resolver, is_partial, root, info, **args):
         for field in get_query_fields(info):
             if to_snake_case(field) in self.model._fields_ordered:
                 required_fields.append(to_snake_case(field))
+        args_copy = args.copy()
         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 + tuple(self.filter_args.keys()):
                         args_copy.pop(arg_name)
@@ -370,13 +371,17 @@ def chained_resolver(self, resolver, is_partial, root, info, **args):
                         return resolved
                     elif not isinstance(resolved[0], DBRef):
                         return resolved
+                    else:
+                        return self.default_resolver(root, info, required_fields, **args_copy)
                 elif isinstance(resolved, QuerySet):
                     args.update(resolved._query)
                     args_copy = args.copy()
                     for arg_name, arg in args.copy().items():
                         if arg_name not in self.model._fields_ordered + ('first', 'last', 'before', 'after') + tuple(
                                 self.filter_args.keys()):
                             args_copy.pop(arg_name)
+                            if arg_name == '_id' and isinstance(arg, dict):
+                                args_copy['pk__in'] = arg['$in']
                             if '.' in arg_name:
                                 operation = list(arg.keys())[0]
                                 args_copy[arg_name.replace('.', '__') + operation.replace('$', '__')] = arg[operation]