bulk_update Only mark updated records with changed applied

lib/activerecord-bulk_update.rb

@@ -7,7 +7,9 @@
 require "activerecord-bulk_update/activerecord/bulk_update"
 require "activerecord-bulk_update/activerecord/querying"
 require "activerecord-bulk_update/activerecord/relation"
+require "activerecord-bulk_update/arel/tree_manager"
 require "activerecord-bulk_update/arel/nodes/cast"
 require "activerecord-bulk_update/arel/nodes/from"
+require "activerecord-bulk_update/arel/nodes/returning"
 require "activerecord-bulk_update/arel/nodes/update_statement"
 require "activerecord-bulk_update/arel/visitors/postgresql"

lib/activerecord-bulk_update/activerecord/bulk_update.rb

@@ -19,23 +19,30 @@ def update_records
       return updates if values.none?
 
       touch_all if touch
-      execute
-      updates.each(&:changes_applied)
+
+      sql, binds = model.connection.send(:to_sql_and_binds, update_manager)
+      result = model.connection.exec_query(sql, "Bulk Update", binds).tap { model.reset }
+      updated_ids = result.rows.flatten(1)
+
+      updates.each do |update|
+        if updated_ids.include?(update.id)
+          update.changes_applied
+        else
+          update.changes.each { |attr, changes| update[attr] = changes.first }
+          update.clear_changes_information
+        end
+      end
     end
 
     def update_by_hash
       extract_values_from_hash
       return 0 if values.none?
 
       touch_all if touch
-      execute
+      model.connection.update(update_manager, "Bulk Update").tap { model.reset }
     end
 
     private
-      def execute
-        model.connection.update(update_manager, "Bulk Update").tap { model.reset }
-      end
-
       def update_manager
         arel.source.left = arel_table
 
@@ -46,6 +53,7 @@ def update_manager
         stmt.offset(arel.offset)
         stmt.order(*arel.orders)
         stmt.wheres = arel.constraints
+        stmt.returning(returning)
 
         if values.uniq.one?
           any_row = values[0]
@@ -56,7 +64,7 @@ def update_manager
         else
           filtering_attributes.each { |attr| arel.where(arel_table[attr].eq(source[attr])) }
           stmt.set(updating_attributes.map { |attr| [arel_table[attr], source["_#{attr}"]] })
-          stmt.ast.from = from
+          stmt.from(from)
         end
 
         stmt
@@ -66,6 +74,12 @@ def source
         @source ||= Arel::Table.new("source")
       end
 
+      def returning
+        return [] unless primary_key
+
+        Arel::Nodes::Returning.new([arel_table[primary_key]])
+      end
+
       def from
         Arel::Nodes::From.new(
           values_list,

lib/activerecord-bulk_update/arel/nodes/cast.rb

@@ -2,7 +2,7 @@
 
 module Arel
   module Nodes
-    # Adds a new type of Node that can be used when updating from a values_list
+    # Adds a new type of Node that can be used to cast a value to a given datatype.
     #
     # @example CAST('value' AS datatype)
     class Cast < Node

lib/activerecord-bulk_update/arel/nodes/returning.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Arel
+  module Nodes
+    # Adds a new type of Node that can be used to return specific columns after an insert or update.
+    #
+    # @example RETURNING "column1", "column2"
+    class Returning < Unary
+    end
+  end
+end

lib/activerecord-bulk_update/arel/nodes/update_statement.rb

@@ -4,7 +4,7 @@ module Arel
   module Nodes
     class UpdateStatement
       # New method to be able to assign a FROM clause to the update statement.
-      attr_accessor :from
+      attr_accessor :from, :returning
     end
   end
 end

lib/activerecord-bulk_update/arel/tree_manager.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Arel
+  class TreeManager
+    def from(*expr)
+      @ast.from = expr
+    end
+
+    def returning(*expr)
+      @ast.returning = expr
+    end
+  end
+end

lib/activerecord-bulk_update/arel/visitors/postgresql.rb

@@ -23,9 +23,15 @@ def visit_Arel_Nodes_From(o, collector)
         collector << ")"
       end
 
+      # New method to cast the Returning node to a partial sql statement.
+      def visit_Arel_Nodes_Returning(o, collector)
+        collector << "RETURNING "
+        collect_nodes_for o.expr, collector, "", ", "
+      end
+
       # MONKEYPATCH
       #
-      # In order to be able to include the optional FROM statement the existing method needs to be overriden.
+      # In order to be able to include the optionals FROM and RETURNING statements the existing method needs to be overriden.
       def visit_Arel_Nodes_UpdateStatement(o, collector)
         o = prepare_update_statement(o)
 
@@ -36,6 +42,7 @@ def visit_Arel_Nodes_UpdateStatement(o, collector)
         collect_nodes_for o.wheres, collector, " WHERE ", " AND "
         collect_nodes_for o.orders, collector, " ORDER BY "
         maybe_visit o.limit, collector
+        maybe_visit o.returning, collector # MONKEYPATCH
       end
     end
   end

test/activerecord/bulk_update_test.rb

@@ -37,10 +37,14 @@ def update_records
           assert_change(-> { FakeRecord.find_by!(name: "third").rank }, from: 3, to: 4) { update_records }
         end
 
-        it "marks all changes as persisted" do
+        it "marks all changes as handled" do
           assert_change(-> { @updates.count(&:has_changes_to_save?) }, from: 3, to: 0) { update_records }
         end
 
+        it "marks all changes as previously saved" do
+          assert_change(-> { @updates.count { |update| update.previous_changes.any? } }, from: 0, to: 3) { update_records }
+        end
+
         it "touches the updated_at" do
           assert_change(-> { fake_records(:first).reload.updated_at }) { update_records }
         end
@@ -121,6 +125,25 @@ def update_records
         end
       end
 
+      describe "when a filter in the query prevents the updates" do
+        before do
+          @updates = [fake_records(:first).tap { |record| record.rank = 9 }]
+          @model = FakeRecord.where(rank: 13)
+        end
+
+        it "resets the unpersisted attribute to the previous value" do
+          assert_change(-> { @updates.first.rank }, from: 9, to: 1) { update_records }
+        end
+
+        it "removes all pending changes" do
+          assert_change(-> { @updates.count(&:has_changes_to_save?) }, from: 1, to: 0) { update_records }
+        end
+
+        it "does not mark the changes as previously saved" do
+          refute_change(-> { @updates.count { |u| u.previous_changes.any? } }, from: 0) { update_records }
+        end
+      end
+
       describe "when given an empty ActiveRecord::Relation" do
         before { @updates = FakeRecord.none }