Skip to content
This repository
Browse code

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: #6500 (comment)
  • Loading branch information...
commit 6aae17e85613fe8c2816ba278f9348f168692479 1 parent 2e379c1
Carlos Antonio da Silva authored June 19, 2012
39  activerecord/lib/active_record/relation/calculations.rb
@@ -107,7 +107,6 @@ def calculate(operation, column_name, options = {})
107 107
       relation = with_default_scope
108 108
 
109 109
       if relation.equal?(self)
110  
-
111 110
         if has_include?(column_name)
112 111
           construct_relation_for_association_calculations.calculate(operation, column_name, options)
113 112
         else
@@ -139,9 +138,6 @@ def calculate(operation, column_name, options = {})
139 138
     #   # SELECT people.id FROM people
140 139
     #   # => [1, 2, 3]
141 140
     #
142  
-    #   Person.pluck([:id, :name])
143  
-    #   # SELECT people.id, people.name FROM people
144  
-    #
145 141
     #   Person.pluck(:id, :name)
146 142
     #   # SELECT people.id, people.name FROM people
147 143
     #
@@ -158,23 +154,19 @@ def calculate(operation, column_name, options = {})
158 154
     #   # => ['0', '27761', '173']
159 155
     #
160 156
     def pluck(*column_names)
161  
-      column_names = column_names.flatten
162  
-
163  
-      if column_names.first.is_a?(Symbol) && self.column_names.include?(column_names.first.to_s)
164  
-        if column_names.one?
165  
-          column_names = "#{table_name}.#{column_names.first}"
  157
+      column_names.map! do |column_name|
  158
+        if column_name.is_a?(Symbol) && self.column_names.include?(column_name.to_s)
  159
+          "#{table_name}.#{column_name}"
166 160
         else
167  
-          column_names = column_names.collect{|column_name| "#{table_name}.#{column_name}"}
  161
+          column_name
168 162
         end
169 163
       end
170 164
 
171  
-      if has_include?(column_names)
172  
-        construct_relation_for_association_calculations.pluck(column_names)
  165
+      if has_include?(column_names.first)
  166
+        construct_relation_for_association_calculations.pluck(*column_names)
173 167
       else
174  
-        result = klass.connection.select_all(select(column_names).arel, nil, bind_values)
175  
-
176  
-        keys   = column_names.is_a?(Array) && !column_names.one? ? result.columns : [result.columns.first]
177  
-        columns = keys.map do |key|
  168
+        result  = klass.connection.select_all(select(column_names).arel, nil, bind_values)
  169
+        columns = result.columns.map do |key|
178 170
           klass.column_types.fetch(key) {
179 171
             result.column_types.fetch(key) {
180 172
               Class.new { def type_cast(v); v; end }.new
@@ -182,19 +174,14 @@ def pluck(*column_names)
182 174
           }
183 175
         end
184 176
 
185  
-        result.map do |attributes|
186  
-          if attributes.one?
187  
-            value = klass.initialize_attributes(attributes).values.first
188  
-
189  
-            columns.first.type_cast(value)
190  
-          else
191  
-            values = klass.initialize_attributes(attributes).values
  177
+        result = result.map do |attributes|
  178
+          values = klass.initialize_attributes(attributes).values
192 179
 
193  
-            values.each_with_index.map do |value, i|
194  
-              columns[i].type_cast(value)
195  
-            end
  180
+          columns.zip(values).map do |column, value|
  181
+            column.type_cast(value)
196 182
           end
197 183
         end
  184
+        columns.one? ? result.map!(&:first) : result
198 185
       end
199 186
     end
200 187
 
26  activerecord/test/cases/calculations_test.rb
@@ -532,10 +532,6 @@ def test_pluck_with_selection_clause
532 532
     assert_equal [50 + 53 + 55 + 60], Account.pluck('SUM(DISTINCT(credit_limit)) as credit_limit')
533 533
   end
534 534
 
535  
-  def test_pluck_expects_a_multiple_selection_as_array
536  
-    assert_raise(ArgumentError) { Account.pluck 'id, credit_limit' }
537  
-  end
538  
-
539 535
   def test_plucks_with_ids
540 536
     assert_equal Company.all.map(&:id).sort, Company.ids.sort
541 537
   end
@@ -551,14 +547,24 @@ def test_pluck_multiple_columns
551 547
     assert_equal [
552 548
       [1, "The First Topic"], [2, "The Second Topic of the day"],
553 549
       [3, "The Third Topic of the day"], [4, "The Fourth Topic of the day"]
554  
-    ], Topic.order(:id).pluck([:id, :title])
  550
+    ], Topic.order(:id).pluck(:id, :title)
555 551
     assert_equal [
556 552
       [1, "The First Topic", "David"], [2, "The Second Topic of the day", "Mary"],
557 553
       [3, "The Third Topic of the day", "Carl"], [4, "The Fourth Topic of the day", "Carl"]
558  
-    ], Topic.order(:id).pluck([:id, :title, :author_name])
559  
-    assert_equal [
560  
-      [1, "The First Topic"], [2, "The Second Topic of the day"],
561  
-      [3, "The Third Topic of the day"], [4, "The Fourth Topic of the day"]
562  
-    ], Topic.order(:id).pluck(:id, :title)
  554
+    ], Topic.order(:id).pluck(:id, :title, :author_name)
  555
+  end
  556
+
  557
+  def test_pluck_with_multiple_columns_and_selection_clause
  558
+    assert_equal [[1, 50], [2, 50], [3, 50], [4, 60], [5, 55], [6, 53]],
  559
+      Account.pluck('id, credit_limit')
  560
+  end
  561
+
  562
+  def test_pluck_with_multiple_columns_and_includes
  563
+    Company.create!(:name => "test", :contracts => [Contract.new(:developer_id => 7)])
  564
+    companies_and_developers = Company.order('companies.id').includes(:contracts).pluck(:name, :developer_id)
  565
+
  566
+    assert_equal Company.count, companies_and_developers.length
  567
+    assert_equal ["37signals", nil], companies_and_developers.first
  568
+    assert_equal ["test", 7], companies_and_developers.last
563 569
   end
564 570
 end

0 notes on commit 6aae17e

Please sign in to comment.
Something went wrong with that request. Please try again.