Allow 2-arity preload function (#4438)

integration_test/cases/preload.exs

@@ -329,7 +329,7 @@ defmodule Ecto.Integration.PreloadTest do
 
   ## With queries
 
-  test "preload with function" do
+  test "preload with 1-arity function" do
     p1 = TestRepo.insert!(%Post{title: "1"})
     p2 = TestRepo.insert!(%Post{title: "2"})
     p3 = TestRepo.insert!(%Post{title: "3"})
@@ -347,6 +347,23 @@ defmodule Ecto.Integration.PreloadTest do
     assert [] = pe3.comments
   end
 
+  test "preload with 2-arity function" do
+    p = TestRepo.insert!(%Post{title: "1"})
+    c1 = TestRepo.insert!(%Comment{post_id: p.id})
+    c2 = TestRepo.insert!(%Comment{post_id: p.id})
+
+    # making a simple preloader so that it works across all adapters
+    preloader = fn parent_ids, assoc ->
+      %{related_key: related_key, queryable: queryable} = assoc
+
+      from(q in queryable, where: field(q, ^related_key) in ^parent_ids, order_by: q.id)
+      |> TestRepo.all()
+    end
+
+    assert p = TestRepo.preload(p, comments: preloader)
+    assert [^c1, ^c2] = p.comments
+  end
+
   test "preload many_to_many with function" do
     p1 = TestRepo.insert!(%Post{title: "1"})
     p2 = TestRepo.insert!(%Post{title: "2"})

lib/ecto/query.ex

@@ -2724,9 +2724,11 @@ defmodule Ecto.Query do
 
   ## Preload functions
 
-  Preload also allows functions to be given. In such cases, the function
-  receives the IDs of the parent association and it must return the associated
-  data. Ecto then will map this data and sort it by the relationship key:
+  Preload also allows functions to be given. If the function has an arity of 1,
+  it receives only the IDs of the parent association. If it has an arity of 2, it
+  receives the IDS of the parent association as the first argument and the association
+  metadata as the second argument. Both functions must return the associated data.
+  Ecto then will map this data and sort it by the relationship key:
 
       comment_preloader = fn post_ids -> fetch_comments_by_post_ids(post_ids) end
       Repo.all from p in Post, preload: [comments: ^comment_preloader]
@@ -2759,6 +2761,33 @@ defmodule Ecto.Query do
       function, the function will receive a list of "post_ids" as the argument
       and it must return a tuple in the format of `{post_id, tag}`
 
+  The 2-arity version of the function is especially useful if you would like to
+  build a general preloader that works across all associations. For example, if
+  you would like to build a preloader for lateral joins that finds the newest
+  associations you may do the following:
+
+      lateral_preloader = fn ids, assoc -> newest_records(ids, assoc, 5) end
+
+      def newest_records(parent_ids, assoc, n) do
+        %{related_key: related_key, queryable: queryable} = assoc
+
+        squery =
+          from q in queryable,
+            where: field(q, ^related_key) == parent_as(:parent_ids).id,
+            order_by: {:desc, :created_at},
+            limit: ^n
+
+        query =
+          from f in fragment("SELECT id from UNNEST(?::int[]) AS id", ^parent_ids), as: :parent_ids,
+            inner_lateral_join: s in subquery(squery), on: true,
+            select: s
+
+        Repo.all(query)
+      end
+
+  For the list of available metadata, see the module documentation of the association types.
+  For example, see `Ecto.Association.BelongsTo`.
+
   ## Dynamic preloads
 
   Preloads can also be specified dynamically using the `dynamic` macro:

lib/ecto/repo/preloader.ex

@@ -277,12 +277,13 @@ defmodule Ecto.Repo.Preloader do
     end
   end
 
-  defp fetch_query(ids, assoc, _repo_name, query, _prefix, related_key, _take, _tuplet) when is_function(query, 1) do
+  defp fetch_query(ids, assoc, _repo_name, query, _prefix, related_key, _take, _tuplet)
+       when is_function(query, 1) or is_function(query, 2) do
     # Note we use an explicit sort because we don't want
     # to reorder based on the struct. Only the ID.
     ids
-    |> Enum.uniq
-    |> query.()
+    |> Enum.uniq()
+    |> preload_function(assoc, query)
     |> fetched_records_to_tuple_ids(assoc, related_key)
     |> Enum.sort(fn {id1, _}, {id2, _} -> id1 <= id2 end)
     |> unzip_ids([], [])
@@ -323,6 +324,9 @@ defmodule Ecto.Repo.Preloader do
     unzip_ids Ecto.Repo.Queryable.all(repo_name, query, tuplet), [], []
   end
 
+  defp preload_function(ids, _assoc, query) when is_function(query, 1), do: query.(ids)
+  defp preload_function(ids, assoc, query) when is_function(query, 2), do: query.(ids, assoc)
+
   defp fetched_records_to_tuple_ids([], _assoc, _related_key),
     do: []
 
@@ -573,13 +577,13 @@ defmodule Ecto.Repo.Preloader do
   end
 
   defp normalize_each({atom, {query, list}}, acc, take, original)
-       when is_atom(atom) and (is_map(query) or is_function(query, 1)) do
+       when is_atom(atom) and (is_map(query) or is_function(query, 1) or is_function(query, 2)) do
     fields = take(take, atom)
     [{atom, {fields, query!(query), normalize_each(wrap(list, original), [], fields, original)}}|acc]
   end
 
   defp normalize_each({atom, query}, acc, take, _original)
-       when is_atom(atom) and (is_map(query) or is_function(query, 1)) do
+       when is_atom(atom) and (is_map(query) or is_function(query, 1) or is_function(query, 2)) do
     [{atom, {take(take, atom), query!(query), []}}|acc]
   end
 
@@ -597,6 +601,7 @@ defmodule Ecto.Repo.Preloader do
   end
 
   defp query!(query) when is_function(query, 1), do: query
+  defp query!(query) when is_function(query, 2), do: query
   defp query!(%Ecto.Query{} = query), do: query
 
   defp take(take, field) do