Refactor counter cache to extract `decrement_counters_before_last_save` on the belongs_to association

2 files changed, 31 insertions, 48 deletions

diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 544aec5e8b..3346725f2d 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -34,14 +34,28 @@ module ActiveRecord
         @updated
       end
 
-      def decrement_counters # :nodoc:
+      def decrement_counters
         update_counters(-1)
       end
 
-      def increment_counters # :nodoc:
+      def increment_counters
         update_counters(1)
       end
 
+      def decrement_counters_before_last_save
+        if reflection.polymorphic?
+          model_was = owner.attribute_before_last_save(reflection.foreign_type).try(:constantize)
+        else
+          model_was = klass
+        end
+
+        foreign_key_was = owner.attribute_before_last_save(reflection.foreign_key)
+
+        if foreign_key_was && model_was < ActiveRecord::Base
+          update_counters_via_scope(model_was, foreign_key_was, -1)
+        end
+      end
+
       def target_changed?
         owner.saved_change_to_attribute?(reflection.foreign_key)
       end
@@ -64,11 +78,16 @@ module ActiveRecord
             if target && !stale_target?
               target.increment!(reflection.counter_cache_column, by, touch: reflection.options[:touch])
             else
-              counter_cache_target.update_counters(reflection.counter_cache_column => by, touch: reflection.options[:touch])
+              update_counters_via_scope(klass, owner._read_attribute(reflection.foreign_key), by)
             end
           end
         end
 
+        def update_counters_via_scope(klass, foreign_key, by)
+          scope = klass.unscoped.where!(primary_key(klass) => foreign_key)
+          scope.update_counters(reflection.counter_cache_column => by, touch: reflection.options[:touch])
+        end
+
         def find_target?
           !loaded? && foreign_key_present? && klass
         end
@@ -78,11 +97,11 @@ module ActiveRecord
         end
 
         def replace_keys(record)
-          owner[reflection.foreign_key] = record ? record._read_attribute(primary_key(record)) : nil
+          owner[reflection.foreign_key] = record ? record._read_attribute(primary_key(record.class)) : nil
         end
 
-        def primary_key(record)
-          reflection.association_primary_key(record.class)
+        def primary_key(klass)
+          reflection.association_primary_key(klass)
         end
 
         def foreign_key_present?
@@ -96,11 +115,6 @@ module ActiveRecord
           inverse && inverse.has_one?
         end
 
-        def counter_cache_target
-          primary_key = reflection.association_primary_key(klass)
-          klass.unscoped.where!(primary_key => owner._read_attribute(reflection.foreign_key))
-        end
-
         def stale_state
           result = owner._read_attribute(reflection.foreign_key) { |n| owner.send(:missing_attribute, n, caller) }
           result && result.to_s
diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 374247ffec..5e18f0e16f 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -21,47 +21,16 @@ module ActiveRecord::Associations::Builder # :nodoc:
       add_default_callbacks(model, reflection)       if reflection.options[:default]
     end
 
-    def self.define_accessors(mixin, reflection)
-      super
-      add_counter_cache_methods mixin
-    end
-
-    def self.add_counter_cache_methods(mixin)
-      return if mixin.method_defined? :belongs_to_counter_cache_after_update
-
-      mixin.class_eval do
-        def belongs_to_counter_cache_after_update(reflection)
-          if association(reflection.name).target_changed?
-            if reflection.polymorphic?
-              model_was = attribute_before_last_save(reflection.foreign_type).try(:constantize)
-            else
-              model_was = reflection.klass
-            end
-
-            foreign_key_was = attribute_before_last_save(reflection.foreign_key)
-            cache_column = reflection.counter_cache_column
-
-            association(reflection.name).increment_counters
-
-            if foreign_key_was && model_was < ActiveRecord::Base
-              counter_cache_target(reflection, model_was, foreign_key_was).update_counters(cache_column => -1)
-            end
-          end
-        end
-
-        private
-          def counter_cache_target(reflection, model, foreign_key)
-            primary_key = reflection.association_primary_key(model)
-            model.unscoped.where!(primary_key => foreign_key)
-          end
-      end
-    end
-
     def self.add_counter_cache_callbacks(model, reflection)
       cache_column = reflection.counter_cache_column
 
       model.after_update lambda { |record|
-        record.belongs_to_counter_cache_after_update(reflection)
+        association = association(reflection.name)
+
+        if association.target_changed?
+          association.increment_counters
+          association.decrement_counters_before_last_save
+        end
       }
 
       klass = reflection.class_name.safe_constantize