Refactored AssociationCollection#count for uniformity and Ruby 1.8.7 support.

[#831 state:resolved]

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

4 files changed, 39 insertions, 36 deletions

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])