Add Migration#reversible for reversible data operations [#8267]

3 files changed, 85 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index c605c2da7f..806b367c6b 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -439,6 +439,47 @@ module ActiveRecord
       @connection.respond_to?(:reverting) && @connection.reverting
     end
 
+    class ReversibleBlockHelper < Struct.new(:reverting)
+      def up
+        yield unless reverting
+      end
+
+      def down
+        yield if reverting
+      end
+    end
+
+    # Used to specify an operation that can be run in one direction or another.
+    # Call the methods +up+ and +down+ of the yielded object to run a block
+    # only in one given direction.
+    # The whole block will be called in the right order within the migration.
+    #
+    # In the following example, the looping on users will always be done
+    # when the three columns 'first_name', 'last_name' and 'full_name' exist,
+    # even when migrating down:
+    #
+    #    class SplitNameMigration < ActiveRecord::Migration
+    #      def change
+    #        add_column :users, :first_name, :string
+    #        add_column :users, :last_name, :string
+    #
+    #        reversible do |dir|
+    #          User.reset_column_information
+    #          User.all.each do |u|
+    #            dir.up   { u.first_name, u.last_name = u.full_name.split(' ') }
+    #            dir.down { u.full_name = "#{u.first_name} #{u.last_name}" }
+    #            u.save
+    #          end
+    #        end
+    #
+    #        revert { add_column :users, :full_name, :string }
+    #      end
+    #    end
+    def reversible
+      helper = ReversibleBlockHelper.new(reverting?)
+      transaction{ yield helper }
+    end
+
     # Runs the given migration classes.
     # Last argument can specify options:
     # - :direction (default is :up)
diff --git a/activerecord/lib/active_record/migration/command_recorder.rb b/activerecord/lib/active_record/migration/command_recorder.rb
index a79428b69a..d0e54240d3 100644
--- a/activerecord/lib/active_record/migration/command_recorder.rb
+++ b/activerecord/lib/active_record/migration/command_recorder.rb
@@ -72,7 +72,7 @@ module ActiveRecord
 
       [:create_table, :create_join_table, :change_table, :rename_table, :add_column, :remove_column,
         :rename_index, :rename_column, :add_index, :remove_index, :add_timestamps, :remove_timestamps,
-        :change_column, :change_column_default, :add_reference, :remove_reference,
+        :change_column, :change_column_default, :add_reference, :remove_reference, :transaction,
       ].each do |method|
         class_eval <<-EOV, __FILE__, __LINE__ + 1
           def #{method}(*args, &block)          # def create_table(*args, &block)
@@ -85,6 +85,10 @@ module ActiveRecord
 
       private
 
+      def invert_transaction(args, &block)
+        [:transaction, args, block]
+      end
+
       def invert_create_table(args)
         [:drop_table, [args.first]]
       end
diff --git a/activerecord/test/cases/invertible_migration_test.rb b/activerecord/test/cases/invertible_migration_test.rb
index 6b81e600af..484b4c0fea 100644
--- a/activerecord/test/cases/invertible_migration_test.rb
+++ b/activerecord/test/cases/invertible_migration_test.rb
@@ -28,6 +28,26 @@ module ActiveRecord
       end
     end
 
+    class InvertibleByPartsMigration < SilentMigration
+      attr_writer :test
+      def change
+        create_table("new_horses") do |t|
+          t.column :breed, :string
+        end
+        reversible do |dir|
+          @test.yield :both
+          dir.up    { @test.yield :up }
+          dir.down  { @test.yield :down }
+        end
+        revert do
+          create_table("horses") do |t|
+            t.column :content, :text
+            t.column :remind_at, :datetime
+          end
+        end
+      end
+    end
+
     class NonInvertibleMigration < SilentMigration
       def change
         create_table("horses") do |t|
@@ -107,6 +127,25 @@ module ActiveRecord
       assert !migration.connection.table_exists?("horses")
     end
 
+    def test_migrate_revert_by_part
+      InvertibleMigration.new.migrate :up
+      received = []
+      migration = InvertibleByPartsMigration.new
+      migration.test = ->(dir){
+        assert migration.connection.table_exists?("horses")
+        assert migration.connection.table_exists?("new_horses")
+        received << dir
+      }
+      migration.migrate :up
+      assert_equal [:both, :up], received
+      assert !migration.connection.table_exists?("horses")
+      assert migration.connection.table_exists?("new_horses")
+      migration.migrate :down
+      assert_equal [:both, :up, :both, :down], received
+      assert migration.connection.table_exists?("horses")
+      assert !migration.connection.table_exists?("new_horses")
+    end
+
     def test_migrate_revert_whole_migration
       migration = InvertibleMigration.new
       [LegacyMigration, InvertibleMigration].each do |klass|