Support composite primary keys during transaction rollback

activerecord/lib/active_record/transactions.rb

@@ -416,8 +416,16 @@ def restore_transaction_record_state(force_restore_state = false)
             end
             @mutations_from_database = nil
             @mutations_before_last_save = nil
-            if @attributes.fetch_value(@primary_key) != restore_state[:id]
-              @attributes.write_from_user(@primary_key, restore_state[:id])
+            if self.class.composite_primary_key?
+              if restore_state[:id] != @primary_key.map { |col| @attributes.fetch_value(col) }
+                @primary_key.zip(restore_state[:id]).each do |col, val|
+                  @attributes.write_from_user(col, val)
+                end
+              end
+            else
+              if @attributes.fetch_value(@primary_key) != restore_state[:id]
+                @attributes.write_from_user(@primary_key, restore_state[:id])
+              end
             end
             freeze if restore_state[:frozen?]
           end

activerecord/test/cases/transactions_test.rb

@@ -9,6 +9,7 @@
 require "models/author"
 require "models/post"
 require "models/movie"
+require "models/cpk"
 
 class TransactionTest < ActiveRecord::TestCase
   self.use_transactional_tests = false
@@ -926,6 +927,32 @@ def test_restore_previously_new_record_after_double_save
     assert_predicate topic, :previously_new_record?
   end
 
+  def test_restore_composite_id_after_rollback
+    book = Cpk::Book.create!(author_id: 1, number: 2)
+
+    Cpk::Book.transaction do
+      book.update!(author_id: 42, number: 42)
+      raise ActiveRecord::Rollback
+    end
+
+    assert_equal [1, 2], book.id
+  ensure
+    Cpk::Book.delete_all
+  end
+
+  def test_rollback_on_composite_key_model
+    Cpk::Book.create!(author_id: 1, number: 3, title: "Charlotte's Web")
+    book_two_unpersisted = Cpk::Book.new(author_id: 1, number: 3)
+
+    assert_raise(ActiveRecord::RecordNotUnique) do
+      Cpk::Book.transaction do
+        book_two_unpersisted.save!
+      end
+    end
+  ensure
+    Cpk::Book.delete_all
+  end
+
   def test_restore_id_after_rollback
     topic = Topic.new