improvement: add bulk_metadata ref to changesets

Author: DGollings

lib/data_layer.ex

@@ -1560,9 +1560,9 @@ defmodule AshPostgres.DataLayer do
                       :ok
                   end
 
-                  {:ok, [changeset.data]}
+                  {:ok, [add_bulk_metadata(changeset.data, changeset)]}
                 else
-                  {:ok, repo.all(query, repo_opts)}
+                  {:ok, Enum.map(repo.all(query, repo_opts), &add_bulk_metadata(&1, changeset))}
                 end
               else
                 :ok
@@ -1612,9 +1612,10 @@ defmodule AshPostgres.DataLayer do
                   Map.merge(changeset.data, Map.take(result, modifying))
                   |> Map.update!(:aggregates, &Map.merge(&1, result.aggregates))
                   |> Map.update!(:calculations, &Map.merge(&1, result.calculations))
+                  |> add_bulk_metadata(changeset)
                   |> then(&{:ok, [&1]})
                 else
-                  {:ok, results}
+                  {:ok, Enum.map(results, &add_bulk_metadata(&1, changeset))}
                 end
               else
                 :ok
@@ -1887,7 +1888,8 @@ defmodule AshPostgres.DataLayer do
             end)
 
           if options[:return_records?] do
-            {:ok, AshSql.Query.remap_mapped_fields(results, query)}
+            results = AshSql.Query.remap_mapped_fields(results, query)
+            {:ok, Enum.map(results, &add_bulk_metadata(&1, changeset))}
           else
             :ok
           end
@@ -3749,4 +3751,33 @@ defmodule AshPostgres.DataLayer do
       resource
     end
   end
+
+  defp add_bulk_metadata(record, changeset) do
+    cond do
+      changeset.context[:bulk_update] ->
+        record
+        |> Ash.Resource.put_metadata(
+          :bulk_update_index,
+          changeset.context.bulk_update.index
+        )
+        |> Ash.Resource.put_metadata(
+          :bulk_action_ref,
+          changeset.context.bulk_update.ref
+        )
+
+      changeset.context[:bulk_destroy] ->
+        record
+        |> Ash.Resource.put_metadata(
+          :bulk_destroy_index,
+          changeset.context.bulk_destroy.index
+        )
+        |> Ash.Resource.put_metadata(
+          :bulk_action_ref,
+          changeset.context.bulk_destroy.ref
+        )
+
+      true ->
+        record
+    end
+  end
 end

test/data_layer_metadata_test.exs

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: 2019 ash_postgres contributors <https://github.com/ash-project/ash_postgres/graphs.contributors>
+#
+# SPDX-License-Identifier: MIT
+
+defmodule AshPostgres.DataLayerMetadataTest do
+  use AshPostgres.RepoCase, async: false
+  alias AshPostgres.Test.Post
+
+  require Ash.Query
+
+  test "update_query sets bulk metadata on returned records" do
+    _post1 = Ash.create!(Post, %{title: "title1"})
+    _post2 = Ash.create!(Post, %{title: "title2"})
+    _post3 = Ash.create!(Post, %{title: "title3"})
+
+    ash_query = Ash.Query.new(Post)
+    {:ok, query} = Ash.Query.data_layer_query(ash_query)
+
+    ref = make_ref()
+
+    changeset = %Ash.Changeset{
+      resource: Post,
+      action_type: :update,
+      data: %Post{},
+      attributes: %{title: "updated"},
+      atomics: [],
+      filter: nil,
+      context: %{bulk_update: %{index: 0, ref: ref}},
+      domain: AshPostgres.Test.Domain,
+      tenant: nil,
+      timeout: :infinity
+    }
+
+    {:ok, results} =
+      AshPostgres.DataLayer.update_query(query, changeset, Post, return_records?: true)
+
+    assert is_list(results)
+    assert length(results) > 0
+
+    Enum.each(results, fn result ->
+      assert is_integer(result.__metadata__.bulk_update_index)
+      assert result.__metadata__.bulk_action_ref == ref
+    end)
+  end
+
+  test "destroy_query sets bulk metadata on returned records" do
+    _post1 = Ash.create!(Post, %{title: "title1"})
+    _post2 = Ash.create!(Post, %{title: "title2"})
+    _post3 = Ash.create!(Post, %{title: "title3"})
+
+    ash_query = Ash.Query.new(Post)
+    {:ok, query} = Ash.Query.data_layer_query(ash_query)
+
+    ref = make_ref()
+
+    changeset = %Ash.Changeset{
+      resource: Post,
+      action_type: :destroy,
+      data: %Post{},
+      attributes: %{},
+      atomics: [],
+      filter: nil,
+      context: %{bulk_destroy: %{index: 0, ref: ref}},
+      domain: AshPostgres.Test.Domain,
+      tenant: nil,
+      timeout: :infinity
+    }
+
+    {:ok, results} =
+      AshPostgres.DataLayer.destroy_query(query, changeset, Post, return_records?: true)
+
+    assert is_list(results)
+    assert length(results) > 0
+
+    Enum.each(results, fn result ->
+      assert is_integer(result.__metadata__.bulk_destroy_index)
+      assert result.__metadata__.bulk_action_ref == ref
+    end)
+  end
+end