Merge pull request rails#49866 from Shopify/prebuild-counter-cache-association-names-after-schema-load

Author: byroot

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

@@ -38,6 +38,7 @@ def self.add_counter_cache_callbacks(model, reflection)
 
       klass = reflection.class_name.safe_constantize
       klass._counter_cache_columns |= [cache_column] if klass && klass.respond_to?(:_counter_cache_columns)
+      model.counter_cached_association_names |= [reflection.name]
     end
 
     def self.touch_record(o, changes, foreign_key, name, touch) # :nodoc:

activerecord/lib/active_record/counter_cache.rb

@@ -7,6 +7,7 @@ module CounterCache
 
     included do
       class_attribute :_counter_cache_columns, instance_accessor: false, default: []
+      class_attribute :counter_cached_association_names, instance_writer: false, default: []
     end
 
     module ClassMethods
@@ -180,14 +181,26 @@ def decrement_counter(counter_name, id, by: 1, touch: nil)
       def counter_cache_column?(name) # :nodoc:
         _counter_cache_columns.include?(name)
       end
+
+      def load_schema # :nodoc:
+        super
+
+        association_names = _reflections.filter_map do |name, reflection|
+          next unless reflection.belongs_to? && reflection.counter_cache_column
+
+          name.to_sym
+        end
+
+        self.counter_cached_association_names |= association_names
+      end
     end
 
     private
       def _create_record(attribute_names = self.attribute_names)
         id = super
 
-        each_counter_cached_associations do |association|
-          association.increment_counters
+        counter_cached_association_names.each do |association_name|
+          association(association_name).increment_counters
         end
 
         id
@@ -197,7 +210,8 @@ def destroy_row
         affected_rows = super
 
         if affected_rows > 0
-          each_counter_cached_associations do |association|
+          counter_cached_association_names.each do |association_name|
+            association = association(association_name)
             foreign_key = association.reflection.foreign_key.to_sym
             unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == foreign_key
               association.decrement_counters
@@ -207,11 +221,5 @@ def destroy_row
 
         affected_rows
       end
-
-      def each_counter_cached_associations
-        _reflections.each do |name, reflection|
-          yield association(name.to_sym) if reflection.belongs_to? && reflection.counter_cache_column
-        end
-      end
   end
 end