Support for composite foreign keys for belongs_to

Author: Leo Batyuk

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
 

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"})
@@ -750,6 +751,12 @@ 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

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{})

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,7 @@ 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]
   end
   def changeset(schema, params) do
     cast(schema, params, ~w(a b name)a)

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
       [] ->

lib/ecto/association.ex

@@ -35,7 +35,7 @@ defmodule Ecto.Association do
                required(:cardinality) => :one | :many,
                required(:relationship) => :parent | :child,
                required(:owner) => atom,
-               required(:owner_key) => atom,
+               required(:owner_key) => list(atom),
                required(:field) => atom,
                required(:unique) => boolean,
                optional(atom) => any}
@@ -71,7 +71,8 @@ defmodule Ecto.Association do
 
     * `:owner` - the owner module of the association
 
-    * `:owner_key` - the key in the owner with the association value
+    * `:owner_key` - the key in the owner with the association value, or a
+      list of keys for composite keys
 
     * `:relationship` - if the relationship to the specified schema is
       of a `:child` or a `:parent`
@@ -235,8 +236,15 @@ defmodule Ecto.Association do
     # for the final WHERE clause with values.
     {_, query, _, dest_out_key} = Enum.reduce(joins, {source, query, counter, source.out_key}, fn curr_rel, {prev_rel, query, counter, _} ->
       related_queryable = curr_rel.schema
-
-      next = join(query, :inner, [{src, counter}], dest in ^related_queryable, on: field(src, ^prev_rel.out_key) == field(dest, ^curr_rel.in_key))
+      # TODO remove this once all relations store keys in lists
+      in_keys = List.wrap(curr_rel.in_key)
+      out_keys = List.wrap(prev_rel.out_key)
+      next = query
+        # join on the first field of the foreign key
+        |> join(:inner, [{src, counter}], dest in ^related_queryable, on: field(src, ^hd(out_keys)) == field(dest, ^hd(in_keys)))
+        # add the rest of the foreign key fields, if any
+        |> composite_joins_query(counter, counter + 1, tl(out_keys), tl(in_keys))
+        # consider where clauses on assocs
         |> combine_joins_query(curr_rel.where, counter + 1)
 
       {curr_rel, next, counter + 1, curr_rel.out_key}
@@ -320,6 +328,16 @@ defmodule Ecto.Association do
     end)
   end
 
+  # TODO docs
+  def composite_joins_query(query, _binding_src, _binding_dst, [], []) do
+    query
+  end
+  def composite_joins_query(query, binding_src, binding_dst, [src_key | src_keys], [dst_key | dst_keys]) do
+    # TODO
+    [query, binding_src, binding_dst, [src_key | src_keys], [dst_key | dst_keys]] |> IO.inspect(label: :composite_joins_query)
+    query
+  end
+
   @doc """
   Add the default assoc query where clauses to a join.
 
@@ -335,6 +353,16 @@ defmodule Ecto.Association do
     %{query | joins: joins ++ [%{join_expr | on: %{join_on | expr: expr, params: params}}]}
   end
 
+  # TODO docs
+  def composite_assoc_query(query, _binding_src, [], []) do
+    query
+  end
+  def composite_assoc_query(query, binding_dst, [dst_key | dst_keys], [value | values]) do
+    # TODO
+    [query, binding_dst, [dst_key | dst_keys], [value | values]] |> IO.inspect(label: :composite_assoc_query)
+    query
+  end
+
   @doc """
   Add the default assoc query where clauses a provided query.
   """
@@ -632,6 +660,10 @@ defmodule Ecto.Association do
 
   defp primary_key!(nil), do: []
   defp primary_key!(struct), do: Ecto.primary_key!(struct)
+
+  def missing_fields(queryable, related_key) do
+    Enum.filter related_key, &is_nil(queryable.__schema__(:type, &1))
+  end
 end
 
 defmodule Ecto.Association.Has do
@@ -644,8 +676,8 @@ defmodule Ecto.Association.Has do
     * `field` - The name of the association field on the schema
     * `owner` - The schema where the association was defined
     * `related` - The schema that is associated
-    * `owner_key` - The key on the `owner` schema used for the association
-    * `related_key` - The key on the `related` schema used for the association
+    * `owner_key` - The list of columns that form the key on the `owner` schema used for the association
+    * `related_key` - The list of columns that form the key on the `related` schema used for the association
     * `queryable` - The real query to use for querying association
     * `on_delete` - The action taken on associations when schema is deleted
     * `on_replace` - The action taken on associations when schema is replaced
@@ -673,8 +705,8 @@ defmodule Ecto.Association.Has do
         {:error, "associated schema #{inspect queryable} does not exist"}
       not function_exported?(queryable, :__schema__, 2) ->
         {:error, "associated module #{inspect queryable} is not an Ecto schema"}
-      is_nil queryable.__schema__(:type, related_key) ->
-        {:error, "associated schema #{inspect queryable} does not have field `#{related_key}`"}
+      [] != (missing_fields = Ecto.Association.missing_fields(queryable, related_key)) ->
+        {:error, "associated schema #{inspect queryable} does not have field(s) `#{inspect missing_fields}`"}
       true ->
         :ok
     end
@@ -686,14 +718,17 @@ defmodule Ecto.Association.Has do
     cardinality = Keyword.fetch!(opts, :cardinality)
     related = Ecto.Association.related_from_query(queryable, name)
 
-    ref =
+    refs =
       module
       |> Module.get_attribute(:primary_key)
       |> get_ref(opts[:references], name)
+      |> List.wrap()
 
-    unless Module.get_attribute(module, :ecto_fields)[ref] do
-      raise ArgumentError, "schema does not have the field #{inspect ref} used by " <>
-        "association #{inspect name}, please set the :references option accordingly"
+    for ref <- refs do
+      unless Module.get_attribute(module, :ecto_fields)[ref] do
+        raise ArgumentError, "schema does not have the field #{inspect ref} used by " <>
+          "association #{inspect name}, please set the :references option accordingly"
+      end
     end
 
     if opts[:through] do
@@ -725,13 +760,19 @@ defmodule Ecto.Association.Has do
       raise ArgumentError, "expected `:where` for #{inspect name} to be a keyword list, got: `#{inspect where}`"
     end
 
+    foreign_key = case opts[:foreign_key] do
+      nil -> Enum.map(refs, &Ecto.Association.association_key(module, &1))
+      key when is_atom(key) -> [key]
+      keys when is_list(keys) -> keys
+    end
+
     %__MODULE__{
       field: name,
       cardinality: cardinality,
       owner: module,
       related: related,
-      owner_key: ref,
-      related_key: opts[:foreign_key] || Ecto.Association.association_key(module, ref),
+      owner_key: refs,
+      related_key: foreign_key,
       queryable: queryable,
       on_delete: on_delete,
       on_replace: on_replace,
@@ -756,19 +797,23 @@ defmodule Ecto.Association.Has do
 
   @impl true
   def joins_query(%{related_key: related_key, owner: owner, owner_key: owner_key, queryable: queryable} = assoc) do
-    from(o in owner, join: q in ^queryable, on: field(q, ^related_key) == field(o, ^owner_key))
+    # TODO find out how to handle a dynamic list of fields here
+    from(o in owner, join: q in ^queryable, on: field(q, ^hd(related_key)) == field(o, ^hd(owner_key)))
+    |> Ecto.Association.composite_joins_query(0, 1, tl(related_key), tl(owner_key))
     |> Ecto.Association.combine_joins_query(assoc.where, 1)
   end
 
   @impl true
   def assoc_query(%{related_key: related_key, queryable: queryable} = assoc, query, [value]) do
-    from(x in (query || queryable), where: field(x, ^related_key) == ^value)
+    from(x in (query || queryable), where: field(x, ^hd(related_key)) == ^hd(value))
+    |> Ecto.Association.composite_assoc_query(0, tl(related_key), tl(value))
     |> Ecto.Association.combine_assoc_query(assoc.where)
   end
 
   @impl true
   def assoc_query(%{related_key: related_key, queryable: queryable} = assoc, query, values) do
-    from(x in (query || queryable), where: field(x, ^related_key) in ^values)
+    from(x in (query || queryable), where: field(x, ^hd(related_key)) in ^Enum.map(values, &hd/1))
+    |> Ecto.Association.composite_assoc_query(0, tl(related_key), Enum.map(values, &tl/1))
     |> Ecto.Association.combine_assoc_query(assoc.where)
   end
 
@@ -807,16 +852,21 @@ defmodule Ecto.Association.Has do
     %{data: parent, repo: repo} = parent_changeset
     %{action: action, changes: changes} = changeset
 
-    {key, value} = parent_key(assoc, parent)
-    changeset = update_parent_key(changeset, action, key, value)
-    changeset = Ecto.Association.update_parent_prefix(changeset, parent)
+    parent_keys = parent_keys(assoc, parent)
+    changeset = Enum.reduce parent_keys, changeset, fn {key, value}, changeset ->
+      changeset = update_parent_key(changeset, action, key, value)
+      Ecto.Association.update_parent_prefix(changeset, parent)
+    end
 
     case apply(repo, action, [changeset, opts]) do
       {:ok, _} = ok ->
         if action == :delete, do: {:ok, nil}, else: ok
       {:error, changeset} ->
-        original = Map.get(changes, key)
-        {:error, put_in(changeset.changes[key], original)}
+        changeset = Enum.reduce parent_keys, changeset, fn {key, _}, changeset ->
+          original = Map.get(changes, key)
+          put_in(changeset.changes[key], original)
+        end
+        {:error, changeset}
     end
   end
 
@@ -825,11 +875,21 @@ defmodule Ecto.Association.Has do
   defp update_parent_key(changeset, _action, key, value),
     do: Ecto.Changeset.put_change(changeset, key, value)
 
-  defp parent_key(%{related_key: related_key}, nil) do
-    {related_key, nil}
+  defp parent_keys(%{related_key: related_keys}, nil) when is_list(related_keys) do
+    Enum.map related_keys, fn related_key -> {related_key, nil} end
+  end
+  defp parent_keys(%{related_key: related_key}, nil) do
+    [{related_key, nil}]
+  end
+  defp parent_keys(%{owner_key: owner_keys, related_key: related_keys}, owner) when is_list(owner_keys) and is_list(related_keys) do
+    owner_keys
+    |> Enum.zip(related_keys)
+    |> Enum.map(fn {owner_key, related_key} ->
+      {related_key, Map.get(owner, owner_key)}
+    end)
   end
-  defp parent_key(%{owner_key: owner_key, related_key: related_key}, owner) do
-    {related_key, Map.get(owner, owner_key)}
+  defp parent_keys(%{owner_key: owner_key, related_key: related_key}, owner) do
+    [{related_key, Map.get(owner, owner_key)}]
   end
 
   ## Relation callbacks
@@ -982,16 +1042,16 @@ defmodule Ecto.Association.BelongsTo do
         {:error, "associated schema #{inspect queryable} does not exist"}
       not function_exported?(queryable, :__schema__, 2) ->
         {:error, "associated module #{inspect queryable} is not an Ecto schema"}
-      is_nil queryable.__schema__(:type, related_key) ->
-        {:error, "associated schema #{inspect queryable} does not have field `#{related_key}`"}
+      [] != (missing_fields = Ecto.Association.missing_fields(queryable, related_key)) ->
+        {:error, "associated schema #{inspect queryable} does not have field(s) `#{inspect missing_fields}`"}
       true ->
         :ok
     end
   end
 
   @impl true
   def struct(module, name, opts) do
-    ref = if ref = opts[:references], do: ref, else: :id
+    refs = if ref = opts[:references], do: List.wrap(ref), else: [:id]
     queryable = Keyword.fetch!(opts, :queryable)
     related = Ecto.Association.related_from_query(queryable, name)
     on_replace = Keyword.get(opts, :on_replace, :raise)
@@ -1013,8 +1073,8 @@ defmodule Ecto.Association.BelongsTo do
       field: name,
       owner: module,
       related: related,
-      owner_key: Keyword.fetch!(opts, :foreign_key),
-      related_key: ref,
+      owner_key: List.wrap(Keyword.fetch!(opts, :foreign_key)),
+      related_key: refs,
       queryable: queryable,
       on_replace: on_replace,
       defaults: defaults,
@@ -1031,19 +1091,22 @@ defmodule Ecto.Association.BelongsTo do
 
   @impl true
   def joins_query(%{related_key: related_key, owner: owner, owner_key: owner_key, queryable: queryable} = assoc) do
-    from(o in owner, join: q in ^queryable, on: field(q, ^related_key) == field(o, ^owner_key))
+    from(o in owner, join: q in ^queryable, on: field(q, ^hd(related_key)) == field(o, ^hd(owner_key)))
+    |> Ecto.Association.composite_joins_query(0, 1, tl(related_key), tl(owner_key))
     |> Ecto.Association.combine_joins_query(assoc.where, 1)
   end
 
   @impl true
   def assoc_query(%{related_key: related_key, queryable: queryable} = assoc, query, [value]) do
-    from(x in (query || queryable), where: field(x, ^related_key) == ^value)
+    from(x in (query || queryable), where: field(x, ^hd(related_key)) == ^hd(value))
+    |> Ecto.Association.composite_assoc_query(0, tl(related_key), tl(value))
     |> Ecto.Association.combine_assoc_query(assoc.where)
   end
 
   @impl true
   def assoc_query(%{related_key: related_key, queryable: queryable} = assoc, query, values) do
-    from(x in (query || queryable), where: field(x, ^related_key) in ^values)
+    from(x in (query || queryable), where: field(x, ^hd(related_key)) in ^Enum.map(values, &hd/1))
+    |> Ecto.Association.composite_assoc_query(0, tl(related_key), Enum.map(values, &tl/1))
     |> Ecto.Association.combine_assoc_query(assoc.where)
   end
 
@@ -1264,11 +1327,12 @@ defmodule Ecto.Association.ManyToMany do
 
     owner_key_type = owner.__schema__(:type, owner_key)
 
+    # TODO fix the hd(values)
     # We only need to join in the "join table". Preload and Ecto.assoc expressions can then filter
     # by &1.join_owner_key in ^... to filter down to the associated entries in the related table.
     from(q in (query || queryable),
       join: j in ^join_through, on: field(q, ^related_key) == field(j, ^join_related_key),
-      where: field(j, ^join_owner_key) in type(^values, {:in, ^owner_key_type})
+      where: field(j, ^join_owner_key) in type(^hd(values), {:in, ^owner_key_type})
     )
     |> Ecto.Association.combine_assoc_query(assoc.where)
     |> Ecto.Association.combine_joins_query(assoc.join_where, 1)