fix: correctly handle return values of batch callbacks (#1424)

* improvement: factor out static branches at compile time

---------

Co-authored-by: Hannes Wüthrich <hannes.wuethrich@zebbra.ch>

lib/ash/resource/change/change.ex

@@ -209,50 +209,70 @@ defmodule Ash.Resource.Change do
         if Module.defines?(__MODULE__, {:batch_change, 3}, :def) do
           @impl true
           def change(changeset, opts, context) do
-            if has_before_batch?() do
+            changeset
+            |> simulate_before_batch(opts, context)
+            |> Ash.Changeset.before_action(fn changeset ->
+              Enum.at(batch_change([changeset], opts, context), 0)
+            end)
+            |> simulate_after_batch(opts, context)
+          end
+
+          if Module.defines?(__MODULE__, {:before_batch, 3}, :def) do
+            defp simulate_before_batch(changeset, opts, context) do
               Ash.Changeset.before_action(changeset, fn changeset ->
                 {[changeset], notifications} =
                   Enum.split_with(
                     apply(__MODULE__, :before_batch, [[changeset], opts, context]),
-                    fn %Ash.Notifier.Notification{} ->
-                      false
+                    fn
+                      %Ash.Notifier.Notification{} ->
+                        false
+
+                      %Ash.Changeset{} ->
+                        true
+
+                      other ->
+                        raise "Expected before_batch/3 to return a list of changesets and notifications, got: #{inspect(other)}"
                     end
                   )
 
                 {changeset, %{notifications: notifications}}
               end)
-            else
+            end
+          else
+            defp simulate_before_batch(changeset, _opts, _context) do
               changeset
             end
-            |> Ash.Changeset.before_action(fn changeset ->
-              Enum.at(batch_change([changeset], opts, context), 0)
-            end)
-            |> then(fn changeset ->
-              if has_after_batch?() do
-                Ash.Changeset.after_action(changeset, fn changeset, result ->
-                  apply(__MODULE__, :after_batch, [[{changeset, result}], opts, context])
-                  |> Enum.reduce({[], [], []}, fn item, {records, errors, notifications} ->
-                    case item do
-                      {:ok, record} -> {[record | records], errors, notifications}
-                      {:error, error} -> {records, [error | errors], notifications}
-                      %Ash.Notifier.Notification{} -> {records, errors, [item | notifications]}
-                    end
-                  end)
-                  |> case do
-                    {[record], [], notifications} ->
-                      {:ok, record, notifications}
+          end
 
-                    {other, [], _notifications} ->
-                      raise "Invalid return value from `after_batch/3`. Expected exactly one record: #{inspect(other)}"
-
-                    {_, errors, _notifications} ->
-                      {:error, errors}
+          if Module.defines?(__MODULE__, {:after_batch, 3}, :def) do
+            defp simulate_after_batch(changeset, opts, context) do
+              Ash.Changeset.after_action(changeset, fn changeset, result ->
+                apply(__MODULE__, :after_batch, [[{changeset, result}], opts, context])
+                |> then(fn
+                  :ok -> [{:ok, result}]
+                  other -> other
+                end)
+                |> Enum.reduce({[], [], []}, fn item, {records, errors, notifications} ->
+                  case item do
+                    {:ok, record} -> {[record | records], errors, notifications}
+                    {:error, error} -> {records, [error | errors], notifications}
+                    %Ash.Notifier.Notification{} -> {records, errors, [item | notifications]}
                   end
                 end)
-              else
-                changeset
-              end
-            end)
+                |> case do
+                  {[record], [], notifications} ->
+                    {:ok, record, notifications}
+
+                  {other, [], _notifications} ->
+                    raise "Invalid return value from `after_batch/3`. Expected exactly one record: #{inspect(other)}"
+
+                  {_, errors, _notifications} ->
+                    {:error, errors}
+                end
+              end)
+            end
+          else
+            defp simulate_after_batch(changeset, _opts, _context), do: changeset
           end
 
           @impl true

lib/ash/type/enum.ex

@@ -86,9 +86,9 @@ defmodule Ash.Type.Enum do
   ```
   """
   @doc "The list of valid values (not all input types that match them)"
-  @callback values() :: [atom]
+  @callback values() :: [atom | String.t()]
   @doc "The description of the value, if existing"
-  @callback description(atom) :: String.t() | nil
+  @callback description(atom | String.t()) :: String.t() | nil
   @doc "true if a given term matches a value"
   @callback match?(term) :: boolean
   @doc "finds the valid value that matches a given input term"

test/actions/bulk/bulk_create_test.exs

@@ -18,10 +18,17 @@ defmodule Ash.Test.Actions.BulkCreateTest do
   defmodule AddAfterToTitle do
     use Ash.Resource.Change
 
+    @impl Ash.Resource.Change
     def batch_change(changesets, _, _) do
       changesets
     end
 
+    @impl Ash.Resource.Change
+    def before_batch(changesets, _, _) do
+      changesets
+    end
+
+    @impl Ash.Resource.Change
     def after_batch(results, _, _) do
       Stream.map(results, fn {_changeset, result} ->
         {:ok, %{result | title: result.title <> "_after"}}
@@ -32,14 +39,12 @@ defmodule Ash.Test.Actions.BulkCreateTest do
   defmodule AddBeforeToTitle do
     use Ash.Resource.Change
 
-    def change(changeset, _, %{bulk?: true}) do
-      changeset
-    end
-
+    @impl Ash.Resource.Change
     def batch_change(changesets, _, _) do
       changesets
     end
 
+    @impl Ash.Resource.Change
     def before_batch(changesets, _, _) do
       changesets
       |> Stream.map(fn changeset ->
@@ -47,6 +52,11 @@ defmodule Ash.Test.Actions.BulkCreateTest do
         Ash.Changeset.force_change_attribute(changeset, :title, "before_" <> title)
       end)
     end
+
+    @impl Ash.Resource.Change
+    def after_batch(_, _, _) do
+      :ok
+    end
   end
 
   defmodule ChangeMessage do
@@ -505,6 +515,15 @@ defmodule Ash.Test.Actions.BulkCreateTest do
                sorted?: true,
                authorize?: false
              )
+
+    assert %{title: "before_title_after"} =
+             Ash.create!(
+               Post,
+               %{title: "title"},
+               action: :create_with_after_batch,
+               tenant: org.id,
+               authorize?: false
+             )
   end
 
   test "will return error count" do