Repository: strawberry-graphql/strawberry

Issue:
N + 1 problem - Support for dataloaders?
Hi 👋

How to solve N+1 problem when using strawberry?


Discussion:
Hello! Good question, you might be able to use the same dataloader that works for Graphene, but not 100% sure as I have never used it.

We are doing some work around N+1 at work with Strawberry, so we might add some support for this.

I need to get familiar with dataloaders, but I guess that would be one solution :)

@hyzyla can you provide more context on what you're doing? maybe there's other ways to solve your problem :)


Currently, I just seeking a library for a new project and this library seems promising. I have experience in solving the N+1 problem for Graphene by using [aiodatalodaer](https://github.com/syrusakbary/aiodataloader). But the solution was non-intuitive and solved just by adopting some tutorials from the web for non-async data loader.

As example I can give you most iconic version of query for blog.
```graphql
{
    blog(id: "1") {
        text,
        comments {
            text,
            date,
            author {   # resolving authors for N comments causes N requests
                  name,
                  avatarURL
            }
        }
    }
}
```
Yeah, from what I've seen you either need async or to use a Promise library. I'll think about this a bit, for example we use Django and using async is not possible with the database right now, so we need some other solution for this. 🤔 
Looks like you can get around the async DB issue in Django by using the `async_to_sync` wrapper:

```python
from asgiref.sync import sync_to_async

class FooLoader(aiodataloader.DataLoader):
    @sync_to_async
    def fetch_foo(self, keys):
        return Foo.objects.in_bulk(keys)

    async def batch_load_fn(self, keys):
        foos = await self.fetch_foo(keys)
        # some other processing
        return foos
```

https://github.com/graphql-python/graphene-django/issues/705#issuecomment-670511313
@jkimbo that's a good point, but I guess when wouldn't be able to use that with old version of django, right? 

Plus it needs to work with the async view (which we aren't providing yet[1]), so might not be that easy.

I wonder if we can make this work without requiring async, I mean with django we can still do bulk fetching of data, no? I was naively thinking that we could have some sort of sync dataloader that first collects the keys/ids and then our GraphQL executor could do the final resolve of the data, some sort of mix of promises and awaitable if that makes sense, but specific for this use case. It might be a stupid idea though :D

In any case the aiodataloader might be a good interim solution.

[1] We have an ASGI one, so should be easy to add one for Django too.
I'm going to mark this issue as required for v1. I think it is pretty important to have it, at least a basic support for dataloaders :)
Hi, I created an [example](https://gist.github.com/hyzyla/68d42951cf5788dd4fc7331d98c4a2be) of using aiodataloader with strawberry. May be helpfull for async integration of strawberry
That's pretty cool! I'm going to check it in the next few days. I need to work on caching some of our data on redis and I think we'll probably use the dataloader pattern for that. so this snipped will be handy. I definitely want to add first class support for dataloaders in strawberry, but I need to work with them first (I'm still unfamiliar with them).

I also saw this blog post that might useful: https://medium.com/@menecio/using-djangos-orm-from-aiohttp-fb9538072ba0
A non-async dataloader is installed with Graphene, it's part of the author's [promise](https://github.com/syrusakbary/promise) library, I used it with some success. I don't _think_ it should require any special magic to use it with Strawberry instead. `aiodataloader` is from the same person and has basically the same API but without the underlying promise implementation.
@AndrewIngram graphql-core v3 dropped support for the promise library so it won't work with Strawberry. The aiodataloader does work but because Django doesn't have full async support getting it working with the orm is a problem. All the solutions I've come across require code changes in the application and not the just in Strawberry which is not great.
😢 

Resolving patch:
diff --git a/strawberry/dataloader.py b/strawberry/dataloader.py
new file mode 100644
--- /dev/null
+++ b/strawberry/dataloader.py
@@ -0,0 +1,127 @@
+import dataclasses
+from asyncio import create_task, get_event_loop
+from asyncio.events import AbstractEventLoop
+from asyncio.futures import Future
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, Dict, Generic, List, Optional, TypeVar
+
+from .exceptions import WrongNumberOfResultsReturned
+
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+@dataclass
+class LoaderTask(Generic[K, T]):
+    key: K
+    future: Future
+
+
+@dataclass
+class Batch(Generic[K, T]):
+    tasks: List[LoaderTask] = dataclasses.field(default_factory=list)
+    dispatched: bool = False
+
+    def add_task(self, key: Any, future: Future):
+        task = LoaderTask[K, T](key, future)
+        self.tasks.append(task)
+
+    def __len__(self) -> int:
+        return len(self.tasks)
+
+
+class DataLoader(Generic[K, T]):
+    queue: List[LoaderTask] = []
+    batch: Optional[Batch[K, T]] = None
+    cache: bool = False
+    cache_map: Dict[K, Future]
+
+    def __init__(
+        self,
+        load_fn: Callable[[List[K]], Awaitable[List[T]]],
+        max_batch_size: Optional[int] = None,
+        cache: bool = True,
+        loop: AbstractEventLoop = None,
+    ):
+        self.load_fn = load_fn
+        self.max_batch_size = max_batch_size
+
+        self.loop = loop or get_event_loop()
+
+        self.cache = cache
+
+        if self.cache:
+            self.cache_map = {}
+
+    def load(self, key: K) -> Awaitable[T]:
+        if self.cache:
+            future = self.cache_map.get(key)
+
+            if future:
+                return future
+
+        future = self.loop.create_future()
+
+        if self.cache:
+            self.cache_map[key] = future
+
+        batch = get_current_batch(self)
+        batch.add_task(key, future)
+
+        return future
+
+
+def should_create_new_batch(loader: DataLoader, batch: Batch) -> bool:
+    if (
+        batch.dispatched
+        or loader.max_batch_size
+        and len(batch) >= loader.max_batch_size
+    ):
+        return True
+
+    return False
+
+
+def get_current_batch(loader: DataLoader) -> Batch:
+    if loader.batch and not should_create_new_batch(loader, loader.batch):
+        return loader.batch
+
+    loader.batch = Batch()
+
+    dispatch(loader, loader.batch)
+
+    return loader.batch
+
+
+def dispatch(loader: DataLoader, batch: Batch):
+    async def dispatch():
+        await dispatch_batch(loader, batch)
+
+    loader.loop.call_soon(create_task, dispatch())
+
+
+async def dispatch_batch(loader: DataLoader, batch: Batch) -> None:
+    batch.dispatched = True
+
+    keys = [task.key for task in batch.tasks]
+
+    # TODO: check if load_fn return an awaitable and it is a list
+
+    try:
+        values = await loader.load_fn(keys)
+        values = list(values)
+
+        if len(values) != len(batch):
+            raise WrongNumberOfResultsReturned(
+                expected=len(batch), received=len(values)
+            )
+
+        for task, value in zip(batch.tasks, values):
+            if isinstance(value, BaseException):
+                task.future.set_exception(value)
+            else:
+                task.future.set_result(value)
+    except Exception as e:
+        for task in batch.tasks:
+            task.future.set_exception(e)
diff --git a/strawberry/exceptions.py b/strawberry/exceptions.py
--- a/strawberry/exceptions.py
+++ b/strawberry/exceptions.py
@@ -127,3 +127,13 @@ def __init__(self, scalar_name: str):
         message = f"Scalar `{scalar_name}` has already been registered"
 
         super().__init__(message)
+
+
+class WrongNumberOfResultsReturned(Exception):
+    def __init__(self, expected: int, received: int):
+        message = (
+            "Received wrong number of results in dataloader, "
+            f"expected: {expected}, received: {received}"
+        )
+
+        super().__init__(message)
diff --git a/strawberry/ext/mypy_plugin.py b/strawberry/ext/mypy_plugin.py
--- a/strawberry/ext/mypy_plugin.py
+++ b/strawberry/ext/mypy_plugin.py
@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 from mypy.nodes import (
     GDEF,
@@ -23,7 +23,13 @@
 
 
 class InvalidNodeTypeException(Exception):
-    pass
+    def __init__(self, node: Any) -> None:
+        self.message = f"Invalid node type: {str(node)}"
+
+        super().__init__()
+
+    def __str__(self) -> str:
+        return self.message
 
 
 def lazy_type_analyze_callback(ctx: AnalyzeTypeContext) -> Type:
@@ -51,7 +57,7 @@ def _get_type_for_expr(expr: Expression, api: SemanticAnalyzerPluginInterface):
             sym = api.lookup_fully_qualified_or_none(expr.fullname)
 
             if sym and isinstance(sym.node, Var):
-                raise InvalidNodeTypeException()
+                raise InvalidNodeTypeException(sym.node)
 
         return api.named_type(expr.name)
 
@@ -65,7 +71,7 @@ def _get_type_for_expr(expr: Expression, api: SemanticAnalyzerPluginInterface):
         if expr.fullname:
             return api.named_type(expr.fullname)
         else:
-            raise InvalidNodeTypeException()
+            raise InvalidNodeTypeException(expr)
 
     raise ValueError(f"Unsupported expression {type(expr)}")
 
@@ -136,14 +142,22 @@ def enum_hook(ctx: DynamicClassDefContext) -> None:
             )
             return
 
-    enum_type = _get_type_for_expr(first_argument, ctx.api)
+    try:
+        enum_type = _get_type_for_expr(first_argument, ctx.api)
 
-    type_alias = TypeAlias(
-        enum_type,
-        fullname=ctx.api.qualified_name(ctx.name),
-        line=ctx.call.line,
-        column=ctx.call.column,
-    )
+        type_alias = TypeAlias(
+            enum_type,
+            fullname=ctx.api.qualified_name(ctx.name),
+            line=ctx.call.line,
+            column=ctx.call.column,
+        )
+    except InvalidNodeTypeException:
+        type_alias = TypeAlias(
+            AnyType(TypeOfAny.from_error),
+            fullname=ctx.api.qualified_name(ctx.name),
+            line=ctx.call.line,
+            column=ctx.call.column,
+        )
 
     ctx.api.add_symbol_table_node(
         ctx.name, SymbolTableNode(GDEF, type_alias, plugin_generated=False)


Test patch:
diff --git a/tests/mypy/mypy.ini b/mypy_tests.ini
similarity index 90%
rename from tests/mypy/mypy.ini
rename to mypy_tests.ini
--- a/tests/mypy/mypy.ini
+++ b/mypy_tests.ini
@@ -2,7 +2,7 @@
 plugins = strawberry.ext.mypy_plugin
 check_untyped_defs = True
 ignore_errors = False
-ignore_missing_imports = True
+; ignore_missing_imports = True
 strict_optional = True
 implicit_reexport = False
 ; Disabled because of this bug: https://github.com/python/mypy/issues/9689
diff --git a/tests/django/test_dataloaders.py b/tests/django/test_dataloaders.py
new file mode 100644
--- /dev/null
+++ b/tests/django/test_dataloaders.py
@@ -0,0 +1,95 @@
+import json
+from typing import List
+
+import pytest
+
+from asgiref.sync import sync_to_async
+
+import django
+from django.test.client import RequestFactory
+
+import strawberry
+from strawberry.dataloader import DataLoader
+from strawberry.django.views import AsyncGraphQLView
+
+from .app.models import Example
+
+
+pytestmark = [
+    pytest.mark.asyncio,
+    pytest.mark.skipif(
+        django.VERSION < (3, 1),
+        reason="Async views are only supported in Django >= 3.1",
+    ),
+]
+
+
+def _prepare_db():
+    ids = []
+
+    for index in range(5):
+        ids.append(Example.objects.create(name=f"This is a demo async {index}").id)
+
+    return ids
+
+
+@pytest.mark.django_db
+async def test_fetch_data_from_db(mocker):
+    def _sync_batch_load(keys):
+        data = Example.objects.filter(id__in=keys)
+
+        return list(data)
+
+    prepare_db = sync_to_async(_prepare_db)
+    batch_load = sync_to_async(_sync_batch_load)
+
+    ids = await prepare_db()
+
+    async def idx(keys) -> List[Example]:
+        return await batch_load(keys)
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader)
+
+    @strawberry.type
+    class Query:
+        hello: str = "strawberry"
+
+        @strawberry.field
+        async def get_example(self, id: strawberry.ID) -> str:
+            example = await loader.load(id)
+
+            return example.name
+
+    schema = strawberry.Schema(query=Query)
+
+    query = f"""{{
+        a: getExample(id: "{ids[0]}")
+        b: getExample(id: "{ids[1]}")
+        c: getExample(id: "{ids[2]}")
+        d: getExample(id: "{ids[3]}")
+        e: getExample(id: "{ids[4]}")
+    }}"""
+
+    factory = RequestFactory()
+    request = factory.post(
+        "/graphql/", {"query": query}, content_type="application/json"
+    )
+
+    response = await AsyncGraphQLView.as_view(schema=schema)(request)
+    data = json.loads(response.content.decode())
+
+    assert not data.get("errors")
+    assert data["data"] == {
+        "a": "This is a demo async 0",
+        "b": "This is a demo async 1",
+        "c": "This is a demo async 2",
+        "d": "This is a demo async 3",
+        "e": "This is a demo async 4",
+    }
+
+    reset_db = sync_to_async(lambda: Example.objects.all().delete())
+    await reset_db()
+
+    mock_loader.assert_called_once_with([str(id_) for id_ in ids])
diff --git a/tests/mypy/test_dataloaders.yml b/tests/mypy/test_dataloaders.yml
new file mode 100644
--- /dev/null
+++ b/tests/mypy/test_dataloaders.yml
@@ -0,0 +1,39 @@
+- case: test_dataloader_no_types
+  main: |
+    from strawberry.dataloader import DataLoader
+
+    def load(keys):
+        return keys
+
+    loader = DataLoader(load)
+
+    reveal_type(loader)
+
+    async def run() -> None:
+      user = await loader.load(1)
+
+      reveal_type(user)
+  out: |
+    main:8: note: Revealed type is 'strawberry.dataloader.DataLoader[Any, Any]'
+    main:13: note: Revealed type is 'Any'
+
+- case: test_dataloader
+  main: |
+    from typing import List
+
+    from strawberry.dataloader import DataLoader
+
+    async def load(keys: List[int]) -> List[str]:
+        return [str(k) for k in keys]
+
+    loader = DataLoader[int, str](load)
+
+    reveal_type(loader)
+
+    async def run() -> None:
+      user = await loader.load(1)
+
+      reveal_type(user)
+  out: |
+    main:10: note: Revealed type is 'strawberry.dataloader.DataLoader[builtins.int*, builtins.str*]'
+    main:15: note: Revealed type is 'builtins.str*'
diff --git a/tests/mypy/test_enum.yml b/tests/mypy/test_enum.yml
--- a/tests/mypy/test_enum.yml
+++ b/tests/mypy/test_enum.yml
@@ -43,18 +43,3 @@
   out: |
     main:19: note: Revealed type is 'Any'
     main:20: note: Revealed type is 'Any'
-
-
-- case: test_enum_from_namespace
-  main: |
-    import enums
-    import strawberry
-
-    Flavour = strawberry.enum(enums.IceCreamFlavour)
-
-    a: Flavour = Flavour.VANILLA
-    reveal_type(Flavour)
-    reveal_type(a)
-  out: |
-    main:7: note: Revealed type is 'def (value: builtins.object) -> enums.IceCreamFlavour'
-    main:8: note: Revealed type is 'enums.IceCreamFlavour'
diff --git a/tests/mypy/test_fields.yml b/tests/mypy/test_fields.yml
--- a/tests/mypy/test_fields.yml
+++ b/tests/mypy/test_fields.yml
@@ -4,7 +4,7 @@
 
     @strawberry.type
     class User:
-        name: str = strawberry.field(description='Example')
+        name: str = strawberry.field(description='Example')  # type: ignore
 
     User(name="Patrick")
     User(n="Patrick")
@@ -21,9 +21,9 @@
     @strawberry.type
     class Example:
         a: str
-        b: str = strawberry.field(name="b")
-        c: str = strawberry.field(name="c", resolver=some_resolver)
-        d: str = strawberry.field(resolver=some_resolver)
+        b: str = strawberry.field(name="b")  # type: ignore
+        c: str = strawberry.field(name="c", resolver=some_resolver)  # type: ignore
+        d: str = strawberry.field(resolver=some_resolver)  # type: ignore
 
         @strawberry.field(description="ABC")
         def e(self, info) -> str:
@@ -44,8 +44,8 @@
     main:22: note: Revealed type is 'builtins.str'
     main:23: note: Revealed type is 'builtins.str'
     main:24: note: Revealed type is 'builtins.str'
-    main:25: note: Revealed type is 'Any'
-    main:26: note: Revealed type is 'Any'
+    main:25: note: Revealed type is 'def (self: main.Example, info: Any) -> builtins.str'
+    main:26: note: Revealed type is 'def (self: main.Example, info: Any) -> builtins.str'
 
 - case: test_private_field
   main: |
diff --git a/tests/schema/test_dataloaders.py b/tests/schema/test_dataloaders.py
new file mode 100644
--- /dev/null
+++ b/tests/schema/test_dataloaders.py
@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from typing import List
+
+import pytest
+
+import strawberry
+from strawberry.dataloader import DataLoader
+
+
+@pytest.mark.asyncio
+async def test_can_use_dataloaders(mocker):
+    @dataclass
+    class User:
+        id: str
+
+    async def idx(keys) -> List[User]:
+        return [User(key) for key in keys]
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader)
+
+    @strawberry.type
+    class Query:
+        @strawberry.field
+        async def get_user(self, id: strawberry.ID) -> str:
+            user = await loader.load(id)
+
+            return user.id
+
+    schema = strawberry.Schema(query=Query)
+
+    query = """{
+        a: getUser(id: "1")
+        b: getUser(id: "2")
+    }"""
+
+    result = await schema.execute(query)
+
+    assert not result.errors
+    assert result.data == {
+        "a": "1",
+        "b": "2",
+    }
+
+    mock_loader.assert_called_once_with(["1", "2"])
diff --git a/tests/test_dataloaders.py b/tests/test_dataloaders.py
new file mode 100644
--- /dev/null
+++ b/tests/test_dataloaders.py
@@ -0,0 +1,196 @@
+import asyncio
+
+import pytest
+
+from strawberry.dataloader import DataLoader
+from strawberry.exceptions import WrongNumberOfResultsReturned
+
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_loading():
+    async def idx(keys):
+        return keys
+
+    loader = DataLoader(load_fn=idx)
+
+    value_a = await loader.load(1)
+    value_b = await loader.load(2)
+    value_c = await loader.load(3)
+
+    assert value_a == 1
+    assert value_b == 2
+    assert value_c == 3
+
+
+async def test_gathering(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader)
+
+    [value_a, value_b, value_c] = await asyncio.gather(
+        loader.load(1),
+        loader.load(2),
+        loader.load(3),
+    )
+
+    mock_loader.assert_called_once_with([1, 2, 3])
+
+    assert value_a == 1
+    assert value_b == 2
+    assert value_c == 3
+
+
+async def test_max_batch_size(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader, max_batch_size=2)
+
+    [value_a, value_b, value_c] = await asyncio.gather(
+        loader.load(1),
+        loader.load(2),
+        loader.load(3),
+    )
+
+    mock_loader.assert_has_calls([mocker.call([1, 2]), mocker.call([3])])
+
+    assert value_a == 1
+    assert value_b == 2
+    assert value_c == 3
+
+
+async def test_error():
+    async def idx(keys):
+        return [ValueError()]
+
+    loader = DataLoader(load_fn=idx)
+
+    with pytest.raises(ValueError):
+        await loader.load(1)
+
+
+async def test_error_and_values():
+    async def idx(keys):
+        if keys == [2]:
+            return [2]
+
+        return [ValueError()]
+
+    loader = DataLoader(load_fn=idx)
+
+    with pytest.raises(ValueError):
+        await loader.load(1)
+
+    assert await loader.load(2) == 2
+
+
+async def test_when_raising_error_in_loader():
+    async def idx(keys):
+        raise ValueError()
+
+    loader = DataLoader(load_fn=idx)
+
+    with pytest.raises(ValueError):
+        await loader.load(1)
+
+    with pytest.raises(ValueError):
+        await asyncio.gather(
+            loader.load(1),
+            loader.load(2),
+            loader.load(3),
+        )
+
+
+async def test_returning_wrong_number_of_results():
+    async def idx(keys):
+        return [1, 2]
+
+    loader = DataLoader(load_fn=idx)
+
+    with pytest.raises(
+        WrongNumberOfResultsReturned,
+        match=(
+            "Received wrong number of results in dataloader, "
+            "expected: 1, received: 2"
+        ),
+    ):
+        await loader.load(1)
+
+
+async def test_caches_by_id(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader, cache=True)
+
+    a = loader.load(1)
+    b = loader.load(1)
+
+    assert a == b
+
+    assert await a == 1
+    assert await b == 1
+
+    mock_loader.assert_called_once_with([1])
+
+
+async def test_caches_by_id_when_loading_many(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader, cache=True)
+
+    a = loader.load(1)
+    b = loader.load(1)
+
+    assert a == b
+
+    assert [1, 1] == await asyncio.gather(a, b)
+
+    mock_loader.assert_called_once_with([1])
+
+
+async def test_cache_disabled(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader, cache=False)
+
+    a = loader.load(1)
+    b = loader.load(1)
+
+    assert a != b
+
+    assert await a == 1
+    assert await b == 1
+
+    mock_loader.assert_has_calls([mocker.call([1, 1])])
+
+
+async def test_cache_disabled_immediate_await(mocker):
+    async def idx(keys):
+        return keys
+
+    mock_loader = mocker.Mock(side_effect=idx)
+
+    loader = DataLoader(load_fn=mock_loader, cache=False)
+
+    a = await loader.load(1)
+    b = await loader.load(1)
+
+    assert a == b
+
+    mock_loader.assert_has_calls([mocker.call([1]), mocker.call([1])])