Feat/async mongoengine [Only For Review] [Do not Merge]

.github/workflows/ci.yml

@@ -15,8 +15,8 @@ jobs:
   build:
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
-        python: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
+        os: [ubuntu-latest]
+        python: ["3.10", "3.11", "3.12", "3.13", "3.14"]
     runs-on: ${{ matrix.os }}
     steps:
     - uses: actions/checkout@v3
@@ -31,6 +31,10 @@ jobs:
         poetry install --with dev
     - name: Lint with ruff
       run: make lint
+    - name: Start MongoDB
+      uses: supercharge/mongodb-github-action@1.12.1
+      with:
+        mongodb-replica-set: test-rs
     - name: Run Tests
       run: make test
     - name: Build Package

.github/workflows/lint.yml

@@ -11,7 +11,7 @@ jobs:
   build:
     strategy:
       matrix:
-        python: ["3.12"]
+        python: ["3.14"]
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3

.github/workflows/publish.yml

@@ -11,8 +11,8 @@ jobs:
   build:
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
-        python: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
+        os: [ubuntu-latest]
+        python: ["3.10", "3.11", "3.12", "3.13", "3.14"]
     runs-on: ${{ matrix.os }}
     steps:
     - uses: actions/checkout@v3
@@ -28,6 +28,10 @@ jobs:
     - name: Lint
       run: |
         make lint
+    - name: Start MongoDB
+      uses: supercharge/mongodb-github-action@1.12.1
+      with:
+        mongodb-replica-set: test-rs
     - name: Run Tests
       run: make test
     - name: Build Package

README.md

@@ -74,8 +74,6 @@ To create a GraphQL schema and async executor; for it you simply have to write t
 import graphene
 
 from graphene_mongo import AsyncMongoengineObjectType
-from graphene_mongo.utils import sync_to_async
-from concurrent.futures import ThreadPoolExecutor
 
 from .models import User as UserModel
 
@@ -89,8 +87,7 @@ class Query(graphene.ObjectType):
     users = graphene.List(User)
 
     async def resolve_users(self, info):
-        return await sync_to_async(list, thread_sensitive=False,
-                             executor=ThreadPoolExecutor())(UserModel.objects.all())
+        return await UserModel.aobjects.to_list()
 
 
 schema = graphene.Schema(query=Query)

README.rst

@@ -71,7 +71,7 @@ Then you can simply query the schema:
             }
         }
     '''
-    result = await schema.execute_async(query)
+    result = await schema.execute_async(query, context_value=dict())
 
 To learn more check out the `Flask MongoEngine example <https://github.com/graphql-python/graphene-mongo/tree/master/examples/flask_mongoengine>`__
 

docs/fields.rst

@@ -15,9 +15,7 @@ Mongoengine Fields
 - FloatField
 - GenericReferenceField
 - IntField
-- LazyReferenceField
 - ListField
-- LongField
 - MapField
 - MultiPolygonField
 - ObjectIdField

examples/django_mongoengine/bike/tests.py

@@ -207,7 +207,6 @@ def test_update_bike_mutation():
     request = factory.post(reverse("graphql-query"))
     client = Client(schema)
     result = client.execute(query, context=request)
-    print(result)
     assert result == expected
 
 

graphene_mongo/__init__.py

@@ -3,7 +3,7 @@
 from .types import MongoengineInputType, MongoengineInterfaceType, MongoengineObjectType
 from .types_async import AsyncMongoengineObjectType
 
-__version__ = "0.4.4"
+__version__ = "0.5.0"
 
 __all__ = [
     "__version__",

graphene_mongo/converter.py

@@ -1,25 +1,24 @@
+from functools import singledispatch
 import sys
 
 import graphene
-import mongoengine
-
 from graphene.types.json import JSONString
 from graphene.utils.str_converters import to_camel_case
-from mongoengine.base import get_document
+import mongoengine
+
 from . import advanced_types
-from .utils import (
-    get_field_description,
-    get_field_is_required,
-    get_field_resolver,
-    ExecutorEnum,
-)
 from .field_resolvers import (
-    DynamicLazyFieldResolver,
     DynamicReferenceFieldResolver,
     ListFieldResolver,
     UnionFieldResolver,
 )
-from functools import singledispatch
+from .utils import (
+    ExecutorEnum,
+    get_field_description,
+    get_field_is_required,
+    get_field_resolver,
+    get_document,
+)
 
 
 class MongoEngineConversionError(Exception):
@@ -53,7 +52,6 @@ def convert_field_to_id(field, registry=None, executor: ExecutorEnum = ExecutorE
 
 
 @convert_mongoengine_field.register(mongoengine.IntField)
-@convert_mongoengine_field.register(mongoengine.LongField)
 @convert_mongoengine_field.register(mongoengine.SequenceField)
 def convert_field_to_int(field, registry=None, executor: ExecutorEnum = ExecutorEnum.SYNC):
     return graphene.Int(
@@ -235,15 +233,13 @@ def convert_field_to_union(field, registry=None, executor: ExecutorEnum = Execut
     Meta = type("Meta", (object,), {"types": tuple(_types)})
     _union = type(name, (graphene.Union,), {"Meta": Meta})
 
-    if isinstance(field, mongoengine.GenericReferenceField) or isinstance(
-        field, mongoengine.GenericLazyReferenceField
-    ):
+    if isinstance(field, mongoengine.GenericReferenceField):
         field_resolver = None
         required = False
         if field.db_field is not None:
             required = get_field_is_required(field, registry)
             resolver_function = getattr(
-                registry.get_type_for_model(field.owner_document, executor=executor),
+                registry.get_type_for_model(field.owner_document),
                 "resolve_" + field.db_field,
                 None,
             )
@@ -270,12 +266,11 @@ def convert_field_to_union(field, registry=None, executor: ExecutorEnum = Execut
 
 @convert_mongoengine_field.register(mongoengine.EmbeddedDocumentField)
 @convert_mongoengine_field.register(mongoengine.ReferenceField)
-@convert_mongoengine_field.register(mongoengine.CachedReferenceField)
 def convert_field_to_dynamic(field, registry=None, executor: ExecutorEnum = ExecutorEnum.SYNC):
     model = field.document_type
 
     def dynamic_type():
-        _type = registry.get_type_for_model(model, executor=executor)
+        _type = registry.get_type_for_model(model)
         if not _type:
             return None
         if isinstance(field, mongoengine.EmbeddedDocumentField):
@@ -289,7 +284,7 @@ def dynamic_type():
         if field.db_field is not None:
             required = get_field_is_required(field, registry)
             resolver_function = getattr(
-                registry.get_type_for_model(field.owner_document, executor=executor),
+                registry.get_type_for_model(field.owner_document),
                 "resolve_" + field.db_field,
                 None,
             )
@@ -314,45 +309,6 @@ def dynamic_type():
     return graphene.Dynamic(dynamic_type)
 
 
-@convert_mongoengine_field.register(mongoengine.LazyReferenceField)
-def convert_lazy_field_to_dynamic(field, registry=None, executor: ExecutorEnum = ExecutorEnum.SYNC):
-    model = field.document_type
-
-    def dynamic_type():
-        _type = registry.get_type_for_model(model, executor=executor)
-        if not _type:
-            return None
-        field_resolver = None
-        required = False
-        if field.db_field is not None:
-            required = get_field_is_required(field, registry)
-            resolver_function = getattr(
-                registry.get_type_for_model(field.owner_document, executor=executor),
-                "resolve_" + field.db_field,
-                None,
-            )
-            if resolver_function and callable(resolver_function):
-                field_resolver = resolver_function
-
-        return graphene.Field(
-            _type,
-            resolver=get_field_resolver(
-                field_resolver=field_resolver,
-                default_sync_resolver=DynamicLazyFieldResolver.lazy_resolver(
-                    field=field, registry=registry, executor=executor
-                ),
-                default_async_resolver=DynamicLazyFieldResolver.lazy_resolver_async(
-                    field=field, registry=registry, executor=executor
-                ),
-                executor=executor,
-            ),
-            description=get_field_description(field, registry),
-            required=required,
-        )
-
-    return graphene.Dynamic(dynamic_type)
-
-
 if sys.version_info >= (3, 6):
 
     @convert_mongoengine_field.register(mongoengine.EnumField)

graphene_mongo/dataloader.py

@@ -0,0 +1,127 @@
+"""
+MongoEngine DataLoader for GraphQL resolution.
+
+Usage:
+    # ====================== Type level ======================
+
+    class UserType(AsyncMongoengineObjectType):
+        class Meta:
+            model = User
+
+        @classmethod
+        async def dataloader_resolver(
+            cls, info: GraphQLResolveInfo, ids: list[ObjectId], projections: list[str] | None = None
+        ):
+            docs = User.aobjects.filter(id__in=ids)
+            if projections:
+                docs = docs.only(*projections)
+            return await docs.to_list()
+
+    # ====================== Field Level ======================
+    from graphene_mongo.utils import get_dataloader
+
+    loader = get_dataloader(info)
+
+    # Single fetch
+    user = await loader.model(User, projections).get(user_id)
+
+    # Pre-load in bulk (e.g. in a resolver that already has many ids)
+    await loader.model(Post, projections).load_many([id1, id2, id3])
+    post = await loader.model(Post).get(id1)  # served from cache
+
+    # This attached the loader to your GraphQL context so it's
+    # shared across all resolvers in a single request, giving you automatic
+    # N+1 batching.
+"""
+
+from __future__ import annotations
+
+from typing import Generic, Type, TypeVar
+
+from aiodataloader import DataLoader
+from bson import ObjectId
+from graphql import GraphQLResolveInfo
+from mongoengine import Document
+
+T = TypeVar("T", bound=Document)
+
+
+class ModelLoader(DataLoader, Generic[T]):
+    """
+    A per-model DataLoader that batches `.get(id)` calls into a single
+    MongoDB query per tick of the event loop.
+    """
+
+    def __init__(self, model: Type[T], info: GraphQLResolveInfo, projections: set[str], **kwargs):
+        super().__init__(cache=True, **kwargs)
+        self._info = info
+        self._model = model
+        self._projections: set[str] = projections
+
+        from .registry import get_global_async_registry
+        from .types_async import AsyncMongoengineObjectType
+
+        registry = get_global_async_registry()
+        self._gql_type: AsyncMongoengineObjectType = registry.get_type_for_document_model(model)
+        if self._gql_type is None:
+            raise NotImplementedError(f"Please define AsyncMongoengineObjectType for {model}")
+
+    async def batch_load_fn(self, keys: list[str]) -> list[T | None]:
+        """
+        Called once per event-loop tick with all ids accumulated so far.
+        Executes a single `filter(id__in=...)` query and maps results back
+        to the original key order (DataLoader requires 1-to-1 ordering).
+        """
+        # Normalise to ObjectId so MongoEngine is happy either way
+        object_ids: list[ObjectId] = []
+        for k in keys:
+            try:
+                object_ids.append(ObjectId(k) if not isinstance(k, ObjectId) else k)
+            except Exception:
+                object_ids.append(k)  # let Mongo surface the error naturally
+
+        docs = await self._gql_type.dataloader_resolver(
+            info=self._info,
+            ids=object_ids,
+            projections=list(self._projections) if self._projections else None,
+        )
+
+        id_map: dict[str, T] = {str(doc.id): doc for doc in docs}
+
+        # Preserve key order; missing ids resolve to None
+        return [id_map.get(str(k)) for k in keys]
+
+
+class MongoDataLoader:
+    """
+    Request-scoped container for per-model DataLoaders.
+    """
+
+    def __init__(self, info: GraphQLResolveInfo):
+        self._info = info
+        self._loaders: dict[tuple[type, frozenset[str]], ModelLoader] = {}
+
+    def model(self, model_class: Type[T], projections: set[str]) -> ModelLoader[T]:
+        """
+        Returns (or creates) the DataLoader for the given MongoEngine model.
+        Loaders are cached per model class for the lifetime of this object.
+        """
+        requested_projections = frozenset(projections)
+        key = (model_class, requested_projections)
+
+        if self._loaders.get(key):
+            return self._loaders[key]
+
+        # Trying to find a loader which has superset projections
+        existing_key = next((k for k in self._loaders if k[0] == model_class), None)
+        if existing_key:
+            existing_loader = self._loaders[existing_key]
+            existing_projections = existing_key[1]
+            if projections.issubset(existing_projections):
+                # reuse old loader as its projections are subset
+                return existing_loader
+
+        self._loaders[key] = ModelLoader(
+            model=model_class, info=self._info, projections=set(requested_projections)
+        )
+        return self._loaders[key]

graphene_mongo/field_resolvers/__init__.py

@@ -1,10 +1,8 @@
-from .dynamic_lazy_field_resolver import DynamicLazyFieldResolver
 from .dynamic_reference_field_resolver import DynamicReferenceFieldResolver
 from .list_field_resolver import ListFieldResolver
 from .union_resolver import UnionFieldResolver
 
 __all__ = [
-    "DynamicLazyFieldResolver",
     "DynamicReferenceFieldResolver",
     "ListFieldResolver",
     "UnionFieldResolver",