From d8ede0c150998ea70194ebb1e099f6b3f817d860 Mon Sep 17 00:00:00 2001
From: Greg Rychlewski <greg.rychlewski@gmail.com>
Date: Mon, 16 Oct 2023 11:46:48 -0400
Subject: [PATCH] Reset belongs_to association if foreign key update results in
 a mismatch (#4299)

---
 lib/ecto.ex                        |  2 ++
 lib/ecto/repo/schema.ex            | 46 +++++++++++++++++++++---------
 test/ecto/repo/belongs_to_test.exs | 15 ++++++++++
 3 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/lib/ecto.ex b/lib/ecto.ex
index ae45c9753b..54b902ad3a 100644
--- a/lib/ecto.ex
+++ b/lib/ecto.ex
@@ -592,6 +592,8 @@ defmodule Ecto do
 
   """
   @spec reset_fields(Ecto.Schema.t(), list()) :: Ecto.Schema.t()
+  def reset_fields(struct, []), do: struct
+
   def reset_fields(%{__struct__: schema} = struct, fields) do
     default_struct = schema.__struct__()
     default_fields = Map.take(default_struct, fields)
diff --git a/lib/ecto/repo/schema.ex b/lib/ecto/repo/schema.ex
index c03b07e1fd..70fdee1b4a 100644
--- a/lib/ecto/repo/schema.ex
+++ b/lib/ecto/repo/schema.ex
@@ -347,8 +347,8 @@ defmodule Ecto.Repo.Schema do
       assoc_opts = assoc_opts(assocs, opts)
       user_changeset = run_prepare(changeset, prepare)
 
-      {changeset, parents, children} = pop_assocs(user_changeset, assocs)
-      changeset = process_parents(changeset, user_changeset, parents, adapter, assoc_opts)
+      {changeset, parents, children, _} = pop_assocs(user_changeset, assocs)
+      changeset = process_parents(changeset, user_changeset, parents, [], adapter, assoc_opts)
 
       if changeset.valid? do
         embeds = Ecto.Embedded.prepare(changeset, embeds, adapter, :insert)
@@ -439,8 +439,8 @@ defmodule Ecto.Repo.Schema do
         assoc_opts = assoc_opts(assocs, opts)
         user_changeset = run_prepare(changeset, prepare)
 
-        {changeset, parents, children} = pop_assocs(user_changeset, assocs)
-        changeset = process_parents(changeset, user_changeset, parents, adapter, assoc_opts)
+        {changeset, parents, children, reset_parents} = pop_assocs(user_changeset, assocs)
+        changeset = process_parents(changeset, user_changeset, parents, reset_parents, adapter, assoc_opts)
 
         if changeset.valid? do
           embeds = Ecto.Embedded.prepare(changeset, embeds, adapter, :update)
@@ -865,29 +865,46 @@ defmodule Ecto.Repo.Schema do
   end
 
   defp pop_assocs(changeset, []) do
-    {changeset, [], []}
+    {changeset, [], [], []}
   end
 
-  defp pop_assocs(%{changes: changes, types: types} = changeset, assocs) do
-    {changes, parent, child} =
-      Enum.reduce assocs, {changes, [], []}, fn assoc, {changes, parent, child} ->
+  defp pop_assocs(%{changes: changes, types: types, data: data} = changeset, assocs) do
+    {changes, parent, child, reset} =
+      Enum.reduce(assocs, {changes, [], [], []}, fn assoc, {changes, parent, child, reset} ->
         case changes do
           %{^assoc => value} ->
             changes = Map.delete(changes, assoc)
 
             case types do
               %{^assoc => {:assoc, %{relationship: :parent} = refl}} ->
-                {changes, [{refl, value} | parent], child}
+                {changes, [{refl, value} | parent], child, reset}
+
               %{^assoc => {:assoc, %{relationship: :child} = refl}} ->
-                {changes, parent, [{refl, value} | child]}
+                {changes, parent, [{refl, value} | child], reset}
             end
 
           %{} ->
-            {changes, parent, child}
+            with %{^assoc => {:assoc, %{relationship: :parent} = refl}} <- types,
+                 true <- reset_parent?(changes, data, refl) do
+              {changes, parent, child, [assoc | reset]}
+            else
+              _ -> {changes, parent, child, reset}
+            end
         end
-      end
+      end)
 
-    {%{changeset | changes: changes}, parent, child}
+    {%{changeset | changes: changes}, parent, child, reset}
+  end
+
+  defp reset_parent?(changes, data, assoc) do
+    %{field: field, owner_key: owner_key, related_key: related_key} = assoc
+
+    with %{^owner_key => owner_value} <- changes,
+         %{^field => %{^related_key => related_value}} when owner_value != related_value <- data do
+      true
+    else
+      _ -> false
+    end
   end
 
   # Don't mind computing options if there are no assocs
@@ -897,7 +914,7 @@ defmodule Ecto.Repo.Schema do
     Keyword.take(opts, [:timeout, :log, :telemetry_event, :prefix])
   end
 
-  defp process_parents(changeset, user_changeset, assocs, adapter, opts) do
+  defp process_parents(changeset, user_changeset, assocs, reset_assocs, adapter, opts) do
     %{changes: changes, valid?: valid?} = changeset
 
     # Even if the changeset is invalid, we want to run parent callbacks
@@ -905,6 +922,7 @@ defmodule Ecto.Repo.Schema do
     case Ecto.Association.on_repo_change(changeset, assocs, adapter, opts) do
       {:ok, struct} when valid? ->
         changes = change_parents(changes, struct, assocs)
+        struct = Ecto.reset_fields(struct, reset_assocs)
         %{changeset | changes: changes, data: struct}
 
       {:ok, _} ->
diff --git a/test/ecto/repo/belongs_to_test.exs b/test/ecto/repo/belongs_to_test.exs
index fceaac9e77..b9d965130d 100644
--- a/test/ecto/repo/belongs_to_test.exs
+++ b/test/ecto/repo/belongs_to_test.exs
@@ -481,4 +481,19 @@ defmodule Ecto.Repo.BelongsToTest do
     loaded = put_in schema.__meta__.state, :loaded
     TestRepo.insert!(loaded)
   end
+
+  test "reset assoc when foreign key update results in a mismatch" do
+    schema = TestRepo.insert!(%MySchema{assoc_id: 1, assoc: %MyAssoc{id: 1}})
+    assert schema.assoc_id == 1
+    assert %MyAssoc{id: 1} = schema.assoc
+
+    updated_schema =
+      schema
+      |> Ecto.Changeset.change(%{assoc_id: 2})
+      |> TestRepo.update!()
+
+
+    assert updated_schema.assoc_id == 2;
+    assert %Ecto.Association.NotLoaded{} = updated_schema.assoc
+  end
 end