Forward ActiveRecord::Relation#count to Enumerable#count if block given

2 files changed, 14 insertions, 5 deletions

diff --git a/activerecord/lib/active_record/associations/collection_association.rb b/activerecord/lib/active_record/associations/collection_association.rb
index 2dca6b612e..00355f3e89 100644
--- a/activerecord/lib/active_record/associations/collection_association.rb
+++ b/activerecord/lib/active_record/associations/collection_association.rb
@@ -246,9 +246,12 @@ module ActiveRecord
         end
       end
 
-      # Count all records using SQL.  Construct options and pass them with
-      # scope to the target class's +count+.
+      # Returns the number of records. If no arguments are given, it counts all
+      # columns using SQL. If one argument is given, it counts only the passed
+      # column using SQL. If a block is given, it counts the number of records
+      # yielding a true value.
       def count(column_name = nil)
+        return super if block_given?
         relation = scope
         if association_scope.distinct_value
           # This is needed because 'SELECT count(DISTINCT *)..' is not valid SQL.
diff --git a/activerecord/lib/active_record/associations/collection_proxy.rb b/activerecord/lib/active_record/associations/collection_proxy.rb
index b9aed05135..9350064028 100644
--- a/activerecord/lib/active_record/associations/collection_proxy.rb
+++ b/activerecord/lib/active_record/associations/collection_proxy.rb
@@ -715,12 +715,13 @@ module ActiveRecord
       end
       alias uniq distinct
 
-      # Count all records using SQL.
+      # Count all records.
       #
       #   class Person < ActiveRecord::Base
       #     has_many :pets
       #   end
       #
+      #   # This will perform the count using SQL.
       #   person.pets.count # => 3
       #   person.pets
       #   # => [
@@ -728,8 +729,13 @@ module ActiveRecord
       #   #       #<Pet id: 2, name: "Spook", person_id: 1>,
       #   #       #<Pet id: 3, name: "Choo-Choo", person_id: 1>
       #   #    ]
-      def count(column_name = nil)
-        @association.count(column_name)
+      #
+      # Passing a block will select all of a person's pets in SQL and then
+      # perform the count using Ruby.
+      #
+      #   person.pets.count { |pet| pet.name.include?('-') } # => 2
+      def count(column_name = nil, &block)
+        @association.count(column_name, &block)
       end
 
       # Returns the size of the collection. If the collection hasn't been loaded,