Add Relation#count

3 files changed, 82 insertions, 0 deletions

diff --git a/activerecord/CHANGELOG b/activerecord/CHANGELOG
index 93c4696eca..5af66c3519 100644
--- a/activerecord/CHANGELOG
+++ b/activerecord/CHANGELOG
@@ -1,5 +1,12 @@
 *Edge*
 
+* Add Relation#count. [Pratik Naik]
+
+  legends = People.where("age > 100")
+  legends.count
+  legends.count(:age, :distinct => true)
+  legends.select('id').count
+
 * Add Model.readonly and association_collection#readonly finder method. [Pratik Naik]
 
   Post.readonly.to_a # Load all posts in readonly mode
diff --git a/activerecord/lib/active_record/relation.rb b/activerecord/lib/active_record/relation.rb
index 93f7b74c68..cb252eea70 100644
--- a/activerecord/lib/active_record/relation.rb
+++ b/activerecord/lib/active_record/relation.rb
@@ -204,6 +204,20 @@ module ActiveRecord
       end
     end
 
+    def count(*args)
+      column_name, options = construct_count_options_from_args(*args)
+      distinct = options[:distinct] ? true : false
+
+      column = if @klass.column_names.include?(column_name.to_s)
+        Arel::Attribute.new(@relation.table, column_name)
+      else
+        Arel::SqlLiteral.new(column_name == :all ? "*" : column_name.to_s)
+      end
+
+      relation = select(column.count(distinct))
+      @klass.connection.select_value(relation.to_sql).to_i
+    end
+
     def destroy_all
       to_a.each {|object| object.destroy}
       reset
@@ -337,5 +351,36 @@ module ActiveRecord
       }.join(',')
     end
 
+    def construct_count_options_from_args(*args)
+      options     = {}
+      column_name = :all
+
+      # We need to handle
+      #   count()
+      #   count(:column_name=:all)
+      #   count(options={})
+      #   count(column_name=:all, options={})
+      #   selects specified by scopes
+
+      # TODO : relation.projections only works when .select() was last in the chain. Fix it!
+      case args.size
+      when 0
+        column_name = @relation.send(:select_clauses).join(', ') if @relation.respond_to?(:projections) && @relation.projections.present?
+      when 1
+        if args[0].is_a?(Hash)
+          column_name = @relation.send(:select_clauses).join(', ') if @relation.respond_to?(:projections) && @relation.projections.present?
+          options = args[0]
+        else
+          column_name = args[0]
+        end
+      when 2
+        column_name, options = args
+      else
+        raise ArgumentError, "Unexpected parameters passed to count(): #{args.inspect}"
+      end
+
+      [column_name || :all, options]
+    end
+
   end
 end
diff --git a/activerecord/test/cases/relations_test.rb b/activerecord/test/cases/relations_test.rb
index cf6708cc8b..ded4f2f479 100644
--- a/activerecord/test/cases/relations_test.rb
+++ b/activerecord/test/cases/relations_test.rb
@@ -358,4 +358,34 @@ class RelationTest < ActiveRecord::TestCase
     assert_raises(ArgumentError) { Post.scoped & Developer.scoped }
   end
 
+  def test_count
+    posts = Post.scoped
+
+    assert_equal 7, posts.count
+    assert_equal 7, posts.count(:all)
+    assert_equal 7, posts.count(:id)
+
+    assert_equal 1, posts.where('comments_count > 1').count
+    assert_equal 5, posts.where(:comments_count => 0).count
+  end
+
+  def test_count_with_distinct
+    posts = Post.scoped
+
+    assert_equal 3, posts.count(:comments_count, :distinct => true)
+    assert_equal 7, posts.count(:comments_count, :distinct => false)
+
+    assert_equal 3, posts.select(:comments_count).count(:distinct => true)
+    assert_equal 7, posts.select(:comments_count).count(:distinct => false)
+  end
+
+  def test_count_explicit_columns
+    Post.update_all(:comments_count => nil)
+    posts = Post.scoped
+
+    assert_equal 7, posts.select('comments_count').count('id')
+    assert_equal 0, posts.select('comments_count').count
+    assert_equal 0, posts.count(:comments_count)
+    assert_equal 0, posts.count('comments_count')
+  end
 end