Fix `transaction` reverting for migrations

[fatkodima & David Verhasselt]

2 files changed, 33 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/migration/command_recorder.rb b/activerecord/lib/active_record/migration/command_recorder.rb
index dea6d4ec08..82f5121d94 100644
--- a/activerecord/lib/active_record/migration/command_recorder.rb
+++ b/activerecord/lib/active_record/migration/command_recorder.rb
@@ -108,11 +108,17 @@ module ActiveRecord
         yield delegate.update_table_definition(table_name, self)
       end
 
+      def replay(migration)
+        commands.each do |cmd, args, block|
+          migration.send(cmd, *args, &block)
+        end
+      end
+
       private
 
         module StraightReversions # :nodoc:
           private
-            { transaction:       :transaction,
+            {
               execute_block:     :execute_block,
               create_table:      :drop_table,
               create_join_table: :drop_join_table,
@@ -133,6 +139,17 @@ module ActiveRecord
 
         include StraightReversions
 
+        def invert_transaction(args)
+          sub_recorder = CommandRecorder.new(delegate)
+          sub_recorder.revert { yield }
+
+          invertions_proc = proc {
+            sub_recorder.replay(self)
+          }
+
+          [:transaction, args, invertions_proc]
+        end
+
         def invert_drop_table(args, &block)
           if args.size == 1 && block == nil
             raise ActiveRecord::IrreversibleMigration, "To avoid mistakes, drop_table is only reversible if given options or a block (can be empty)."
diff --git a/activerecord/lib/active_record/migration/compatibility.rb b/activerecord/lib/active_record/migration/compatibility.rb
index 0edaaa0cf9..8f6fcfcaea 100644
--- a/activerecord/lib/active_record/migration/compatibility.rb
+++ b/activerecord/lib/active_record/migration/compatibility.rb
@@ -16,6 +16,21 @@ module ActiveRecord
       V6_0 = Current
 
       class V5_2 < V6_0
+        module CommandRecorder
+          def invert_transaction(args, &block)
+            [:transaction, args, block]
+          end
+        end
+
+        private
+
+          def command_recorder
+            recorder = super
+            class << recorder
+              prepend CommandRecorder
+            end
+            recorder
+          end
       end
 
       class V5_1 < V5_2