mongoid 9.1.0 → 9.1.1

data/lib/mongoid/association/eager_load/inclusion_tree.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+require 'mongoid/association/eager_load/discriminated_inclusion'
+require 'mongoid/association/eager_load/inclusion'
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # The tree of nested inclusions an eager load asks to load. Built from the
+      # criteria's inclusions, it contributes each root node's stages to the
+      # pipeline, with each node already carrying its own children.
+      #
+      # Each root branch is built from its own copy of the inclusions, and an
+      # inclusion is removed as it is placed, so it lands once per branch even if
+      # more than one parent in that branch points at it, and a circular chain of
+      # inclusions can't loop forever.
+      #
+      # @api private
+      class InclusionTree
+        class << self
+          # Builds the tree for the criteria's inclusions.
+          #
+          # @param [ Array<Mongoid::Association::Relatable> ] inclusions The inclusions.
+          # @param [ LookupPipeline ] pipeline The pipeline being built.
+          #
+          # @return [ InclusionTree ] The tree.
+          def from(inclusions, pipeline)
+            new(inclusions, pipeline, inclusions.to_h { |association| [ association.name, association ] })
+          end
+
+          private :new
+        end
+
+        def initialize(inclusions, pipeline, by_name)
+          @inclusions = inclusions
+          @pipeline = pipeline
+          @by_name = by_name
+        end
+
+        # Contribute each root inclusion's stages to the pipeline. Each root carries
+        # its own children, so the whole tree is appended by recursion from the
+        # roots downward.
+        #
+        # @param [ Array<Hash> ] destination The pipeline the stages are appended to.
+        def contribute_to(destination)
+          roots.each { |root| root.contribute(destination, []) }
+        end
+
+        private
+
+        # A name that more than one subclass defines (with different targets) can't
+        # share one $lookup field, so its nodes -- each carrying its own children --
+        # are grouped and routed by the discriminator instead of becoming separate,
+        # overwriting roots.
+        def roots
+          top_level.group_by(&:name).map do |_name, associations|
+            nodes = associations.map { |association| node(association, @inclusions.dup) }
+            nodes.one? ? nodes.first : DiscriminatedInclusion.new(nodes)
+          end
+        end
+
+        # The inclusions that no other inclusion is the parent of.
+        def top_level
+          @inclusions.reject do |association|
+            association.parent_inclusions.any? { |name| @by_name.key?(name) }
+          end
+        end
+
+        def node(association, available)
+          children = take_children(association, available).map { |child| node(child, available) }
+          AssociationInclusion.for(association, @pipeline, children)
+        end
+
+        # The still-available inclusions parented to +association+, removed as they
+        # are taken so each lands once on this branch. A child belongs here when it
+        # names this association as its parent and its owner shares the target's
+        # class hierarchy, which tells apart children of two unrelated subclasses
+        # that share an association name.
+        def take_children(association, available)
+          children = available.select do |candidate|
+            candidate.parent_inclusions.include?(association.name) &&
+              same_hierarchy?(association.klass, candidate.inverse_class)
+          end
+          available.reject! { |candidate| children.include?(candidate) }
+          children
+        end
+
+        # Whether two classes belong to the same inheritance chain (one is the
+        # other, an ancestor of it, or a descendant of it).
+        def same_hierarchy?(one, other)
+          one <= other || other <= one
+        end
+      end
+    end
+  end
+end

data/lib/mongoid/association/eager_load/lookup_pipeline.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require 'mongoid/association/eager_load/embedded_distributor'
+require 'mongoid/association/eager_load/inclusion_tree'
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # Builds the aggregation pipeline that eager-loads a criteria's inclusions
+      # with $lookup.
+      #
+      # It starts with the criteria's own match/sort/skip/limit, then lets each
+      # root of the inclusion tree contribute its stages. This object owns the
+      # stage-building helpers; how each inclusion contributes is the inclusion's
+      # own business (see Inclusion).
+      #
+      # For Band.eager_load(albums: :tracks) the result is roughly:
+      #
+      #   [ <criteria match / sort / skip / limit>,
+      #     { '$lookup' => {                 # JoinedInclusion(:albums)
+      #       'from' => 'albums',
+      #       'localField' => '_id',
+      #       'foreignField' => 'band_id',
+      #       'as' => 'albums',
+      #       'pipeline' => [
+      #         { '$sort' => {
+      #           '_id' => 1
+      #         } },
+      #         { '$lookup' => {             # JoinedInclusion(:tracks), nested
+      #           'as' => 'tracks',
+      #           'pipeline' => [
+      #             { '$sort' => {
+      #               '_id' => 1
+      #             } }
+      #           ]
+      #         } }
+      #       ]
+      #     } } ]
+      #
+      # @api private
+      class LookupPipeline
+        def initialize(criteria)
+          @criteria = criteria
+        end
+
+        # @return [ Array<Hash> ] The aggregation pipeline stages.
+        def stages
+          pipeline = @criteria.selector.to_pipeline
+          pipeline.concat(@criteria.options.to_pipeline_for_lookup)
+          InclusionTree.from(@criteria.inclusions, self).contribute_to(pipeline)
+          pipeline
+        end
+
+        # The $lookup stage for a referenced inclusion: its key fields, a
+        # discriminator match when the target shares its collection with sibling
+        # subclasses, and an order. Children are added by the inclusion itself.
+        #
+        # @param [ Mongoid::Association::Relatable ] association The inclusion.
+        #
+        # @return [ Hash ] The $lookup stage.
+        def lookup_stage_for(association)
+          local_field, foreign_field = lookup_fields(association)
+          sub_pipeline = []
+          sub_pipeline << discriminator_match(association) if association.klass.hereditary?
+          sub_pipeline << order(association)
+          { '$lookup' => {
+            'from' => association.klass.collection.name,
+            'localField' => local_field,
+            'foreignField' => foreign_field,
+            'as' => association.name.to_s,
+            'pipeline' => sub_pipeline
+          } }
+        end
+
+        # Builds the stages that distribute a referenced inclusion living inside an
+        # embedded document onto that document.
+        #
+        # @param [ Mongoid::Association::Relatable ] association The inclusion.
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded path.
+        # @param [ Hash ] lookup_stage The $lookup stage for the association.
+        #
+        # @return [ Array<Hash> ] The stages to append.
+        def distribute(association, chain, lookup_stage)
+          EmbeddedDistributor.for(association: association, chain: chain, lookup_stage: lookup_stage).stages
+        end
+
+        private
+
+        # When the association stores the foreign key on the current document
+        # (belongs_to, has_and_belongs_to_many) the local field is that key; for
+        # the others (has_many, has_one) the key is on the related document.
+        def lookup_fields(association)
+          if association.stores_foreign_key?
+            [ association.foreign_key, association.primary_key ]
+          else
+            [ association.primary_key, association.foreign_key ]
+          end
+        end
+
+        def discriminator_match(association)
+          target = association.klass
+          { '$match' => {
+            target.discriminator_key => { '$in' => target._types }
+          } }
+        end
+
+        def order(association)
+          unless association.order
+            return { '$sort' => {
+              '_id' => 1
+            } }
+          end
+
+          ordering = association.order
+          ordering = { ordering => 1 } unless ordering.is_a?(Hash)
+          { '$sort' => ordering }
+        end
+      end
+    end
+  end
+end

data/lib/mongoid/association/eager_load/polymorphic_preloader.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'mongoid/association/eager_load/polymorphic_targets'
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # Resolves a polymorphic belongs_to onto already-materialized root documents.
+      #
+      # A polymorphic belongs_to can't be expressed as a $lookup: its target
+      # collection varies per document. So once the roots are materialized, the
+      # foreign keys are grouped by type, PolymorphicTargets resolves the documents
+      # for those keys, and the result is set on each document.
+      #
+      # @api private
+      class PolymorphicPreloader
+        def initialize(association, root_class)
+          @association = association
+          @root_class = root_class
+        end
+
+        # Resolve and assign the polymorphic target on each of the documents.
+        #
+        # @param [ Array<Mongoid::Document> ] documents The materialized root documents.
+        def preload_into(documents)
+          targets = PolymorphicTargets.for(@association, keys_by_type(documents), @root_class)
+          assign(documents, targets)
+        end
+
+        private
+
+        # The foreign keys on the documents grouped by polymorphic type,
+        # e.g. { "Printer" => [ id1 ], "Scanner" => [ id2 ] }.
+        def keys_by_type(documents)
+          documents.each_with_object({}) do |document, grouped|
+            type, key = reference_on(document)
+            (grouped[type] ||= []) << key if type && key
+          end
+        end
+
+        def assign(documents, targets)
+          documents.each do |document|
+            type, key = reference_on(document)
+            target = targets.dig(type, key) if type && key
+            document.set_relation(@association.name, target)
+          end
+        end
+
+        # The [ type, key ] reference stored on the document for this association,
+        # or an empty pair when the document holds no such reference.
+        def reference_on(document)
+          type_field = @association.inverse_type
+          key_field = @association.foreign_key
+          return [] unless document.respond_to?(type_field) && document.respond_to?(key_field)
+
+          [ document.send(type_field), document.send(key_field) ]
+        end
+      end
+    end
+  end
+end

data/lib/mongoid/association/eager_load/polymorphic_targets.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # The targets of a polymorphic belongs_to, indexed as
+      # { type => { primary_key => document } }. Each subclass reaches the types that
+      # live in one place (the root's database or elsewhere); .for resolves the whole
+      # set, routing each type to the subclass that can reach it.
+      #
+      # @api private
+      class PolymorphicTargets
+        class << self
+          # Resolve every polymorphic target for the foreign keys grouped by type.
+          # The types whose documents share the root's database are fetched together
+          # in one $facet; those living elsewhere are read through their own models.
+          #
+          # @param [ Mongoid::Association::Relatable ] association The polymorphic inclusion.
+          # @param [ Hash ] keys_by_type The foreign keys grouped by type.
+          # @param [ Class ] root_class The class being queried.
+          #
+          # @return [ Hash ] The targets, as { type => { primary_key => document } }.
+          def for(association, keys_by_type, root_class)
+            here, elsewhere = keys_by_type.partition do |type, _keys|
+              in_root_database?(association, type, root_class)
+            end
+            same_database = SameDatabaseTargets.new(association, here.to_h, root_class)
+            other_databases = OtherDatabaseTargets.new(association, elsewhere.to_h)
+            same_database.fetch.merge(other_databases.fetch)
+          end
+
+          private
+
+          # Whether the type's model shares the root's database (and client): exactly
+          # what a $lookup from the root collection can reach.
+          def in_root_database?(association, type, root_class)
+            model = association.resolver.model_for(type)
+            model.client_name == root_class.client_name &&
+              model.database_name == root_class.database_name
+          end
+        end
+
+        def initialize(association, keys_by_type)
+          @association = association
+          @keys_by_type = keys_by_type
+        end
+
+        # @return [ Hash ] The targets, as { type => { primary_key => document } }.
+        def fetch
+          raise NotImplementedError
+        end
+
+        private
+
+        def primary_key
+          @association.primary_key
+        end
+
+        def model_for(type)
+          @association.resolver.model_for(type)
+        end
+
+        # The raw documents instantiated and indexed by their primary key.
+        def indexed(documents, model)
+          documents.map { |document| Factory.from_db(model, document) }
+                   .index_by { |document| document.send(primary_key) }
+        end
+      end
+
+      # Targets that live in the root's own database. A $lookup can reach them, so
+      # every type is fetched together in one $facet aggregation against the root
+      # collection.
+      #
+      # For { 'Printer' => [ id1 ], 'Scanner' => [ id2 ] } it runs:
+      #
+      #   [
+      #     { '$limit' => 1 },   # one input doc, so each facet branch runs once
+      #     { '$facet' => {      # run one $lookup per type within a single query
+      #       'Printer' => [ { '$lookup' => { 'from' => 'printers', ... } }, ... ],
+      #       'Scanner' => [ { '$lookup' => { 'from' => 'scanners', ... } }, ... ]
+      #     } }
+      #   ]
+      #
+      # @api private
+      class SameDatabaseTargets < PolymorphicTargets
+        def initialize(association, keys_by_type, root_class)
+          super(association, keys_by_type)
+          @root_class = root_class
+        end
+
+        # @return [ Hash ] The targets, as { type => { primary_key => document } }.
+        def fetch
+          return {} if @keys_by_type.empty?
+
+          aggregated = @root_class.collection.aggregate([ { '$limit' => 1 }, { '$facet' => facets } ]).first
+          aggregated.to_h do |type, branch|
+            # $limit => 1 makes each branch yield a single wrapper holding the matches.
+            [ type, indexed(branch.first['matches'], model_for(type)) ]
+          end
+        end
+
+        private
+
+        def facets
+          @keys_by_type.to_h do |type, keys|
+            [ type, branch_for(model_for(type).collection.name, keys) ]
+          end
+        end
+
+        # One $facet branch: the documents in +collection_name+ whose primary key
+        # is among +keys+, exposed under "matches".
+        def branch_for(collection_name, keys)
+          [
+            { '$lookup' => {
+              'from' => collection_name,
+              'pipeline' => [
+                { '$match' => {
+                  primary_key => { '$in' => keys.uniq }
+                } }
+              ],
+              'as' => 'matches'
+            } },
+            { '$project' => {
+              '_id' => 0,
+              'matches' => 1
+            } }
+          ]
+        end
+      end
+
+      # Targets kept in another database (or cluster), which a $lookup cannot
+      # reach. Each type is read directly through its own model, which connects
+      # with that model's client.
+      #
+      # For { 'Scanner' => [ id2 ] } it runs, on the Scanner model's own client:
+      #   scanners.find('_id' => { '$in' => [ id2 ] })
+      #
+      # @api private
+      class OtherDatabaseTargets < PolymorphicTargets
+        # @return [ Hash ] The targets, as { type => { primary_key => document } }.
+        def fetch
+          @keys_by_type.to_h do |type, keys|
+            model = model_for(type)
+            documents = model.collection.find(primary_key => { '$in' => keys.uniq })
+            [ type, indexed(documents, model) ]
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mongoid/association/eager_loadable.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require 'mongoid/association/eager'
+require 'mongoid/association/eager_load/embedded_distributor'
+require 'mongoid/association/eager_load/polymorphic_preloader'
+require 'mongoid/association/eager_load/lookup_pipeline'
 
 module Mongoid
   module Association
@@ -27,24 +30,49 @@ module Mongoid
 
       # Load the associations for the given documents using $lookup.
       #
-      # If any of the associated collections reside in a different cluster than
-      # the root class, falls back to the #includes behavior and logs a warning.
+      # If any of the associated collections reside in a different cluster or
+      # database than the root class, falls back to the #includes behavior and
+      # logs a warning.
       #
       # @return [ Array<Mongoid::Document> ] The given documents.
       def eager_load_with_lookup
-        offenders = cross_cluster_inclusions
+        offenders = inclusions_unreachable_by_lookup
         if offenders.any?
-          root_client = klass.client_name
-          offender_list = offenders.map { |a| "#{a.name} (#{a.klass.client_name})" }.join(', ')
+          offender_descriptions = offenders.map do |offender|
+            "#{offender.name} (client: #{offender.klass.client_name}, database: #{offender.klass.database_name})"
+          end
           Mongoid.logger.warn(
             'eager_load cannot use $lookup aggregation because the following associations ' \
-            "reside in a different cluster than #{klass} (client: #{root_client}): " \
-            "#{offender_list}. Falling back to #includes behavior."
+            "reside in a different cluster or database than #{klass} " \
+            "(client: #{klass.client_name}, database: #{klass.database_name}): " \
+            "#{offender_descriptions.join(', ')}. Falling back to #includes behavior."
+          )
+          return eager_load(docs_for_lookup_fallback)
+        end
+
+        through_inclusions = criteria.inclusions.select do |association|
+          association.is_a?(Association::Referenced::HasOneThrough) ||
+            association.is_a?(Association::Referenced::HasManyThrough)
+        end
+
+        if through_inclusions.any?
+          through_names = through_inclusions.map { |association| ":#{association.name}" }
+          Mongoid.logger.warn(
+            "#{through_names.join(', ')} are :through associations and do not support " \
+            '$lookup-based eager loading. All inclusions for this query will be preloaded ' \
+            'using separate queries.'
           )
           return eager_load(docs_for_lookup_fallback)
         end
 
-        preload_for_lookup(criteria)
+        documents = preload_for_lookup(criteria)
+        # A polymorphic belongs_to cannot be expressed as a $lookup: its target
+        # collection varies per document. It is resolved after the roots are
+        # materialized, each inclusion by its own preloader.
+        criteria.inclusions.select(&:polymorphic?).each do |association|
+          EagerLoad::PolymorphicPreloader.new(association, klass).preload_into(documents)
+        end
+        documents
       end
 
       # Load the associations for the given documents. This will be done
@@ -82,39 +110,16 @@ module Mongoid
         end
       end
 
-      # Load the associations for the given documents. This will be done
-      # recursively to load the associations of the given documents'
-      # associated documents.
+      # Materialize the root documents with their inclusions eager-loaded by a
+      # single $lookup aggregation. The pipeline is built by LookupPipeline; the
+      # polymorphic inclusions it leaves out are resolved by the caller.
       #
-      # @param [ Array<Mongoid::Association::Relatable> ] associations
-      #   The associations to load.
-      # @param [ Array<Mongoid::Document> ] docs The documents.
+      # @param [ Mongoid::Criteria ] criteria The criteria to load.
+      #
+      # @return [ Array<Mongoid::Document> ] The materialized root documents.
       def preload_for_lookup(criteria)
-        assoc_map = criteria.inclusions.group_by(&:inverse_class_name)
-
-        # match first
-        pipeline = criteria.selector.to_pipeline
-        # then sort, skip, limit
-        pipeline.concat(criteria.options.to_pipeline_for_lookup)
-
-        # account for single-collection inheritance
-        root_class = klass.root_class
-
-        if assoc_map[klass.to_s]
-          assoc_map[klass.to_s].each do |assoc|
-            # Create a copy of the mapping for each top-level association to avoid mutation issues
-            pipeline << create_pipeline(assoc, assoc_map.dup)
-          end
-        end
-
-        if klass != root_class && assoc_map[root_class.to_s]
-          assoc_map[root_class.to_s].each do |assoc|
-            # Create a copy of the mapping for each top-level association to avoid mutation issues
-            pipeline << create_pipeline(assoc, assoc_map.dup)
-          end
-        end
-
-        Eager.new(criteria.inclusions, [], true, pipeline).run
+        pipeline = EagerLoad::LookupPipeline.new(criteria).stages
+        Eager.run(criteria.inclusions, [], true, pipeline)
       end
 
       private
@@ -128,71 +133,23 @@ module Mongoid
         raise NotImplementedError, "#{self.class} must implement #docs_for_lookup_fallback"
       end
 
-      # Returns the inclusions whose target class resides in a different cluster
-      # than the root class.
+      # Returns the inclusions whose target class can't be reached by a $lookup
+      # from the root class, which joins only within the same client and database.
       #
       # @return [ Array<Mongoid::Association::Relatable> ] The offending inclusions.
-      def cross_cluster_inclusions
-        root_client = klass.client_name
-        criteria.inclusions.reject { |assoc| assoc.klass.client_name == root_client }
-      end
-
-      public
-
-      def switch_local_and_foreign_fields?(association)
-        association.is_a?(Mongoid::Association::Referenced::BelongsTo) ||
-          association.is_a?(Mongoid::Association::Referenced::HasAndBelongsToMany)
-      end
-
-      def create_pipeline(current_assoc, mapping)
-        # Build nested pipeline for children and ordering
-        pipeline_stages = []
-
-        # For belongs_to and has_and_belongs_to_many, the foreign key is on the current document
-        # For has_many/has_one, the foreign key is on the related document
-        if switch_local_and_foreign_fields?(current_assoc)
-          local_field = current_assoc.foreign_key
-          foreign_field = current_assoc.primary_key
-        else
-          local_field = current_assoc.primary_key
-          foreign_field = current_assoc.foreign_key
+      def inclusions_unreachable_by_lookup
+        # Polymorphic associations have no single resolvable klass and are not
+        # loaded via $lookup, so they are never offenders.
+        criteria.inclusions.reject do |association|
+          association.polymorphic? || reachable_by_lookup?(association.klass)
         end
+      end
 
-        # Build the 'as' field with embedded path prefix if needed
-        as_field = current_assoc.name.to_s
-
-        stage = {
-          '$lookup' => {
-            'from' => current_assoc.klass.collection.name,
-            'localField' => local_field,
-            'foreignField' => foreign_field,
-            'as' => as_field
-          }
-        }
-
-        # Add ordering if defined on the association, or default to _id for consistent order
-        if current_assoc.order
-          sort_spec = current_assoc.order.is_a?(Hash) ? current_assoc.order : { current_assoc.order => 1 }
-          pipeline_stages << { '$sort' => sort_spec }
-        else
-          # Default to sorting by _id to maintain insertion order consistency
-          pipeline_stages << { '$sort' => { '_id' => 1 } }
-        end
-
-        # Add nested lookups for child associations
-        # Child associations don't need the embedded_path prefix since they're referenced from the looked-up document
-        # Remove this class from the mapping to prevent infinite loops with circular references
-        class_name = current_assoc.klass.to_s
-        if child_assocs = mapping.delete(class_name)
-          child_assocs.each do |child|
-            pipeline_stages << create_pipeline(child, mapping)
-          end
-        end
-
-        # Always add pipeline since we always have at least $sort
-        stage['$lookup']['pipeline'] = pipeline_stages
-
-        stage
+      # Whether a $lookup from a query on the root class can reach the model: it
+      # must live in the same client and database.
+      def reachable_by_lookup?(model)
+        model.client_name == klass.client_name &&
+          model.database_name == klass.database_name
       end
     end
   end

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

@@ -601,7 +601,12 @@ module Mongoid
           # @api private
           def update_attributes_hash
             if _target.empty?
-              _base.attributes.delete(_association.store_as)
+              # Only remove the key if it is absent or nil. If it is already
+              # an empty array, the caller explicitly set it to [] (e.g. via
+              # a raw write_attribute call), and that intent must be preserved
+              # so the empty array is persisted to the database.
+              stored = _association.store_as
+              _base.attributes.delete(stored) unless _base.attributes[stored] == []
             else
               _base.attributes.merge!(_association.store_as => _target.map(&:attributes))
             end

data/lib/mongoid/association/embedded/embeds_one.rb

@@ -21,6 +21,7 @@ module Mongoid
           as
           cascade_callbacks
           cyclic
+          fallback
           store_as
         ]
 

data/lib/mongoid/association/macros.rb

@@ -218,7 +218,19 @@ module Mongoid
         private
 
         def define_association!(macro_name, name, options = {}, &block)
-          Association::MACRO_MAPPING[macro_name].new(self, name, options, &block).tap do |assoc|
+          klass = if options[:through]
+                    Association::THROUGH_MACRO_MAPPING[macro_name] ||
+                      raise(
+                        Errors::InvalidRelationOption.new(
+                          self, name, :through,
+                          Association::THROUGH_MACRO_MAPPING.keys
+                        )
+                      )
+                  else
+                    Association::MACRO_MAPPING[macro_name]
+                  end
+
+          klass.new(self, name, options, &block).tap do |assoc|
             assoc.setup!
             self.relations = relations.merge(name => assoc)
             if assoc.embedded? && assoc.respond_to?(:store_as) && assoc.store_as != name