pulled in the assoc helper from absinthe ecto

assoc.ex

defmodule MyWeb.Schema.Shared do
  @moduledoc false

  import Absinthe.Resolution.Helpers
  import Ecto.Query

  def assoc(assoc_key) do
    fn data, _, _ ->
      Lol.Repo |> ecto_batch(data, assoc_key)
    end
  end

  def ecto_batch(repo, %model{} = parent, association, callback \\ &default_callback/1) do
    {assoc_field, query_fun} = normalize(association)
    assoc = model.__schema__(:association, assoc_field)

    %{owner: owner, owner_key: owner_key, field: field} = assoc

    id = Map.fetch!(parent, owner_key)

    query = resolve_query(query_fun, assoc, parent)
    meta = {repo, owner, owner_key, field, query, self()}

    batch({__MODULE__, :perform_batch, meta}, id, fn results ->
      results
      |> Map.get(id)
      |> callback.()
    end)
  end

  defp default_callback(result) do
    {:ok, result}
  end

  defp normalize(field) when is_atom(field), do: {field, nil}
  defp normalize({field, query_fun}), do: {field, query_fun}

  # Query is resolved with `Repo.preload(association_name: query)`, so it must
  # return the association type.
  defp resolve_query(query_fun, %{queryable: queryable}, _) when is_function(query_fun),
    do: query_fun.(from(queryable))

  defp resolve_query(query_fun, %{field: field}, parent) when is_function(query_fun),
    do: query_fun.(Ecto.assoc(parent, field))

  defp resolve_query(_, _, _), do: nil

  @doc false
  # this has to be public because it gets called from the absinthe batcher
  def perform_batch({repo, owner, owner_key, field, query, caller}, ids) do
    unique_ids = ids |> MapSet.new() |> MapSet.to_list()
    preload = if query != nil, do: [{field, query}], else: field

    unique_ids
    |> Enum.map(&Map.put(struct(owner), owner_key, &1))
    |> repo.preload(preload, caller: caller)
    |> Enum.map(&{Map.get(&1, owner_key), Map.get(&1, field)})
    |> Map.new()
  end
end