mongoid 9.0.11 → 9.1.1

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

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # Something an eager load contributes to the pipeline, in the role it plays
+      # while the pipeline is built. A root is asked to contribute and the whole
+      # tree follows by recursion. AssociationInclusion stands for a single
+      # association; DiscriminatedInclusion stands for a name several subclasses
+      # share.
+      #
+      # @api private
+      class Inclusion
+        # Add this inclusion's stages to the destination.
+        #
+        # @param [ Array<Hash> ] destination The pipeline (or sub-pipeline) the
+        #   stages are appended to.
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded path
+        #   accumulated from the ancestors above this inclusion (empty at the top).
+        def contribute(destination, chain)
+          raise NotImplementedError
+        end
+      end
+
+      # An inclusion that stands for a single association. The LookupPipeline holds
+      # the stage-building helpers the kinds lean on, and a node carries its own
+      # children, so the pipeline is built by recursion from the roots downward.
+      #
+      # @api private
+      class AssociationInclusion < Inclusion
+        class << self
+          # Builds the right kind of inclusion for the association. Each subclass
+          # decides whether it handles it (.for?); exactly one does.
+          #
+          # @param [ Mongoid::Association::Relatable ] association The inclusion.
+          # @param [ LookupPipeline ] pipeline The pipeline being built.
+          # @param [ Array<Inclusion> ] children The inclusions nested under it.
+          #
+          # @return [ AssociationInclusion ] The matching kind of inclusion.
+          def for(association, pipeline, children)
+            subclasses.find { |kind| kind.for?(association) }.new(association, pipeline, children)
+          end
+
+          # Whether this kind handles the given association.
+          #
+          # @param [ Mongoid::Association::Relatable ] association The inclusion.
+          #
+          # @return [ true | false ] Whether it handles it.
+          def for?(association)
+            raise NotImplementedError
+          end
+        end
+
+        # @return [ Mongoid::Association::Relatable ] The association this stands for.
+        attr_reader :association
+
+        def initialize(association, pipeline, children)
+          super()
+          @association = association
+          @pipeline = pipeline
+          @children = children
+        end
+      end
+
+      # A referenced inclusion: contributes a $lookup whose sub-pipeline holds its
+      # own children. When it lives inside an embedded document (a non-empty
+      # chain), the $lookup is distributed onto that embedded path instead of
+      # standing at the top level.
+      #
+      # For a has_many :albums it contributes:
+      #
+      #   { '$lookup' => {
+      #     'from' => 'albums',
+      #     'localField' => '_id',        # the band's _id...
+      #     'foreignField' => 'band_id',  # ...matched against each album's band_id
+      #     'as' => 'albums',             # matches are written to this field
+      #     'pipeline' => [
+      #       { '$sort' => {
+      #         '_id' => 1
+      #       } },
+      #       <children>
+      #     ]
+      #   } }
+      #
+      # @api private
+      class JoinedInclusion < AssociationInclusion
+        class << self
+          # The default kind: a referenced, non-polymorphic association, i.e. the
+          # one no sibling kind claims.
+          #
+          # @param [ Mongoid::Association::Relatable ] association The inclusion.
+          #
+          # @return [ true | false ] Whether it handles it.
+          def for?(association)
+            (superclass.subclasses - [ self ]).none? { |kind| kind.for?(association) }
+          end
+        end
+
+        # Append the $lookup, with the children in its sub-pipeline, to the
+        # destination; or distribute it onto the embedded path when nested in one.
+        #
+        # @param [ Array<Hash> ] destination The pipeline (or sub-pipeline) the
+        #   stages are appended to.
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded path
+        #   accumulated from the ancestors above this inclusion (empty at the top).
+        def contribute(destination, chain)
+          stage = @pipeline.lookup_stage_for(@association)
+          @children.each { |child| child.contribute(stage['$lookup']['pipeline'], []) }
+
+          if chain.empty?
+            destination << stage
+          else
+            destination.concat(@pipeline.distribute(@association, chain, stage))
+          end
+        end
+      end
+
+      # An embedded inclusion: it rides inside its own document, so it adds no
+      # stage of its own. Its children contribute to the same destination, with
+      # this document appended to their embedded path.
+      #
+      # For Computer.eager_load(port: :device) the :port inclusion emits nothing;
+      # it hands the path [ :port ] to :device, which EmbeddedDistributor then
+      # turns into stages.
+      #
+      # @api private
+      class EmbeddedInclusion < AssociationInclusion
+        class << self
+          # @param [ Mongoid::Association::Relatable ] association The inclusion.
+          #
+          # @return [ true | false ] Whether the association is embedded.
+          def for?(association)
+            association.embedded?
+          end
+        end
+
+        # Add no stage of its own; hand this document down the embedded path so the
+        # children distribute onto it.
+        #
+        # @param [ Array<Hash> ] destination The pipeline (or sub-pipeline) the
+        #   stages are appended to.
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded path
+        #   accumulated from the ancestors above this inclusion (empty at the top).
+        def contribute(destination, chain)
+          @children.each { |child| child.contribute(destination, chain + [ @association ]) }
+        end
+      end
+
+      # A polymorphic inclusion: its target collection varies per document, so it
+      # can't be a $lookup. It adds nothing here; PolymorphicPreloader resolves it
+      # after the roots are materialized.
+      #
+      # @api private
+      class DeferredInclusion < AssociationInclusion
+        class << self
+          # @param [ Mongoid::Association::Relatable ] association The inclusion.
+          #
+          # @return [ true | false ] Whether the association is polymorphic.
+          def for?(association)
+            association.polymorphic?
+          end
+        end
+
+        # Add nothing; PolymorphicPreloader resolves the association after the
+        # roots are materialized.
+        #
+        # @param [ Array<Hash> ] destination The pipeline (unused).
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded path (unused).
+        def contribute(destination, chain); end
+      end
+    end
+  end
+end

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