From 8780874dbc7420a573f184a99c5fcf16dba43e86 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Mon, 17 Feb 2025 00:40:28 -0500
Subject: [PATCH] fix: ensure atomic `set_attribute` behaves the same as
 non-atomic

---
 lib/ash/changeset/changeset.ex           | 15 ++++++++-
 lib/ash/resource/change/set_attribute.ex | 43 +++++++++++++++++++++---
 test/actions/update_test.exs             |  2 +-
 3 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/lib/ash/changeset/changeset.ex b/lib/ash/changeset/changeset.ex
index 3c836da07..6c7c39220 100644
--- a/lib/ash/changeset/changeset.ex
+++ b/lib/ash/changeset/changeset.ex
@@ -953,7 +953,8 @@ defmodule Ash.Changeset do
         changeset.context
       )
 
-    with {:atomic, changeset, atomic_changes, validations} <-
+    with {:ok, change_opts} <- module.init(change_opts),
+         {:atomic, changeset, atomic_changes, validations} <-
            atomic_with_changeset(
              module.atomic(changeset, change_opts, struct(Ash.Resource.Change.Context, context)),
              changeset
@@ -1805,6 +1806,18 @@ defmodule Ash.Changeset do
       end
     end)
     |> Ash.Changeset.hydrate_atomic_refs(actor, eager?: true)
+    |> then(fn changeset ->
+      if changeset.action.type == :update do
+        attributes =
+          changeset.attributes
+          |> Map.keys()
+          |> Enum.reject(&Ash.Resource.Info.attribute(changeset.resource, &1).allow_nil?)
+
+        require_values(changeset, :update, false, attributes)
+      else
+        changeset
+      end
+    end)
   end
 
   @doc """
diff --git a/lib/ash/resource/change/set_attribute.ex b/lib/ash/resource/change/set_attribute.ex
index 21f1065d7..8d543af84 100644
--- a/lib/ash/resource/change/set_attribute.ex
+++ b/lib/ash/resource/change/set_attribute.ex
@@ -71,8 +71,7 @@ defmodule Ash.Resource.Change.SetAttribute do
         Changeset.force_change_attribute(changeset, opts[:attribute], value)
       end
     else
-      if opts[:set_when_nil?] or
-           Changeset.get_attribute(changeset, opts[:attribute]) != nil do
+      if value != nil or opts[:set_when_nil?] do
         Changeset.force_change_attribute(changeset, opts[:attribute], value)
       else
         changeset
@@ -81,13 +80,49 @@ defmodule Ash.Resource.Change.SetAttribute do
   end
 
   @impl true
-  def atomic(_changeset, opts, _context) do
+  def atomic(changeset, opts, context) do
     value =
       case opts[:value] do
         value when is_function(value) -> value.()
         value -> value
       end
 
-    {:atomic, %{opts[:attribute] => value}}
+    if Ash.Expr.expr?(value) do
+      if opts[:new?] do
+        {:atomic,
+         %{
+           opts[:attribute] =>
+             expr(
+               if is_nil(^atomic_ref(opts[:attribute])) do
+                 ^value
+               else
+                 ^atomic_ref(opts[:attribute])
+               end
+             )
+         }}
+      else
+        if opts[:set_when_nil?] do
+          {:atomic, %{opts[:attribute] => value}}
+        else
+          if is_nil(value) do
+            {:ok, changeset}
+          else
+            {:atomic,
+             %{
+               opts[:attribute] =>
+                 expr(
+                   if is_nil(^value) do
+                     ^atomic_ref(opts[:attribute])
+                   else
+                     ^value
+                   end
+                 )
+             }}
+          end
+        end
+      end
+    else
+      {:ok, change(changeset, opts, context)}
+    end
   end
 end
diff --git a/test/actions/update_test.exs b/test/actions/update_test.exs
index 54b3cf1cd..9423dce81 100644
--- a/test/actions/update_test.exs
+++ b/test/actions/update_test.exs
@@ -81,7 +81,7 @@ defmodule Ash.Test.Actions.UpdateTest do
 
       update :set_private_attribute_to_nil do
         accept []
-        change set_attribute(:non_nil_private, nil)
+        change set_attribute(:non_nil_private, nil, set_when_nil?: true)
       end
 
       update :set_private_attribute_from_arg do