Eager Loading DM associations

eager_loading.rb

# Usage:
#
# customers = Customer.all
# customers.eager_load(Customer.orders.line_items.item, Customer.address)
#

module DataMapper
  module EagerLoading

    def eager_load(*query_paths)
      @eager_loaded_relationships ||= {}

      query_paths.each do |query_path|

        scope = self

        query_path.relationships.each do |relation|
          raise "EagerLoading doesn't work for composite keys yet" unless relation.target_key.count == 1

          if @eager_loaded_relationships[relation]
            scope = @eager_loaded_relationships[relation]
            next
          end

          # for each level in the query path, collect all the resources referencing keys at the current scope
          next_scope = relation.target_model.all(relation.target_key.first.name => scope.collect(&:"#{relation.source_key.first.name}").compact)

          # create a map between the target keys and the resources that exist for them
          links = next_scope.inject({}) do |map, resource|
            map.merge(relation.target_key.first.get(resource) => [resource]) { |k, v1, v2| v1 + v2 }
          end

          # now pre-load those from the map
          scope.each do |parent|
            if links.key?(relation.source_key.first.get(parent))
              if relation.is_a?(DataMapper::Associations::ManyToOne::Relationship)
                raise "Found more than one target element for a DataMapper::Associations::ManyToOne::Relationship (#{relation.inspect})" if links[relation.source_key.first.get(parent)].count > 1
                parent.instance_variable_set(:"@#{relation.name}", links[relation.source_key.first.get(parent)].first)
              else
                parent.instance_variable_set(:"@#{relation.name}", links[relation.source_key.first.get(parent)])
              end
            end
          end

          # and step into the next nesting level
          scope = next_scope
          @eager_loaded_relationships[relation] = scope
        end

      end

      self
    end

  end
end

DataMapper::Collection.send(:include, DataMapper::EagerLoading)