MONGOID-5900 Fix #pluck on associations

lib/mongoid/association/embedded/embeds_many/proxy.rb

@@ -14,6 +14,7 @@ class EmbedsMany
         # the array of child documents.
         class Proxy < Association::Many
           include Batchable
+          extend Forwardable
 
           # Class-level methods for the Proxy class.
           module ClassMethods
@@ -54,6 +55,8 @@ def foreign_key_suffix
 
           extend ClassMethods
 
+          def_delegators :criteria, :find, :pluck
+
           # Instantiate a new embeds_many association.
           #
           # @example Create the new association.
@@ -312,35 +315,6 @@ def exists?(id_or_conditions = :none)
             end
           end
 
-          # Finds a document in this association through several different
-          # methods.
-          #
-          # This method delegates to +Mongoid::Criteria#find+. If this method is
-          # not given a block, it returns one or many documents for the provided
-          # _id values.
-          #
-          # If this method is given a block, it returns the first document
-          # of those found by the current Criteria object for which the block
-          # returns a truthy value.
-          #
-          # @example Find a document by its id.
-          #   person.addresses.find(BSON::ObjectId.new)
-          #
-          # @example Find documents for multiple ids.
-          #   person.addresses.find([ BSON::ObjectId.new, BSON::ObjectId.new ])
-          #
-          # @example Finds the first matching document using a block.
-          #   person.addresses.find { |addr| addr.state == 'CA' }
-          #
-          # @param [ Object... ] *args Various arguments.
-          # @param &block Optional block to pass.
-          # @yield [ Object ] Yields each enumerable element to the block.
-          #
-          # @return [ Document | Array<Document> | nil ] A document or matching documents.
-          def find(...)
-            criteria.find(...)
-          end
-
           # Get all the documents in the association that are loaded into memory.
           #
           # @example Get the in memory documents.

lib/mongoid/association/referenced/has_many/enumerable.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 # rubocop:todo all
 
+require 'mongoid/pluckable'
+
 module Mongoid
   module Association
     module Referenced
@@ -12,6 +14,7 @@ class HasMany
         class Enumerable
           extend Forwardable
           include ::Enumerable
+          include Pluckable
 
           # The three main instance variables are collections of documents.
           #
@@ -374,6 +377,43 @@ def marshal_load(data)
             @_added, @_loaded, @_unloaded, @executed = data
           end
 
+          # Plucks the given field names from the documents in the target.
+          # If the collection has been loaded, it plucks from the loaded
+          # documents; otherwise, it plucks from the unloaded criteria.
+          # Regardless, it also plucks from any added documents.
+          #
+          # @param [ Symbol... ] *fields The field names to pluck.
+          #
+          # @return [ Array | Array<Array> ] The array of field values. If
+          #   multiple fields are given, an array of arrays is returned.
+          def pluck(*keys)
+            [].tap do |results|
+              if _loaded? || _added.any?
+                klass = @_association.klass
+                prepared = prepare_pluck(keys, document_class: klass)
+              end
+
+              if _loaded?
+                docs = _loaded.values.map { |v| BSON::Document.new(v.attributes) }
+                results.concat pluck_from_documents(docs, prepared[:field_names], document_class: klass)
+              elsif _unloaded
+                criteria = if _added.any?
+                  ids_to_exclude = _added.keys
+                  _unloaded.not(:_id.in => ids_to_exclude)
+                else
+                  _unloaded
+                end
+
+                results.concat criteria.pluck(*keys)
+              end
+
+              if _added.any?
+                docs = _added.values.map { |v| BSON::Document.new(v.attributes) }
+                results.concat pluck_from_documents(docs, prepared[:field_names], document_class: klass)
+              end
+            end
+          end
+
           # Reset the enumerable back to its persisted state.
           #
           # @example Reset the enumerable.

lib/mongoid/association/referenced/has_many/proxy.rb

@@ -33,7 +33,7 @@ def embedded?
 
           extend ClassMethods
 
-          def_delegator :criteria, :count
+          def_delegators :criteria, :count
           def_delegators :_target, :first, :in_memory, :last, :reset, :uniq
 
           # Instantiate a new references_many association. Will set the foreign key
@@ -281,6 +281,22 @@ def nullify
 
           alias nullify_all nullify
 
+          # Plucks the given field names from the documents in the
+          # association. It is safe to use whether the association is
+          # loaded or not, and whether there are unsaved documents in the
+          # association or not.
+          #
+          # @example Pluck the titles of all posts.
+          #   person.posts.pluck(:title)
+          #
+          # @param [ Symbol... ] *fields The field names to pluck.
+          #
+          # @return [ Array | Array<Array> ] The array of field values. If
+          #   multiple fields are given, an array of arrays is returned.
+          def pluck(*fields)
+            _target.pluck(*fields)
+          end
+
           # Clear the association. Will delete the documents from the db if they are
           # already persisted.
           #

lib/mongoid/contextual/mongo.rb

@@ -2,6 +2,7 @@
 # rubocop:todo all
 
 require 'mongoid/atomic_update_preparer'
+require 'mongoid/pluckable'
 require "mongoid/contextual/mongo/documents_loader"
 require "mongoid/contextual/atomic"
 require "mongoid/contextual/aggregable/mongo"
@@ -22,6 +23,7 @@ class Mongo
       include Atomic
       include Association::EagerLoadable
       include Queryable
+      include Pluckable
 
       # Options constant.
       OPTIONS = [ :hint,
@@ -331,23 +333,12 @@ def map_reduce(map, reduce)
       #   in the array will be a single value. Otherwise, each
       #   result in the array will be an array of values.
       def pluck(*fields)
-        # Multiple fields can map to the same field name. For example, plucking
-        # a field and its _translations field map to the same field in the database.
-        # because of this, we need to keep track of the fields requested.
-        normalized_field_names = []
-        normalized_select = fields.inject({}) do |hash, f|
-          db_fn = klass.database_field_name(f)
-          normalized_field_names.push(db_fn)
-          hash[klass.cleanse_localized_field_names(f)] = true
-          hash
-        end
-
-        view.projection(normalized_select).reduce([]) do |plucked, doc|
-          values = normalized_field_names.map do |n|
-            extract_value(doc, n)
-          end
-          plucked << (values.size == 1 ? values.first : values)
-        end
+        # Multiple fields can map to the same field name. For example,
+        # plucking a field and its _translations field map to the same
+        # field in the database. because of this, we need to prepare the
+        # projection specifically.
+        prep = prepare_pluck(fields, prepare_projection: true)
+        pluck_from_documents(view.projection(prep[:projection]), prep[:field_names])
       end
 
       # Pick the single field values from the database.
@@ -903,78 +894,6 @@ def acknowledged_write?
         collection.write_concern.nil? || collection.write_concern.acknowledged?
       end
 
-      # Fetch the element from the given hash and demongoize it using the
-      # given field. If the obj is an array, map over it and call this method
-      # on all of its elements.
-      #
-      # @param [ Hash | Array<Hash> ] obj The hash or array of hashes to fetch from.
-      # @param [ String ] meth The key to fetch from the hash.
-      # @param [ Field ] field The field to use for demongoization.
-      #
-      # @return [ Object ] The demongoized value.
-      #
-      # @api private
-      def fetch_and_demongoize(obj, meth, field)
-        if obj.is_a?(Array)
-          obj.map { |doc| fetch_and_demongoize(doc, meth, field) }
-        else
-          res = obj.try(:fetch, meth, nil)
-          field ? field.demongoize(res) : res.class.demongoize(res)
-        end
-      end
-
-      # Extracts the value for the given field name from the given attribute
-      # hash.
-      #
-      # @param [ Hash ] attrs The attributes hash.
-      # @param [ String ] field_name The name of the field to extract.
-      #
-      # @param [ Object ] The value for the given field name
-      def extract_value(attrs, field_name)
-        i = 1
-        num_meths = field_name.count('.') + 1
-        curr = attrs.dup
-
-        klass.traverse_association_tree(field_name) do |meth, obj, is_field|
-          field = obj if is_field
-          is_translation = false
-          # If no association or field was found, check if the meth is an
-          # _translations field.
-          if obj.nil? & tr = meth.match(/(.*)_translations\z/)&.captures&.first
-            is_translation = true
-            meth = tr
-          end
-
-          # 1. If curr is an array fetch from all elements in the array.
-          # 2. If the field is localized, and is not an _translations field
-          #    (_translations fields don't show up in the fields hash).
-          #    - If this is the end of the methods, return the translation for
-          #      the current locale.
-          #    - Otherwise, return the whole translations hash so the next method
-          #      can select the language it wants.
-          # 3. If the meth is an _translations field, do not demongoize the
-          #    value so the full hash is returned.
-          # 4. Otherwise, fetch and demongoize the value for the key meth.
-          curr = if curr.is_a? Array
-            res = fetch_and_demongoize(curr, meth, field)
-            res.empty? ? nil : res
-          elsif !is_translation && field&.localized?
-            if i < num_meths
-              curr.try(:fetch, meth, nil)
-            else
-              fetch_and_demongoize(curr, meth, field)
-            end
-          elsif is_translation
-            curr.try(:fetch, meth, nil)
-          else
-            fetch_and_demongoize(curr, meth, field)
-          end
-
-          i += 1
-        end
-        curr
-      end
-
       # Recursively demongoize the given value. This method recursively traverses
       # the class tree to find the correct field to use to demongoize the value.
       #

lib/mongoid/pluckable.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module Mongoid
+  # Provides shared behavior for any document with "pluck" functionality.
+  #
+  # @api private
+  module Pluckable
+    extend ActiveSupport::Concern
+
+    private
+
+      # Prepares the field names for plucking by normalizing them to their
+      # database field names. Also prepares a projection hash if requested.
+      def prepare_pluck(field_names, document_class: klass, prepare_projection: false)
+        normalized_field_names = []
+        projection = {}
+
+        field_names.each do |f|
+          db_fn = document_class.database_field_name(f)
+          normalized_field_names.push(db_fn)
+
+          if prepare_projection
+            cleaned_name = document_class.cleanse_localized_field_names(f)
+            canonical_name = document_class.database_field_name(cleaned_name)
+            projection[canonical_name] = true
+          end
+        end
+
+        { field_names: normalized_field_names, projection: projection }
+      end
+
+      # Plucks the given field names from the given documents.
+      def pluck_from_documents(documents, field_names, document_class: klass)
+        documents.reduce([]) do |plucked, doc|
+          values = field_names.map { |name| extract_value(doc, name.to_s, document_class) }
+          plucked << (values.size == 1 ? values.first : values)
+        end
+      end
+
+      # Fetch the element from the given hash and demongoize it using the
+      # given field. If the obj is an array, map over it and call this method
+      # on all of its elements.
+      #
+      # @param [ Hash | Array<Hash> ] obj The hash or array of hashes to fetch from.
+      # @param [ String ] key The key to fetch from the hash.
+      # @param [ Field ] field The field to use for demongoization.
+      #
+      # @return [ Object ] The demongoized value.
+      def fetch_and_demongoize(obj, key, field)
+        if obj.is_a?(Array)
+          obj.map { |doc| fetch_and_demongoize(doc, key, field) }
+        else
+          value = obj.try(:fetch, key, nil)
+          field ? field.demongoize(value) : value.class.demongoize(value)
+        end
+      end
+
+      # Extracts the value for the given field name from the given attribute
+      # hash.
+      #
+      # @param [ Hash ] attrs The attributes hash.
+      # @param [ String ] field_name The name of the field to extract.
+      #
+      # @param [ Object ] The value for the given field name
+      def extract_value(attrs, field_name, document_class)
+        i = 1
+        num_meths = field_name.count('.') + 1
+        curr = attrs.dup
+
+        document_class.traverse_association_tree(field_name) do |meth, obj, is_field|
+          field = obj if is_field
+
+          # use the correct document class to check for localized fields on
+          # embedded documents.
+          document_class = obj.klass if obj.respond_to?(:klass)
+
+          is_translation = false
+          # If no association or field was found, check if the meth is an
+          # _translations field.
+          if obj.nil? & tr = meth.match(/(.*)_translations\z/)&.captures&.first
+            is_translation = true
+            meth = document_class.database_field_name(tr)
+          end
+
+          # 1. If curr is an array fetch from all elements in the array.
+          # 2. If the field is localized, and is not an _translations field
+          #    (_translations fields don't show up in the fields hash).
+          #    - If this is the end of the methods, return the translation for
+          #      the current locale.
+          #    - Otherwise, return the whole translations hash so the next method
+          #      can select the language it wants.
+          # 3. If the meth is an _translations field, do not demongoize the
+          #    value so the full hash is returned.
+          # 4. Otherwise, fetch and demongoize the value for the key meth.
+          curr = if curr.is_a? Array
+            res = fetch_and_demongoize(curr, meth, field)
+            res.empty? ? nil : res
+          elsif !is_translation && field&.localized?
+            if i < num_meths
+              curr.try(:fetch, meth, nil)
+            else
+              fetch_and_demongoize(curr, meth, field)
+            end
+          elsif is_translation
+            curr.try(:fetch, meth, nil)
+          else
+            fetch_and_demongoize(curr, meth, field)
+          end
+
+          i += 1
+        end
+        curr
+      end
+
+  end
+end