Refactor pluck with multiple columns

Ensure it works with mix of symbols and strings, and with a select
clause possibly containing more than one column.

Also remove support for pluck with an array of columns, in favor of
passing the list of attributes:

    Model.pluck(:a, :b)

See comments: https://github.com/rails/rails/pull/6500#issuecomment-6030292

2 files changed, 29 insertions, 36 deletions

diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 0abf57cd91..9c27392299 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -107,7 +107,6 @@ module ActiveRecord
       relation = with_default_scope
 
       if relation.equal?(self)
-
         if has_include?(column_name)
           construct_relation_for_association_calculations.calculate(operation, column_name, options)
         else
@@ -139,9 +138,6 @@ module ActiveRecord
     #   # SELECT people.id FROM people
     #   # => [1, 2, 3]
     #
-    #   Person.pluck([:id, :name])
-    #   # SELECT people.id, people.name FROM people
-    #
     #   Person.pluck(:id, :name)
     #   # SELECT people.id, people.name FROM people
     #
@@ -158,23 +154,19 @@ module ActiveRecord
     #   # => ['0', '27761', '173']
     #
     def pluck(*column_names)
-      column_names = column_names.flatten
-
-      if column_names.first.is_a?(Symbol) && self.column_names.include?(column_names.first.to_s)
-        if column_names.one?
-          column_names = "#{table_name}.#{column_names.first}"
+      column_names.map! do |column_name|
+        if column_name.is_a?(Symbol) && self.column_names.include?(column_name.to_s)
+          "#{table_name}.#{column_name}"
         else
-          column_names = column_names.collect{|column_name| "#{table_name}.#{column_name}"}
+          column_name
         end
       end
 
-      if has_include?(column_names)
-        construct_relation_for_association_calculations.pluck(column_names)
+      if has_include?(column_names.first)
+        construct_relation_for_association_calculations.pluck(*column_names)
       else
-        result = klass.connection.select_all(select(column_names).arel, nil, bind_values)
-
-        keys   = column_names.is_a?(Array) && !column_names.one? ? result.columns : [result.columns.first]
-        columns = keys.map do |key|
+        result  = klass.connection.select_all(select(column_names).arel, nil, bind_values)
+        columns = result.columns.map do |key|
           klass.column_types.fetch(key) {
             result.column_types.fetch(key) {
               Class.new { def type_cast(v); v; end }.new
@@ -182,19 +174,14 @@ module ActiveRecord
           }
         end
 
-        result.map do |attributes|
-          if attributes.one?
-            value = klass.initialize_attributes(attributes).values.first
-
-            columns.first.type_cast(value)
-          else
-            values = klass.initialize_attributes(attributes).values
+        result = result.map do |attributes|
+          values = klass.initialize_attributes(attributes).values
 
-            values.each_with_index.map do |value, i|
-              columns[i].type_cast(value)
-            end
+          columns.zip(values).map do |column, value|
+            column.type_cast(value)
           end
         end
+        columns.one? ? result.map!(&:first) : result
       end
     end
 
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index e86cf33b66..4df613488a 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -532,10 +532,6 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal [50 + 53 + 55 + 60], Account.pluck('SUM(DISTINCT(credit_limit)) as credit_limit')
   end
 
-  def test_pluck_expects_a_multiple_selection_as_array
-    assert_raise(ArgumentError) { Account.pluck 'id, credit_limit' }
-  end
-
   def test_plucks_with_ids
     assert_equal Company.all.map(&:id).sort, Company.ids.sort
   end
@@ -551,14 +547,24 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal [
       [1, "The First Topic"], [2, "The Second Topic of the day"],
       [3, "The Third Topic of the day"], [4, "The Fourth Topic of the day"]
-    ], Topic.order(:id).pluck([:id, :title])
+    ], Topic.order(:id).pluck(:id, :title)
     assert_equal [
       [1, "The First Topic", "David"], [2, "The Second Topic of the day", "Mary"],
       [3, "The Third Topic of the day", "Carl"], [4, "The Fourth Topic of the day", "Carl"]
-    ], Topic.order(:id).pluck([:id, :title, :author_name])
-    assert_equal [
-      [1, "The First Topic"], [2, "The Second Topic of the day"],
-      [3, "The Third Topic of the day"], [4, "The Fourth Topic of the day"]
-    ], Topic.order(:id).pluck(:id, :title)
+    ], Topic.order(:id).pluck(:id, :title, :author_name)
+  end
+
+  def test_pluck_with_multiple_columns_and_selection_clause
+    assert_equal [[1, 50], [2, 50], [3, 50], [4, 60], [5, 55], [6, 53]],
+      Account.pluck('id, credit_limit')
+  end
+
+  def test_pluck_with_multiple_columns_and_includes
+    Company.create!(:name => "test", :contracts => [Contract.new(:developer_id => 7)])
+    companies_and_developers = Company.order('companies.id').includes(:contracts).pluck(:name, :developer_id)
+
+    assert_equal Company.count, companies_and_developers.length
+    assert_equal ["37signals", nil], companies_and_developers.first
+    assert_equal ["test", 7], companies_and_developers.last
   end
 end