Refactored AssociationCollection#count for uniformity and Ruby 1.8.7 support.

[#831 state:resolved]

Signed-off-by: Jeremy Kemper <jeremy@bitsweat.net>

6 files changed, 49 insertions, 36 deletions

diff --git a/activerecord/lib/active_record/associations.rb b/activerecord/lib/active_record/associations.rb
index 4d935612ca..98710dee09 100755
--- a/activerecord/lib/active_record/associations.rb
+++ b/activerecord/lib/active_record/associations.rb
@@ -1164,6 +1164,9 @@ module ActiveRecord
       #   If true, duplicate associated objects will be ignored by accessors and query methods.
       # [:finder_sql]
       #   Overwrite the default generated SQL statement used to fetch the association with a manual statement
+      # [:counter_sql]
+      #   Specify a complete SQL statement to fetch the size of the association. If <tt>:finder_sql</tt> is
+      #   specified but not <tt>:counter_sql</tt>, <tt>:counter_sql</tt> will be generated by replacing <tt>SELECT ... FROM</tt> with <tt>SELECT COUNT(*) FROM</tt>.
       # [:delete_sql]
       #   Overwrite the default generated SQL statement used to remove links between the associated
       #   classes with a manual statement.
diff --git a/activerecord/lib/active_record/associations/association_collection.rb b/activerecord/lib/active_record/associations/association_collection.rb
index 9061037b39..5092ccc1dc 100644
--- a/activerecord/lib/active_record/associations/association_collection.rb
+++ b/activerecord/lib/active_record/associations/association_collection.rb
@@ -128,6 +128,35 @@ module ActiveRecord
         end
       end
 
+      # Count all records using SQL. If the +:counter_sql+ option is set for the association, it will
+      # be used for the query. If no +:counter_sql+ was supplied, but +:finder_sql+ was set, the
+      # descendant's +construct_sql+ method will have set :counter_sql automatically.
+      # Otherwise, construct options and pass them with scope to the target class's +count+.
+      def count(*args)
+        if @reflection.options[:counter_sql]
+          @reflection.klass.count_by_sql(@counter_sql)
+        else
+          column_name, options = @reflection.klass.send(:construct_count_options_from_args, *args)
+          if @reflection.options[:uniq]
+            # This is needed because 'SELECT count(DISTINCT *)..' is not valid SQL.
+            column_name = "#{@reflection.quoted_table_name}.#{@reflection.klass.primary_key}" if column_name == :all
+            options.merge!(:distinct => true)
+          end
+
+          value = @reflection.klass.send(:with_scope, construct_scope) { @reflection.klass.count(column_name, options) }
+
+          limit  = @reflection.options[:limit]
+          offset = @reflection.options[:offset]
+
+          if limit || offset
+            [ [value - offset.to_i, 0].max, limit.to_i ].min
+          else
+            value
+          end
+        end
+      end
+
+
       # Remove +records+ from this association.  Does not destroy +records+.
       def delete(*records)
         records = flatten_deeper(records)
diff --git a/activerecord/lib/active_record/associations/has_and_belongs_to_many_association.rb b/activerecord/lib/active_record/associations/has_and_belongs_to_many_association.rb
index e7e433b6b6..3d689098b5 100644
--- a/activerecord/lib/active_record/associations/has_and_belongs_to_many_association.rb
+++ b/activerecord/lib/active_record/associations/has_and_belongs_to_many_association.rb
@@ -78,6 +78,16 @@ module ActiveRecord
           end
 
           @join_sql = "INNER JOIN #{@owner.connection.quote_table_name @reflection.options[:join_table]} ON #{@reflection.quoted_table_name}.#{@reflection.klass.primary_key} = #{@owner.connection.quote_table_name @reflection.options[:join_table]}.#{@reflection.association_foreign_key}"
+
+          if @reflection.options[:counter_sql]
+            @counter_sql = interpolate_sql(@reflection.options[:counter_sql])
+          elsif @reflection.options[:finder_sql]
+            # replace the SELECT clause with COUNT(*), preserving any hints within /* ... */
+            @reflection.options[:counter_sql] = @reflection.options[:finder_sql].sub(/SELECT (\/\*.*?\*\/ )?(.*)\bFROM\b/im) { "SELECT #{$1}COUNT(*) FROM" }
+            @counter_sql = interpolate_sql(@reflection.options[:counter_sql])
+          else
+            @counter_sql = @finder_sql
+          end
         end
 
         def construct_scope
diff --git a/activerecord/lib/active_record/associations/has_many_association.rb b/activerecord/lib/active_record/associations/has_many_association.rb
index 1535995410..1838021d40 100644
--- a/activerecord/lib/active_record/associations/has_many_association.rb
+++ b/activerecord/lib/active_record/associations/has_many_association.rb
@@ -1,32 +1,6 @@
 module ActiveRecord
   module Associations
     class HasManyAssociation < AssociationCollection #:nodoc:
-      # Count the number of associated records. All arguments are optional.
-      def count(*args)
-        if @reflection.options[:counter_sql]
-          @reflection.klass.count_by_sql(@counter_sql)
-        elsif @reflection.options[:finder_sql]
-          @reflection.klass.count_by_sql(@finder_sql)
-        else
-          column_name, options = @reflection.klass.send(:construct_count_options_from_args, *args)          
-          options[:conditions] = options[:conditions].blank? ?
-            @finder_sql :
-            @finder_sql + " AND (#{sanitize_sql(options[:conditions])})"
-          options[:include] ||= @reflection.options[:include]
-
-          value = @reflection.klass.count(column_name, options)
-
-          limit  = @reflection.options[:limit]
-          offset = @reflection.options[:offset]
-
-          if limit || offset
-            [ [value - offset.to_i, 0].max, limit.to_i ].min
-          else
-            value
-          end
-        end
-      end
-
       protected
         def owner_quoted_id
           if @reflection.options[:primary_key]
diff --git a/activerecord/lib/active_record/associations/has_many_through_association.rb b/activerecord/lib/active_record/associations/has_many_through_association.rb
index 24b02efc35..84fa900f46 100644
--- a/activerecord/lib/active_record/associations/has_many_through_association.rb
+++ b/activerecord/lib/active_record/associations/has_many_through_association.rb
@@ -31,16 +31,6 @@ module ActiveRecord
         return count
       end
       
-      def count(*args)
-        column_name, options = @reflection.klass.send(:construct_count_options_from_args, *args)
-        if @reflection.options[:uniq]
-          # This is needed because 'SELECT count(DISTINCT *)..' is not valid SQL statement.
-          column_name = "#{@reflection.quoted_table_name}.#{@reflection.klass.primary_key}" if column_name == :all
-          options.merge!(:distinct => true) 
-        end
-        @reflection.klass.send(:with_scope, construct_scope) { @reflection.klass.count(column_name, options) } 
-      end
-
       protected
         def construct_find_options!(options)
           options[:select]  = construct_select(options[:select])
diff --git a/activerecord/test/cases/associations/has_and_belongs_to_many_associations_test.rb b/activerecord/test/cases/associations/has_and_belongs_to_many_associations_test.rb
index 0572418e3b..1b31d28679 100644
--- a/activerecord/test/cases/associations/has_and_belongs_to_many_associations_test.rb
+++ b/activerecord/test/cases/associations/has_and_belongs_to_many_associations_test.rb
@@ -703,4 +703,11 @@ class HasAndBelongsToManyAssociationsTest < ActiveRecord::TestCase
     # due to Unknown column 'authors.id'
     assert Category.find(1).posts_with_authors_sorted_by_author_id.find_by_title('Welcome to the weblog')
   end
+
+  def test_counting_on_habtm_association_and_not_array
+    david = Developer.find(1)
+    # Extra parameter just to make sure we aren't falling back to
+    # Array#count in Ruby >=1.8.7, which would raise an ArgumentError
+    assert_nothing_raised { david.projects.count(:all, :conditions => '1=1') }
+  end
 end