Fork of d11wtq's DataMapper EagerLoader plugin

dm_eager_loader.rb

module DataMapper
  class Collection
    #
    # EagerLoader takes a QueryPath object and loads all relationships
    # referenced in the path, into an existing Collection.
    #
    # Using eager-loading allows you to optimize out the classic "n+1"
    # query problem, when you intend to iterate over several arbitrarily deeply
    # nested relationships.  One query per relationship is executed,
    # vs one query per record.
    #
    class EagerLoader
      #
      # Initialize the EagerLoader to pre-load all relationships as deep as
      # `query_path` into `collection`.
      #
      # @param [Collection] collection
      #   the source collection to load related resources into
      #
      # @param [QueryPath] query_path
      #   a valid QueryPath for the target model in the collection
      #
      def initialize(collection, query_path)
        @collection = collection
        @query_path = query_path
      end

      #
      # Perform eager loading immediately and return the collection.
      #
      # The number of queries executed is identical to the number of
      # relationships in the query path. This method should not be invoked
      # before query-building is complete (e.g. before limits have been applied)
      #
      # @return [Collection]
      #   the source collection resources were loaded into
      #
      def load
        scope = @collection

        @query_path.relationships.each do |relation|
          next_scope = relation.target_model.all(target_conditions(scope, relation))
          load_into_collection(scope, next_scope, relation)
          scope = next_scope
        end

        @collection
      end

      #
      # Map target key names to source key values, to create valid query
      # conditions for finding all related resources.
      #
      def target_conditions(collection, relationship)
        conditions = {}

        collection.each do |resource|
          case relationship
          when DataMapper::Associations::OneToMany::Relationship
            keys   = relationship.child_key.collect(&:name)
            ids    = relationship.parent_key.get(resource)

          when DataMapper::Associations::ManyToOne::Relationship
            keys   = relationship.parent_key.collect(&:name)
            ids    = relationship.child_key.get(resource)
          else
            raise(NotImplementedError,"#{relationship.class} not supported")
          end

          Hash[keys.zip(ids)].each do |key,id|
            (conditions[key] ||= []) << id
          end
        end

        conditions
      end

      #
      # Given a set of resources and the relationship from which they derive,
      # map them, referenced by the target key.
      #
      def key_mappings(collection, relationship)
        map = {}

        collection.each do |resource|
          case relationship
          when DataMapper::Associations::OneToMany::Relationship
            foreign_key = relationship.child_key.get(resource)

            (map[foreign_key] ||= []) << resource
          when DataMapper::Associations::ManyToOne::Relationship
            foreign_key = relationship.parent_key.get(resource)

            map[foreign_key] = resource
          else
            raise(NotImplementedError,"#{relationship.class} not supported")
          end
        end

        map
      end

      #
      # For each of the pre-loaded resources in +related_resources+,
      # re-establish their relationships in `collection`.
      #
      def load_into_collection(collection, related_resources, relationship)
        map = key_mappings(related_resources, relationship)

        collection.each do |resource|
          foreign_key = case relationship
                        when DataMapper::Associations::OneToMany::Relationship
                          relationship.child_key.get(resource)
                        when DataMapper::Associations::ManyToOne::Relationship
                          relationship.parent_key.get(resource)
                        else
                          raise(NotImplementedError,"#{relationship.class} not supported")
                        end

          if map.key?(foreign_key)
            relationship.set!(resource,map[foreign_key])
          end
        end
      end
    end

    module EagerLoading
      def eager_load(query_path)
        EagerLoader.new(self, query_path).load
      end
    end

    include EagerLoading
  end
end