Absinthe dataloader usage example

example.exs

Mix.install([:dataloader, :absinthe])

defmodule User do
  defstruct [:id, :name]

  def get(id) do
    Enum.find(list(), &(&1.id == id))
  end

  def list do
    [
      %__MODULE__{id: "1", name: "John"},
      %__MODULE__{id: "2", name: "Alice"},
      %__MODULE__{id: "3", name: "Bob"}
    ]
  end
end

defmodule Post do
  defstruct [:id, :author_id, :content]

  def list(filter \\ %{}) do
    for post <- [
          %__MODULE__{id: "1", author_id: "1", content: "post 1"},
          %__MODULE__{id: "2", author_id: "2", content: "post 2"},
          %__MODULE__{id: "3", author_id: "3", content: "post 3"}
        ],
        Map.take(post, Map.keys(filter)) == filter,
        do: post
  end
end

defmodule Comment do
  defstruct [:author_id, :post_id, :content]

  def list(filter \\ %{}) do
    for comment <- [
          %__MODULE__{post_id: "1", author_id: "2", content: "comment 1"},
          %__MODULE__{post_id: "1", author_id: "2", content: "comment 2"},
          %__MODULE__{post_id: "1", author_id: "3", content: "comment 3"},
          %__MODULE__{post_id: "2", author_id: "1", content: "comment 1"},
          %__MODULE__{post_id: "3", author_id: "2", content: "comment 1"}
        ],
        Map.take(comment, Map.keys(filter)) == filter,
        do: comment
  end
end

defmodule DataSource do
  def new do
    Dataloader.KV.new(&load/2)
  end

  defp load({:author, args}, batch) when map_size(args) == 0 do
    Map.new(batch, fn %{author_id: author_id} = assoc ->
      {assoc, User.get(author_id)}
    end)
  end

  defp load({:posts, args}, batch) when map_size(args) == 0 do
    Map.new(batch, fn %User{} = user ->
      filter = %{author_id: user.id}
      {user, Post.list(filter)}
    end)
  end

  defp load({:comments, args}, batch) do
    Map.new(batch, fn %Post{} = post ->
      filter = Map.merge(args, %{post_id: post.id})
      {post, Comment.list(filter)}
    end)
  end
end

ExUnit.start()
ExUnit.run()

defmodule Schema do
  use Absinthe.Schema
  import Absinthe.Resolution.Helpers, only: [dataloader: 1]

  object :user do
    field(:id, :id)
    field(:name, :string)
    field(:posts, list_of(:post), resolve: dataloader(:example))
  end

  object :post do
    field(:id, :id)
    field(:content, :string)
    field(:author, :user, resolve: dataloader(:example))

    field :comments, list_of(:comment) do
      arg(:author_id, :id)
      resolve(dataloader(:example))
    end
  end

  object :comment do
    field(:id, :id)
    field(:content, :string)
    field(:author, :user, resolve: dataloader(:example))
  end

  query do
    field :user, :user do
      resolve(fn _args, %{context: context} -> {:ok, context.current_user} end)
    end

    field :posts, list_of(:post) do
      arg(:author_id, :id)
      resolve(fn args, _res -> {:ok, Post.list(args)} end)
    end
  end

  def context(ctx) do
    loader = Dataloader.add_source(Dataloader.new(), :example, DataSource.new())
    current_user = User.get("1")

    ctx
    |> Map.put(:loader, loader)
    |> Map.put(:current_user, current_user)
  end

  def plugins do
    [Absinthe.Middleware.Dataloader | Absinthe.Plugin.defaults()]
  end
end

defmodule SchemaTest do
  use ExUnit.Case, async: true

  test "posts" do
    document = """
      query {
        posts {
          content
          author {
            name
          }
          comments {
            content
            author {
              name
            }
          }
        }
      }
    """

    assert {:ok, %{data: response}} = Absinthe.run(document, Schema)

    assert response == %{
             "posts" => [
               %{
                 "author" => %{"name" => "John"},
                 "content" => "post 1",
                 "comments" => [
                   %{"author" => %{"name" => "Alice"}, "content" => "comment 1"},
                   %{"author" => %{"name" => "Alice"}, "content" => "comment 2"},
                   %{"author" => %{"name" => "Bob"}, "content" => "comment 3"}
                 ]
               },
               %{
                 "author" => %{"name" => "Alice"},
                 "content" => "post 2",
                 "comments" => [
                   %{"author" => %{"name" => "John"}, "content" => "comment 1"}
                 ]
               },
               %{
                 "author" => %{"name" => "Bob"},
                 "content" => "post 3",
                 "comments" => [
                   %{"author" => %{"name" => "Alice"}, "content" => "comment 1"}
                 ]
               }
             ]
           }
  end

  test "post comments" do
    document = """
      query {
        posts(author_id: 1) {
          comments(author_id: 2) {
            content
          }
        }
      }
    """

    assert {:ok, %{data: response}} = Absinthe.run(document, Schema)

    assert response == %{
             "posts" => [
               %{
                 "comments" => [
                   %{"content" => "comment 1"},
                   %{"content" => "comment 2"}
                 ]
               }
             ]
           }
  end

  test "user posts" do
    document = """
      query {
        user {
          posts {
            content
          }
        }
      }
    """

    assert {:ok, %{data: response}} = Absinthe.run(document, Schema)

    assert response == %{
             "user" => %{
               "posts" => [
                 %{
                   "content" => "post 1"
                 }
               ]
             }
           }
  end
end