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.

1 files changed, 20 insertions, 0 deletions

diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 97b1ff18e2..fbcb21118d 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -21,11 +21,13 @@ module ActiveRecord::Associations::Builder
 
     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 @@ module ActiveRecord::Associations::Builder
             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)