fix: return skipped upserts in bulk_create (#626)

Author: barnabasJ

lib/data_layer.ex

@@ -662,6 +662,7 @@ defmodule AshPostgres.DataLayer do
   def can?(_, :combine), do: true
   def can?(_, {:combine, _}), do: true
   def can?(_, :bulk_create), do: true
+  def can?(_, :bulk_upsert_return_skipped), do: true
 
   def can?(_, :action_select), do: true
 
@@ -2035,6 +2036,74 @@ defmodule AshPostgres.DataLayer do
           repo.insert_all(source, ecto_changesets, opts)
         end)
 
+      identity = options[:identity]
+      keys = Map.get(identity || %{}, :keys) || Ash.Resource.Info.primary_key(resource)
+
+      # if it's single the return_skipped_upsert? is handled at the
+      # call site https://github.com/ash-project/ash_postgres/blob/0b21d4a99cc3f6d8676947e291ac9b9d57ad6e2e/lib/data_layer.ex#L3046-L3046
+      result =
+        if options[:return_skipped_upsert?] && !opts[:single?] do
+          [changeset | _] = changesets
+
+          results_by_identity =
+            result
+            |> elem(1)
+            |> List.wrap()
+            |> Enum.into(%{}, fn r ->
+              {Map.take(r, keys), r}
+            end)
+
+          ash_query =
+            resource
+            |> Ash.Query.do_filter(
+              or:
+                changesets
+                |> Enum.filter(fn changeset ->
+                  not Map.has_key?(
+                    results_by_identity,
+                    Map.take(changeset.attributes, keys)
+                  )
+                end)
+                |> Enum.map(fn changeset ->
+                  changeset.attributes
+                  |> Map.take(keys)
+                  |> Keyword.new()
+                end)
+            )
+            |> then(fn
+              query when is_nil(identity) or is_nil(identity.where) -> query
+              query -> Ash.Query.do_filter(query, identity.where)
+            end)
+            |> Ash.Query.set_tenant(changeset.tenant)
+
+          skipped_upserts =
+            with {:ok, ecto_query} <- Ash.Query.data_layer_query(ash_query),
+                 {:ok, results} <- run_query(ecto_query, resource) do
+              results
+              |> Enum.map(fn result ->
+                Ash.Resource.put_metadata(result, :upsert_skipped, true)
+              end)
+              |> Enum.reduce(%{}, fn r, acc ->
+                Map.put(acc, Map.take(r, keys), r)
+              end)
+            end
+
+          results =
+            changesets
+            |> Enum.map(fn changeset ->
+              identity =
+                changeset.attributes
+                |> Map.take(keys)
+
+              Map.get(results_by_identity, identity, Map.get(skipped_upserts, identity))
+            end)
+            |> Enum.filter(& &1)
+
+          {length(results), results}
+        else
+          result
+        end
+
       case result do
         {_, nil} ->
           :ok
@@ -2045,25 +2114,43 @@ defmodule AshPostgres.DataLayer do
 
             {:ok, results}
           else
-            {:ok,
-             Stream.zip_with(results, changesets, fn result, changeset ->
-               if !opts[:upsert?] do
-                 maybe_create_tenant!(resource, result)
-               end
-
-               case get_bulk_operation_metadata(changeset) do
-                 {index, metadata_key} ->
-                   Ash.Resource.put_metadata(result, metadata_key, index)
-
-                 nil ->
-                   # Compatibility fallback
-                   Ash.Resource.put_metadata(
-                     result,
-                     :bulk_create_index,
-                     changeset.context[:bulk_create][:index]
-                   )
-               end
-             end)}
+            results_by_identity =
+              results
+              |> Enum.into(%{}, fn r ->
+                {Map.take(r, keys), r}
+              end)
+
+            results =
+              changesets
+              |> Enum.map(fn changeset ->
+                identity =
+                  changeset.attributes
+                  |> Map.take(keys)
+
+                result_for_changeset = Map.get(results_by_identity, identity)
+
+                if result_for_changeset do
+                  if !opts[:upsert?] do
+                    maybe_create_tenant!(resource, result_for_changeset)
+                  end
+
+                  case get_bulk_operation_metadata(changeset) do
+                    {index, metadata_key} ->
+                      Ash.Resource.put_metadata(result_for_changeset, metadata_key, index)
+
+                    nil ->
+                      # Compatibility fallback
+                      Ash.Resource.put_metadata(
+                        result_for_changeset,
+                        :bulk_create_index,
+                        changeset.context[:bulk_create][:index]
+                      )
+                  end
+                end
+              end)
+              |> Enum.filter(& &1)
+
+            {:ok, results}
           end
       end
     rescue

mix.exs

@@ -177,10 +177,10 @@ defmodule AshPostgres.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:ash, ash_version("~> 3.5 and >= 3.5.35")},
+      {:ash, ash_version("~> 3.5 and >= 3.6.2")},
       {:spark, "~> 2.3 and >= 2.3.4"},
       {:ash_sql, ash_sql_version("~> 0.3 and >= 0.3.2")},
-      {:igniter, "~> 0.6 and >= 0.6.14", optional: true},
+      {:igniter, "~> 0.6 and >= 0.6.29", optional: true},
       {:ecto_sql, "~> 3.13"},
       {:ecto, "~> 3.13"},
       {:jason, "~> 1.0"},

test/bulk_create_test.exs

@@ -176,6 +176,67 @@ defmodule AshPostgres.BulkCreateTest do
                end)
     end
 
+    test "bulk upsert returns skipped records with return_skipped_upsert?" do
+      assert [
+               {:ok, %{title: "fredfoo", uniq_if_contains_foo: "1foo", price: 10}},
+               {:ok, %{title: "georgefoo", uniq_if_contains_foo: "2foo", price: 20}},
+               {:ok, %{title: "herbert", uniq_if_contains_foo: "3", price: 30}}
+             ] =
+               Ash.bulk_create!(
+                 [
+                   %{title: "fredfoo", uniq_if_contains_foo: "1foo", price: 10},
+                   %{title: "georgefoo", uniq_if_contains_foo: "2foo", price: 20},
+                   %{title: "herbert", uniq_if_contains_foo: "3", price: 30}
+                 ],
+                 Post,
+                 :create,
+                 return_stream?: true,
+                 return_records?: true
+               )
+               |> Enum.sort_by(fn {:ok, result} -> result.title end)
+
+      results =
+        Ash.bulk_create!(
+          [
+            %{title: "fredfoo", uniq_if_contains_foo: "1foo", price: 10},
+            %{title: "georgefoo", uniq_if_contains_foo: "2foo", price: 20_000},
+            %{title: "herbert", uniq_if_contains_foo: "3", price: 30}
+          ],
+          Post,
+          :upsert_with_no_filter,
+          return_stream?: true,
+          upsert_condition: expr(price != upsert_conflict(:price)),
+          return_errors?: true,
+          return_records?: true,
+          return_skipped_upsert?: true
+        )
+        |> Enum.sort_by(fn
+          {:ok, result} ->
+            result.title
+
+          _ ->
+            nil
+        end)
+
+      assert [
+               {:ok, skipped},
+               {:ok, updated},
+               {:ok, no_conflict}
+             ] = results
+
+      assert skipped.title == "fredfoo"
+      assert skipped.price == 10
+      assert Ash.Resource.get_metadata(skipped, :upsert_skipped) == true
+
+      assert updated.title == "georgefoo"
+      assert updated.price == 20_000
+      refute Ash.Resource.get_metadata(updated, :upsert_skipped)
+
+      assert no_conflict.title == "herbert"
+      assert no_conflict.price == 30
+      refute Ash.Resource.get_metadata(no_conflict, :upsert_skipped)
+    end
+
     # confirmed that this doesn't work because it can't. An upsert must map to a potentially successful insert.
     # leaving this test here for posterity
     # test "bulk creates can upsert with id" do