Call NPlusOneQuery's call_association when calling count on collectin…

…o assocation
flyerhzm · Nov 5, 2023 · b6812c2 · b6812c2
1 parent ac58d67
commit b6812c2
Show file tree

Hide file tree

Showing 9 changed files with 26 additions and 2 deletions.
diff --git a/lib/bullet/active_record4.rb b/lib/bullet/active_record4.rb
@@ -181,6 +181,7 @@ def has_cached_counter?(reflection = reflection())
         def count(column_name = nil, options = {})
           if Bullet.start?
             Bullet::Detector::CounterCache.add_counter_cache(proxy_association.owner, proxy_association.reflection.name)
+            Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
           end
           super(column_name, options)
         end

diff --git a/lib/bullet/active_record41.rb b/lib/bullet/active_record41.rb
@@ -173,6 +173,7 @@ def count_records
         def count(column_name = nil, options = {})
           if Bullet.start?
             Bullet::Detector::CounterCache.add_counter_cache(proxy_association.owner, proxy_association.reflection.name)
+            Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
           end
           super(column_name, options)
         end

diff --git a/lib/bullet/active_record42.rb b/lib/bullet/active_record42.rb
@@ -238,6 +238,7 @@ def count_records
         def count(column_name = nil, options = {})
           if Bullet.start?
             Bullet::Detector::CounterCache.add_counter_cache(proxy_association.owner, proxy_association.reflection.name)
+            Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
           end
           super(column_name, options)
         end

diff --git a/lib/bullet/active_record5.rb b/lib/bullet/active_record5.rb
@@ -269,6 +269,7 @@ def count
                 proxy_association.owner,
                 proxy_association.reflection.name
               )
+              Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
             end
             super
           end

diff --git a/lib/bullet/active_record52.rb b/lib/bullet/active_record52.rb
@@ -251,6 +251,7 @@ def count(column_name = nil)
                 proxy_association.owner,
                 proxy_association.reflection.name
               )
+              Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
             end
             super(column_name)
           end

diff --git a/lib/bullet/active_record61.rb b/lib/bullet/active_record61.rb
@@ -278,6 +278,7 @@ def count(column_name = nil)
                 proxy_association.owner,
                 proxy_association.reflection.name
               )
+              Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
             end
             super(column_name)
           end

diff --git a/lib/bullet/active_record70.rb b/lib/bullet/active_record70.rb
@@ -294,6 +294,7 @@ def count(column_name = nil)
                 proxy_association.owner,
                 proxy_association.reflection.name
               )
+              Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
             end
             super(column_name)
           end

diff --git a/lib/bullet/active_record71.rb b/lib/bullet/active_record71.rb
@@ -294,6 +294,7 @@ def count(column_name = nil)
                 proxy_association.owner,
                 proxy_association.reflection.name
               )
+              Bullet::Detector::NPlusOneQuery.call_association(proxy_association.owner, proxy_association.reflection.name)
             end
             super(column_name)
           end

diff --git a/spec/integration/active_record/association_spec.rb b/spec/integration/active_record/association_spec.rb
@@ -88,13 +88,29 @@
         expect(Bullet::Detector::Association).to be_detecting_unpreloaded_association_for(Post, :comments)
       end
 
-      it 'should not detect unused preload person => pets with empty?' do
-        Person.all.each { |person| person.pets.empty? }
+      it 'should not detect unused preload post => comment with empty?' do
+        Post.includes(:comments).each { |post| post.comments.empty? }
         Bullet::Detector::UnusedEagerLoading.check_unused_preload_associations
         expect(Bullet::Detector::Association).not_to be_has_unused_preload_associations
 
         expect(Bullet::Detector::Association).to be_completely_preloading_associations
       end
+
+      it 'should not detect unused preload post => comment with count' do
+        Post.includes(:comments).each { |post| post.comments.count }
+        Bullet::Detector::UnusedEagerLoading.check_unused_preload_associations
+        expect(Bullet::Detector::Association).not_to be_has_unused_preload_associations
+
+        expect(Bullet::Detector::Association).to be_completely_preloading_associations
+      end
+
+      it 'should detect non preload post => comments with count' do
+        Post.all.each { |post| post.comments.count }
+        Bullet::Detector::UnusedEagerLoading.check_unused_preload_associations
+        expect(Bullet::Detector::Association).not_to be_has_unused_preload_associations
+
+        expect(Bullet::Detector::Association).to be_detecting_unpreloaded_association_for(Post, :comments)
+      end
     end
 
     context 'category => posts => comments' do