MONGOID-5900, MONGOID-5901 Fix #pluck on associations (mongodb#6067)

* use the specialized implementation of pluck

* use a custom pluck implementation

* new file

* incorporate Johnny's changes from mongodb#6044

* use same submodules as master

* Make changes suggested by copilot review

* rubocop appeasement

* fix broken refactoring

Author: jamis

.rubocop.yml

@@ -72,6 +72,15 @@ Metrics/ModuleLength:
 Metrics/MethodLength:
   Max: 20
 
+Metrics/PerceivedComplexity:
+  Max: 10
+
+Metrics/CyclomaticComplexity:
+  Max: 10
+
+Metrics/AbcSize:
+  Max: 20
+
 RSpec/BeforeAfterAll:
   Enabled: false
 

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_and_belongs_to_many/proxy.rb

@@ -53,8 +53,6 @@ def embedded?
           # @param [ Document... ] *args Any number of documents.
           #
           # @return [ Array<Document> ] The loaded docs.
-          #
-          # rubocop:disable Metrics/AbcSize
           def <<(*args)
             docs = args.flatten
             return concat(docs) if docs.size > 1
@@ -89,7 +87,6 @@ def <<(*args)
             end
             unsynced(_base, foreign_key) and self
           end
-          # rubocop:enable Metrics/AbcSize
 
           alias push <<
 
@@ -360,8 +357,6 @@ def clear_target_for_nullify
           #   in bulk
           # @param [ Array ] inserts the list of Hashes of attributes that will
           #   be inserted (corresponding to the ``docs`` list)
-          #
-          # rubocop:disable Metrics/AbcSize
           def append_document(doc, ids, docs, inserts)
             return unless doc
 
@@ -379,7 +374,6 @@ def append_document(doc, ids, docs, inserts)
               unsynced(_base, foreign_key)
             end
           end
-          # rubocop:enable Metrics/AbcSize
         end
       end
     end

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.
@@ -893,78 +884,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.
       #