Merge pull request #9737 from wangjohn/counter_cache_update_attributes_fix

The counter cache will now work correctly when the foreign key is changed. Fixes #9722.

activerecord/CHANGELOG.md

 ## Rails 4.0.0 (unreleased) ##
 
+*   Counter caches on associations will now stay valid when attributes are
+    updated (not just when records are created or destroyed), for example,
+    when calling +update_attributes+. The following code now works:
+
+      class Comment < ActiveRecord::Base
+        belongs_to :post, counter_cache: true
+      end
+
+      class Post < ActiveRecord::Base
+        has_many :comments
+      end
+
+      post = Post.create
+      comment = Comment.create
+
+      post.comments << comment
+      post.save.reload.comments_count     # => 1
+      comment.update_attributes(:post_id => nil)
+
+      post.save.reload.comments_count     # => 0
+
+    Updating the id of a +belongs_to+ object with the id of a new object will
+    also keep the count accurate.
+
+    *John Wang*
+
 *   Referencing join tables implicitly was deprecated. There is a
     possibility that these deprecation warnings are shown even if you
     don't make use of that feature. You can now disable the feature entirely.

activerecord/lib/active_record/associations/builder/belongs_to.rb

@@ -21,11 +21,13 @@ def build
 
     def add_counter_cache_callbacks(reflection)
       cache_column = reflection.counter_cache_column
+      foreign_key = reflection.foreign_key
 
       mixin.class_eval <<-CODE, __FILE__, __LINE__ + 1
         def belongs_to_counter_cache_after_create_for_#{name}
           record = #{name}
           record.class.increment_counter(:#{cache_column}, record.id) unless record.nil?
+          @_after_create_counter_called = true
         end
 
         def belongs_to_counter_cache_before_destroy_for_#{name}
@@ -34,10 +36,28 @@ def belongs_to_counter_cache_before_destroy_for_#{name}
             record.class.decrement_counter(:#{cache_column}, record.id) unless record.nil?
           end
         end
+
+        def belongs_to_counter_cache_after_update_for_#{name}
+          if (@_after_create_counter_called ||= false)
+            @_after_create_counter_called = false
+          elsif self.#{foreign_key}_changed? && !new_record? && defined?(#{name.to_s.camelize})
+            model = #{name.to_s.camelize}
+            foreign_key_was = self.#{foreign_key}_was
+            foreign_key = self.#{foreign_key}
+
+            if foreign_key && model.respond_to?(:increment_counter)
+              model.increment_counter(:#{cache_column}, foreign_key)
+            end
+            if foreign_key_was && model.respond_to?(:decrement_counter)
+              model.decrement_counter(:#{cache_column}, foreign_key_was)
+            end
+          end
+        end
       CODE
 
       model.after_create   "belongs_to_counter_cache_after_create_for_#{name}"
       model.before_destroy "belongs_to_counter_cache_before_destroy_for_#{name}"
+      model.after_update   "belongs_to_counter_cache_after_update_for_#{name}"
 
       klass = reflection.class_name.safe_constantize
       klass.attr_readonly cache_column if klass && klass.respond_to?(:attr_readonly)

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -789,6 +789,37 @@ def test_custom_named_counter_cache
     end
   end
 
+  def test_calling_update_attributes_on_id_changes_the_counter_cache
+    topic = Topic.order("id ASC").first
+    original_count = topic.replies.to_a.size
+    assert_equal original_count, topic.replies_count
+
+    first_reply = topic.replies.first
+    first_reply.update_attributes(:parent_id => nil)
+    assert_equal original_count - 1, topic.reload.replies_count
+
+    first_reply.update_attributes(:parent_id => topic.id)
+    assert_equal original_count, topic.reload.replies_count
+  end
+
+  def test_calling_update_attributes_changing_ids_doesnt_change_counter_cache
+    topic1 = Topic.find(1)
+    topic2 = Topic.find(3)
+    original_count1 = topic1.replies.to_a.size
+    original_count2 = topic2.replies.to_a.size
+
+    reply1 = topic1.replies.first
+    reply2 = topic2.replies.first
+
+    reply1.update_attributes(:parent_id => topic2.id)
+    assert_equal original_count1 - 1, topic1.reload.replies_count
+    assert_equal original_count2 + 1, topic2.reload.replies_count
+
+    reply2.update_attributes(:parent_id => topic1.id)
+    assert_equal original_count1, topic1.reload.replies_count
+    assert_equal original_count2, topic2.reload.replies_count
+  end
+
   def test_deleting_a_collection
     force_signal37_to_load_all_clients_of_firm
     companies(:first_firm).clients_of_firm.create("name" => "Another Client")