Support associations on composite foreign keys

Author: soundmonster

Diff for: Earthfile

@@ -60,7 +60,7 @@ integration-test-base:
         apk del .build-dependencies && rm -f msodbcsql*.sig mssql-tools*.apk
     ENV PATH="/opt/mssql-tools/bin:${PATH}"
 
-    GIT CLONE https://github.com/elixir-ecto/ecto_sql.git /src/ecto_sql
+    GIT CLONE --branch composite_foreign_keys https://github.com/soundmonster/ecto_sql.git /src/ecto_sql
     WORKDIR /src/ecto_sql
     RUN mix deps.get
 

Diff for: integration_test/cases/assoc.exs

@@ -10,6 +10,7 @@ defmodule Ecto.Integration.AssocTest do
   alias Ecto.Integration.PostUser
   alias Ecto.Integration.Comment
   alias Ecto.Integration.Permalink
+  alias Ecto.Integration.CompositePk
 
   test "has_many assoc" do
     p1 = TestRepo.insert!(%Post{title: "1"})
@@ -55,6 +56,22 @@ defmodule Ecto.Integration.AssocTest do
     assert p2.id == pid2
   end
 
+  test "belongs_to assoc with composite key" do
+    TestRepo.insert!(%CompositePk{a: 2, b: 1, name: "foo"})
+    TestRepo.insert!(%CompositePk{a: 2, b: 2, name: "bar"})
+    TestRepo.insert!(%CompositePk{a: 2, b: 3, name: "unused"})
+
+    p1 = TestRepo.insert!(%Post{title: "first", composite_a: 2, composite_b: 1})
+    p2 = TestRepo.insert!(%Post{title: "none"})
+    p3 = TestRepo.insert!(%Post{title: "second", composite_a: 2, composite_b: 2})
+
+    assert [c1, c2] = TestRepo.all Ecto.assoc([p1, p2, p3], :composite)
+    assert c1.a == 2
+    assert c1.b == 1
+    assert c2.a == 2
+    assert c2.b == 2
+  end
+
   test "has_many through assoc" do
     p1 = TestRepo.insert!(%Post{})
     p2 = TestRepo.insert!(%Post{})
@@ -725,6 +742,27 @@ defmodule Ecto.Integration.AssocTest do
     assert perma.post_id == nil
   end
 
+  test "belongs_to changeset assoc on composite key" do
+    changeset =
+      %CompositePk{a: 1, b: 2}
+      |> Ecto.Changeset.change()
+      |> Ecto.Changeset.put_assoc(:posts, [%Post{title: "1"}])
+
+    composite = TestRepo.insert!(changeset)
+    assert [post] = composite.posts
+    assert post.id
+    assert post.composite_a == composite.a
+    assert post.composite_b == composite.b
+    assert post.title == "1"
+    # TODO Repo.get should work with composite keys somehow, right?
+    # composite = TestRepo.get! from(Composite, preload: [:post]), [composite.a, composite.b]
+    # assert composite.post.title == "1"
+
+    post = TestRepo.get! from(Post, preload: [:composite]), post.id
+    assert post.composite.a == 1
+    assert post.composite.b == 2
+  end
+
   test "inserting struct with associations" do
     tree = %Permalink{
       url: "root",
@@ -750,6 +788,14 @@ defmodule Ecto.Integration.AssocTest do
     assert Enum.all?(tree.post.comments, & &1.id)
   end
 
+  test "inserting struct with associations on composite keys" do
+    # creates nested belongs_to
+    %Post{composite: composite} =
+      TestRepo.insert! %Post{title: "1", composite: %CompositePk{a: 1, b: 2, name: "name"}}
+
+    assert %CompositePk{a: 1, b: 2, name: "name"} = composite
+  end
+
   test "inserting struct with empty associations" do
     permalink = TestRepo.insert!(%Permalink{url: "root", post: nil})
     assert permalink.post == nil

Diff for: integration_test/cases/preload.exs

@@ -6,6 +6,7 @@ defmodule Ecto.Integration.PreloadTest do
 
   alias Ecto.Integration.Post
   alias Ecto.Integration.Comment
+  alias Ecto.Integration.CompositePk
   alias Ecto.Integration.Item
   alias Ecto.Integration.Permalink
   alias Ecto.Integration.User
@@ -341,6 +342,25 @@ defmodule Ecto.Integration.PreloadTest do
     assert [] = pe3.comments
   end
 
+  test "preload composite foreign key with function" do
+    c11 = TestRepo.insert!(%CompositePk{a: 1, b: 1, name: "11"})
+    c12 = TestRepo.insert!(%CompositePk{a: 1, b: 2, name: "12"})
+    c22 = TestRepo.insert!(%CompositePk{a: 2, b: 2, name: "22"})
+    c33 = TestRepo.insert!(%CompositePk{a: 3, b: 3, name: "33"})
+
+    TestRepo.insert!(%Post{title: "1", composite_a: 1, composite_b: 1})
+    TestRepo.insert!(%Post{title: "2", composite_a: 1, composite_b: 1})
+    TestRepo.insert!(%Post{title: "3", composite_a: 1, composite_b: 2})
+    TestRepo.insert!(%Post{title: "4", composite_a: 2, composite_b: 2})
+
+    assert [ce12, ce11, ce33, ce22] = TestRepo.preload([c12, c11, c33, c22],
+                                              posts: fn _ -> TestRepo.all(Post) end)
+    assert [%Post{title: "1"}, %Post{title: "2"}] = ce11.posts
+    assert [%Post{title: "3"}] = ce12.posts
+    assert [%Post{title: "4"}] = ce22.posts
+    assert [] = ce33.posts
+  end
+
   test "preload many_to_many with function" do
     p1 = TestRepo.insert!(%Post{title: "1"})
     p2 = TestRepo.insert!(%Post{title: "2"})
@@ -389,6 +409,50 @@ defmodule Ecto.Integration.PreloadTest do
     assert p3.users == [%{id: uid1}, %{id: uid4}]
   end
 
+  test "preload many_to_many on composite foreign keys with function" do
+    c11 = TestRepo.insert!(%CompositePk{a: 1, b: 1, name: "11"})
+    c12 = TestRepo.insert!(%CompositePk{a: 1, b: 2, name: "12"})
+    c22 = TestRepo.insert!(%CompositePk{a: 2, b: 2, name: "22"})
+
+    TestRepo.insert_all "composite_pk_composite_pk", [[a_1: 1, b_1: 1, a_2: 1, b_2: 2],
+                                                      [a_1: 1, b_1: 1, a_2: 2, b_2: 2],
+                                                      [a_1: 1, b_1: 2, a_2: 1, b_2: 1],
+                                                      [a_1: 2, b_1: 2, a_2: 2, b_2: 2]]
+
+    wrong_preloader = fn composite_ids ->
+      composite_ids_a = Enum.map(composite_ids, &Enum.at(&1, 0))
+      composite_ids_b = Enum.map(composite_ids, &Enum.at(&1, 1))
+      TestRepo.all(
+        from c in CompositePk,
+             join: cc in "composite_pk_composite_pk",
+             where: cc.a_1 in ^composite_ids_a and cc.b_1 in ^composite_ids_b and cc.a_2 == c.a and cc.b_2 == c.b,
+             order_by: [c.a, c.b],
+             select: map(c, [:name])
+      )
+    end
+
+    assert_raise RuntimeError, ~r/invalid custom preload for `composites` on `Ecto.Integration.CompositePk`/, fn ->
+      TestRepo.preload([c11, c12, c22], composites: wrong_preloader)
+    end
+
+    right_preloader = fn composite_ids ->
+      composite_ids_a = Enum.map(composite_ids, &Enum.at(&1, 0))
+      composite_ids_b = Enum.map(composite_ids, &Enum.at(&1, 1))
+      TestRepo.all(
+        from c in CompositePk,
+             join: cc in "composite_pk_composite_pk",
+             where: cc.a_1 in ^composite_ids_a and cc.b_1 in ^composite_ids_b and cc.a_2 == c.a and cc.b_2 == c.b,
+             order_by: [c.a, c.b],
+             select: {[cc.a_1, cc.b_1], map(c, [:name])}
+      )
+    end
+
+    [c11, c12, c22] = TestRepo.preload([c11, c12, c22], composites: right_preloader)
+    assert c11.composites == [%{name: "12"}, %{name: "22"}]
+    assert c12.composites == [%{name: "11"}]
+    assert c22.composites == [%{name: "22"}]
+  end
+
   test "preload with query" do
     p1 = TestRepo.insert!(%Post{title: "1"})
     p2 = TestRepo.insert!(%Post{title: "2"})
@@ -604,6 +668,18 @@ defmodule Ecto.Integration.PreloadTest do
     assert TestRepo.preload(updated, [:author], force: true).author == nil
   end
 
+  test "preload raises with association over composite foreign key is set but without id" do
+    p1 = TestRepo.insert!(%Post{title: "1"})
+    c11 = TestRepo.insert!(%CompositePk{a: 1, b: 1, name: "11"})
+    updated = %{p1 | composite: c11, composite_a: nil, composite_b: nil}
+
+    assert ExUnit.CaptureLog.capture_log(fn ->
+      assert TestRepo.preload(updated, [:composite]).composite == c11
+    end) =~ ~r/its association keys `\[composite_a, composite_b\]` are nil/
+
+    assert TestRepo.preload(updated, [:composite], force: true).composite == nil
+  end
+
   test "preload skips already loaded for cardinality one" do
     %Post{id: pid} = TestRepo.insert!(%Post{title: "1"})
 

Diff for: integration_test/cases/repo.exs

@@ -152,6 +152,24 @@ defmodule Ecto.Integration.RepoTest do
     assert TestRepo.all(PostUserCompositePk) == []
   end
 
+  @tag :composite_pk
+  # TODO this needs a better name
+  test "insert, update and delete with associated composite pk #2" do
+    composite = TestRepo.insert!(%CompositePk{a: 1, b: 2, name: "name"})
+    post = TestRepo.insert!(%Post{title: "post title", composite: composite})
+
+    assert post.composite_a == 1
+    assert post.composite_b == 2
+    assert TestRepo.get_by!(CompositePk, [a: 1, b: 2]) == composite
+
+    post = post |> Ecto.Changeset.change(composite: nil) |> TestRepo.update!
+    assert is_nil(post.composite_a)
+    assert is_nil(post.composite_b)
+
+    TestRepo.delete!(post)
+    assert TestRepo.all(CompositePk) == [composite]
+  end
+
   @tag :invalid_prefix
   test "insert, update and delete with invalid prefix" do
     post = TestRepo.insert!(%Post{})

Diff for: integration_test/support/schemas.exs

@@ -54,6 +54,8 @@ defmodule Ecto.Integration.Post do
     has_one :update_permalink, Ecto.Integration.Permalink, foreign_key: :post_id, on_delete: :delete_all, on_replace: :update
     has_many :comments_authors, through: [:comments, :author]
     belongs_to :author, Ecto.Integration.User
+    belongs_to :composite, Ecto.Integration.CompositePk,
+      foreign_key: [:composite_a, :composite_b], references: [:a, :b], type: [:integer, :integer], on_replace: :nilify
     many_to_many :users, Ecto.Integration.User,
       join_through: "posts_users", on_delete: :delete_all, on_replace: :delete
     many_to_many :ordered_users, Ecto.Integration.User, join_through: "posts_users", preload_order: [desc: :name]
@@ -291,6 +293,10 @@ defmodule Ecto.Integration.CompositePk do
     field :a, :integer, primary_key: true
     field :b, :integer, primary_key: true
     field :name, :string
+    has_many :posts, Ecto.Integration.Post, foreign_key: [:composite_a, :composite_b], references: [:a, :b]
+    many_to_many :composites, Ecto.Integration.CompositePk,
+      join_through: "composite_pk_composite_pk", join_keys: [[a_1: :a, b_1: :b], [a_2: :a, b_2: :b]],
+      on_delete: :delete_all, on_replace: :delete
   end
   def changeset(schema, params) do
     cast(schema, params, ~w(a b name)a)

Diff for: lib/ecto.ex

@@ -510,10 +510,15 @@ defmodule Ecto do
     refl = %{owner_key: owner_key} = Ecto.Association.association_from_schema!(schema, assoc)
 
     values =
-      Enum.uniq for(struct <- structs,
-        assert_struct!(schema, struct),
-        key = Map.fetch!(struct, owner_key),
-        do: key)
+      structs
+      |> Enum.filter(&assert_struct!(schema, &1))
+      |> Enum.map(fn struct ->
+        owner_key
+        # TODO remove List.wrap once all assocs use lists
+        |> List.wrap
+        |> Enum.map(&Map.fetch!(struct, &1))
+      end)
+      |> Enum.uniq
 
     case assocs do
       [] ->