Allow select with Arel and count as well as calculations with Arel

It allows a query like `User.select(:name).count` to be written
using Arel as `User.select(User.arel_table[:name]).count`.

It exposes the calculations API to accept Arel nodes:
`User.count(User.arel_table[:name])`, `User.sum(User.arel_table[:id])`,
`Account.average(Account.arel_table[:credit_limit])`,
`Account.maximum(Account.arel_table[:credit_limit])` and
`Account.minimum(Account.arel_table[:credit_limit])`.

2 files changed, 55 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index df72ba7e9c..3b7f7ca151 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -217,6 +217,8 @@ module ActiveRecord
     end
 
     def aggregate_column(column_name)
+      return column_name if Arel::Expressions === column_name
+
       if @klass.column_names.include?(column_name.to_s)
         Arel::Attribute.new(@klass.unscoped.table, column_name)
       else
@@ -366,9 +368,9 @@ module ActiveRecord
       end
     end
 
-    # TODO: refactor to allow non-string `select_values` (eg. Arel nodes).
     def select_for_count
       if select_values.present?
+        return select_values.first if select_values.one?
         select_values.join(", ")
       else
         :all
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index aa10817527..5327849697 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -31,11 +31,20 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal 318, Account.sum(:credit_limit)
   end
 
+  def test_should_sum_arel_attribute
+    assert_equal 318, Account.sum(Account.arel_table[:credit_limit])
+  end
+
   def test_should_average_field
     value = Account.average(:credit_limit)
     assert_equal 53.0, value
   end
 
+  def test_should_average_arel_attribute
+    value = Account.average(Account.arel_table[:credit_limit])
+    assert_equal 53.0, value
+  end
+
   def test_should_resolve_aliased_attributes
     assert_equal 318, Account.sum(:available_credit)
   end
@@ -60,14 +69,26 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal 60, Account.maximum(:credit_limit)
   end
 
+  def test_should_get_maximum_of_arel_attribute
+    assert_equal 60, Account.maximum(Account.arel_table[:credit_limit])
+  end
+
   def test_should_get_maximum_of_field_with_include
     assert_equal 55, Account.where("companies.name != 'Summit'").references(:companies).includes(:firm).maximum(:credit_limit)
   end
 
+  def test_should_get_maximum_of_arel_attribute_with_include
+    assert_equal 55, Account.where("companies.name != 'Summit'").references(:companies).includes(:firm).maximum(Account.arel_table[:credit_limit])
+  end
+
   def test_should_get_minimum_of_field
     assert_equal 50, Account.minimum(:credit_limit)
   end
 
+  def test_should_get_minimum_of_arel_attribute
+    assert_equal 50, Account.minimum(Account.arel_table[:credit_limit])
+  end
+
   def test_should_group_by_field
     c = Account.group(:firm_id).sum(:credit_limit)
     [1,6,2].each do |firm_id|
@@ -131,6 +152,14 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal 3, accounts.select(:firm_id).count
   end
 
+  def test_limit_should_apply_before_count_arel_attribute
+    accounts = Account.limit(3).where('firm_id IS NOT NULL')
+
+    firm_id_attribute = Account.arel_table[:firm_id]
+    assert_equal 3, accounts.count(firm_id_attribute)
+    assert_equal 3, accounts.select(firm_id_attribute).count
+  end
+
   def test_count_should_shortcut_with_limit_zero
     accounts = Account.limit(0)
 
@@ -353,10 +382,23 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal 6, Account.select("DISTINCT accounts.id").includes(:firm).count
   end
 
+  def test_count_selected_arel_attribute
+    assert_equal 5, Account.select(Account.arel_table[:firm_id]).count
+    assert_equal 4, Account.distinct.select(Account.arel_table[:firm_id]).count
+  end
+
   def test_count_with_column_parameter
     assert_equal 5, Account.count(:firm_id)
   end
 
+  def test_count_with_arel_attribute
+    assert_equal 5, Account.count(Account.arel_table[:firm_id])
+  end
+
+  def test_count_with_arel_star
+    assert_equal 6, Account.count(Arel.star)
+  end
+
   def test_count_with_distinct
     assert_equal 4, Account.select(:credit_limit).distinct.count
 
@@ -378,6 +420,16 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal 4, Account.joins(:firm).distinct.count('companies.id')
   end
 
+  def test_count_arel_attribute_in_joined_table_with
+    assert_equal 5, Account.joins(:firm).count(Company.arel_table[:id])
+    assert_equal 4, Account.joins(:firm).distinct.count(Company.arel_table[:id])
+  end
+
+  def test_count_selected_arel_attribute_in_joined_table
+    assert_equal 5, Account.joins(:firm).select(Company.arel_table[:id]).count
+    assert_equal 4, Account.joins(:firm).distinct.select(Company.arel_table[:id]).count
+  end
+
   def test_should_count_field_in_joined_table_with_group_by
     c = Account.group('accounts.firm_id').joins(:firm).count('companies.id')