Make references statements reversible

2 files changed, 40 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/migration/command_recorder.rb b/activerecord/lib/active_record/migration/command_recorder.rb
index 96b62fdd61..95f4360578 100644
--- a/activerecord/lib/active_record/migration/command_recorder.rb
+++ b/activerecord/lib/active_record/migration/command_recorder.rb
@@ -51,13 +51,15 @@ module ActiveRecord
         super || delegate.respond_to?(*args)
       end
 
-      [: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].each do |method|
+      [: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].each do |method|
         class_eval <<-EOV, __FILE__, __LINE__ + 1
           def #{method}(*args)          # def create_table(*args)
             record(:"#{method}", args)  #   record(:create_table, args)
           end                           # end
         EOV
       end
+      alias :add_belongs_to :add_reference
+      alias :remove_belongs_to :remove_reference
 
       private
 
@@ -102,6 +104,16 @@ module ActiveRecord
         [:remove_timestamps, args]
       end
 
+      def invert_add_reference(args)
+        [:remove_reference, args]
+      end
+      alias :invert_add_belongs_to :invert_add_reference
+
+      def invert_remove_reference(args)
+        [:add_reference, args]
+      end
+      alias :invert_remove_belongs_to :invert_remove_reference
+
       # Forwards any missing method call to the \target.
       def method_missing(method, *args, &block)
         @delegate.send(method, *args, &block)
diff --git a/activerecord/test/cases/migration/command_recorder_test.rb b/activerecord/test/cases/migration/command_recorder_test.rb
index 7d026961be..f2213ee6aa 100644
--- a/activerecord/test/cases/migration/command_recorder_test.rb
+++ b/activerecord/test/cases/migration/command_recorder_test.rb
@@ -110,9 +110,9 @@ module ActiveRecord
       end
 
       def test_invert_add_index_with_name
-          @recorder.record :add_index, [:table, [:one, :two], {:name => "new_index"}]
-          remove = @recorder.inverse.first
-          assert_equal [:remove_index, [:table, {:name => "new_index"}]], remove
+        @recorder.record :add_index, [:table, [:one, :two], {:name => "new_index"}]
+        remove = @recorder.inverse.first
+        assert_equal [:remove_index, [:table, {:name => "new_index"}]], remove
       end
 
       def test_invert_add_index_with_no_options
@@ -138,6 +138,30 @@ module ActiveRecord
         add = @recorder.inverse.first
         assert_equal [:add_timestamps, [:table]], add
       end
+
+      def test_invert_add_reference
+        @recorder.record :add_reference, [:table, :taggable, { polymorphic: true }]
+        remove = @recorder.inverse.first
+        assert_equal [:remove_reference, [:table, :taggable, { polymorphic: true }]], remove
+      end
+
+      def test_invert_add_belongs_to_alias
+        @recorder.record :add_belongs_to, [:table, :user]
+        remove = @recorder.inverse.first
+        assert_equal [:remove_reference, [:table, :user]], remove
+      end
+
+      def test_invert_remove_reference
+        @recorder.record :remove_reference, [:table, :taggable, { polymorphic: true }]
+        add = @recorder.inverse.first
+        assert_equal [:add_reference, [:table, :taggable, { polymorphic: true }]], add
+      end
+
+      def test_invert_remove_belongs_to_alias
+        @recorder.record :remove_belongs_to, [:table, :user]
+        add = @recorder.inverse.first
+        assert_equal [:add_reference, [:table, :user]], add
+      end
     end
   end
 end