1 files changed, 21 insertions, 1 deletions
diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 6b479968c8..579d7789bd 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -21,23 +21,43 @@ 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}
-          unless marked_for_destruction?
+          unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == #{foreign_key.to_sym.inspect}
             record = #{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)