diff options
| -rw-r--r-- | activerecord/lib/active_record/relation/calculations.rb | 50 | ||||
| -rw-r--r-- | 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 @@ 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 + # # Person.uniq.pluck(:role) # # SELECT DISTINCT role FROM people # # => ['admin', 'member', 'guest'] @@ -151,29 +157,43 @@ module ActiveRecord # # 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 @@ 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_single_selection + def test_pluck_expects_a_multiple_selection_as_array assert_raise(ArgumentError) { Account.pluck 'id, credit_limit' } end @@ -546,4 +546,19 @@ class CalculationsTest < ActiveRecord::TestCase 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 |