Repository: rails/rails
File: activerecord/lib/active_record/relation/calculations.rb (Ruby)

Code before:
    #   # SELECT people.name FROM people WHERE id = 1 LIMIT 1
    #   # => 'David'
    #
    #   Person.where(id: 1).pick(:name, :email_address)
    #   # SELECT people.name, people.email_address FROM people WHERE id = 1 LIMIT 1
    #   # => [ 'David', 'david@loudthinking.com' ]
    def pick(*column_names)
      if loaded? && all_attributes?(column_names)
        result = records.pick(*column_names)
        return @async ? Promise::Complete.new(result) : result
      end

      limit(1).pluck(*column_names).then(&:first)
    end

    # Same as <tt>#pick</tt> but perform the query asynchronously and returns an <tt>ActiveRecord::Promise</tt>
    def async_pick(*column_names)
      async.pick(*column_names)
    end

    # Returns the base model's ID's for the relation using the table's primary key
    #
    #   Person.ids # SELECT people.id FROM people
    #   Person.joins(:companies).ids # SELECT people.id FROM people INNER JOIN companies ON companies.id = people.company_id
    def ids
      if loaded?
        result = records.pluck(*Array(primary_key))
        return @async ? Promise::Complete.new(result) : result
      end

      if has_include?(primary_key)
        relation = apply_join_dependency.distinct
        return relation.ids
      end

      columns = arel_columns(Array(primary_key))
      relation = spawn
      relation.select_values = columns
      result = if relation.where_clause.contradiction?
        ActiveRecord::Result.empty
      else
        skip_query_cache_if_necessary do
          klass.connection.select_all(relation, "#{klass.name} Ids", async: @async)
        end
      end

      result.then { |result| type_cast_pluck_values(result, columns) }
    end

    # Same as <tt>#ids</tt> but perform the query asynchronously and returns an <tt>ActiveRecord::Promise</tt>
    def async_ids

Review comment (bug):
It's a nitpick but perhaps we could find a bit more suitable name as the value still represents a singular key whether it's a single-column or a set of columns. Perhaps `pk` or `primary_key_array` would be a bit more suitable name


Also I was under wrong impression that we already have a memoized `Array(primary_key)` value defined on the model class and was going to suggest using it here. But it looks like we don't have one yet so it's okay to explicitly wrap it here as eventually it will be refactored once we add a concept like:
```ruby
def composite_primary_key
  @_composite_primary_key ||= Array(primary_key)
end
```

Code after:
    #   # SELECT people.name FROM people WHERE id = 1 LIMIT 1
    #   # => 'David'
    #
    #   Person.where(id: 1).pick(:name, :email_address)
    #   # SELECT people.name, people.email_address FROM people WHERE id = 1 LIMIT 1
    #   # => [ 'David', 'david@loudthinking.com' ]
    def pick(*column_names)
      if loaded? && all_attributes?(column_names)
        result = records.pick(*column_names)
        return @async ? Promise::Complete.new(result) : result
      end

      limit(1).pluck(*column_names).then(&:first)
    end

    # Same as <tt>#pick</tt> but perform the query asynchronously and returns an <tt>ActiveRecord::Promise</tt>
    def async_pick(*column_names)
      async.pick(*column_names)
    end

    # Returns the base model's ID's for the relation using the table's primary key
    #
    #   Person.ids # SELECT people.id FROM people
    #   Person.joins(:companies).ids # SELECT people.id FROM people INNER JOIN companies ON companies.id = people.company_id
    def ids
      primary_key_array = Array(primary_key)

      if loaded?
        result = records.pluck(*primary_key_array)
        return @async ? Promise::Complete.new(result) : result
      end

      if has_include?(primary_key)
        relation = apply_join_dependency.group(*primary_key_array)
        return relation.ids
      end

      columns = arel_columns(primary_key_array)
      relation = spawn
      relation.select_values = columns

      result = if relation.where_clause.contradiction?
        ActiveRecord::Result.empty
      else
        skip_query_cache_if_necessary do
          klass.connection.select_all(relation, "#{klass.name} Ids", async: @async)
        end
      end

      result.then { |result| type_cast_pluck_values(result, columns) }
    end