Allow 2-arity functions as preload function in Query (#4659)

Author: orlera

lib/ecto/query/builder/preload.ex

@@ -309,6 +309,7 @@ defmodule Ecto.Query.Builder.Preload do
 
   defp assert_query_or_fun!(%Ecto.Query{}, _key), do: :ok
   defp assert_query_or_fun!(fun, _key) when is_function(fun, 1), do: :ok
+  defp assert_query_or_fun!(fun, _key) when is_function(fun, 2), do: :ok
 
   defp assert_query_or_fun!(other, key) do
     raise ArgumentError,

test/ecto/query/builder/preload_test.exs

@@ -66,9 +66,9 @@ defmodule Ecto.Query.Builder.PreloadTest do
     test "supports dynamics for join association bindings using named bindings" do
       comments = :comments
 
-      query = 
-        from p in "posts", 
-        join: c in assoc(p, :comments), 
+      query =
+        from p in "posts",
+        join: c in assoc(p, :comments),
         as: ^comments
       preloads = [
         comments: dynamic([{^comments, c}], c)
@@ -120,6 +120,18 @@ defmodule Ecto.Query.Builder.PreloadTest do
       preloads = [comments: outer_dynamic]
       assert %{preloads: [], assocs: [comments: {1, []}]} = preload(query, ^preloads)
     end
+
+    test "supports preload functions" do
+      query = from p in "posts", join: c in assoc(p, :comments), as: :comments
+      function_1 = fn _ -> [] end
+      function_2 = fn _, _ -> [] end
+
+      preloads = [comments: function_1]
+      assert %{preloads: ^preloads, assocs: []} = preload(query, ^preloads)
+
+      preloads = [comments: function_2]
+      assert %{preloads: ^preloads, assocs: []} = preload(query, ^preloads)
+    end
   end
 
   describe "invalid preload" do
@@ -175,5 +187,13 @@ defmodule Ecto.Query.Builder.PreloadTest do
         preload(query, ^[comments: dynamic([_, c], c.field)])
       end
     end
+
+    test "raises when preload function has more than two arguments" do
+      message = ~r"invalid preload for key `:comments`:"
+      assert_raise ArgumentError, message, fn ->
+        query = from p in "posts", join: c in assoc(p, :comments)
+        preload(query, ^[comments: fn _, _, _  -> [] end])
+      end
+    end
   end
 end