From 2e379c1e63b3646f9aff4d7e242ca37b4a57f529 Mon Sep 17 00:00:00 2001
From: jeroeningen <jeroen@jeroen.(none)>
Date: Sat, 26 May 2012 19:17:05 +0200
Subject: [PATCH] ActiveRecord#pluck now accepts multiple columns

---
 .../active_record/relation/calculations.rb    | 50 +++++++++++++------
 activerecord/test/cases/calculations_test.rb  | 17 ++++++-
 2 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 22c3e6a324..0abf57cd91 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -139,6 +139,12 @@ def calculate(operation, column_name, options = {})
     #   # 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
+    #
     #   Person.uniq.pluck(:role)
     #   # SELECT DISTINCT role FROM people
     #   # => ['admin', 'member', 'guest']
@@ -151,29 +157,43 @@ def calculate(operation, column_name, options = {})
     #   # SELECT DATEDIFF(updated_at, created_at) FROM people
     #   # => ['0', '27761', '173']
     #
-    def pluck(column_name)
-      if column_name.is_a?(Symbol) && column_names.include?(column_name.to_s)
-        column_name = "#{table_name}.#{column_name}"
+    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}"
+        else
+          column_names = column_names.collect{|column_name| "#{table_name}.#{column_name}"}
+        end
       end
 
-      if has_include?(column_name)
-        construct_relation_for_association_calculations.pluck(column_name)
+      if has_include?(column_names)
+        construct_relation_for_association_calculations.pluck(column_names)
       else
-        result = klass.connection.select_all(select(column_name).arel, nil, bind_values)
-
-        key    = result.columns.first
-        column = klass.column_types.fetch(key) {
-          result.column_types.fetch(key) {
-            Class.new { def type_cast(v); v; end }.new
+        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|
+          klass.column_types.fetch(key) {
+            result.column_types.fetch(key) {
+              Class.new { def type_cast(v); v; end }.new
+            }
           }
-        }
+        end
 
         result.map do |attributes|
-          raise ArgumentError, "Pluck expects to select just one attribute: #{attributes.inspect}" unless attributes.one?
+          if attributes.one?
+            value = klass.initialize_attributes(attributes).values.first
 
-          value = klass.initialize_attributes(attributes).values.first
+            columns.first.type_cast(value)
+          else
+            values = klass.initialize_attributes(attributes).values
 
-          column.type_cast(value)
+            values.each_with_index.map do |value, i|
+              columns[i].type_cast(value)
+            end
+          end
         end
       end
     end
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index f748b897ee..e86cf33b66 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -532,7 +532,7 @@ def test_pluck_with_selection_clause
     assert_equal [50 + 53 + 55 + 60], Account.pluck('SUM(DISTINCT(credit_limit)) as credit_limit')
   end
 
-  def test_pluck_expects_a_single_selection
+  def test_pluck_expects_a_multiple_selection_as_array
     assert_raise(ArgumentError) { Account.pluck 'id, credit_limit' }
   end
 
@@ -546,4 +546,19 @@ def test_pluck_not_auto_table_name_prefix_if_column_included
     assert_equal Company.count, ids.length
     assert_equal [7], ids.compact
   end
+
+  def test_pluck_multiple_columns
+    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])
+    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)
+  end
 end
-- 
GitLab