Add :eager_loader association option, to specify code to be run when …

…eager loading

:eager_loader is powerful option that allows you to completely
specify how associations are to be eagerly loaded.  It is flexible
enough that it can be used to eagerly load association types that
Sequel doesn't even natively support, such as polymorphic
associations.  The eager loading code has been greatly simplified as
it now just uses the :eager_loader option, which has defaults that
are the same as the previous eager loading behavior.

:eager_loader should be a proc that takes 3 arguments, a key_hash,
an array of records, and a hash of dependent associations.  Since you
are given all of the records, you can do things like filter on
associations that are specified by multiple keys, or do multiple
queries depending on the content of the records (which would be
necessary for polymorphic associations).  Inside the :eager_loader
proc, you should get the related objects and populate the
associations for all objects in the array of records.  The hash
of dependent associations is available for you to cascade the eager
loading down multiple levels, but it is up to you to use it.  The
key_hash is a performance enhancement that is used by the default
code and is also available to you.  It is a hash with keys being
foreign/primary key symbols in the current table, and the values
being hashes where the key is foreign/primary key values and values
being arrays of current model objects having the foreign/primary key
value associated with the key.  This is hard to visualize, so I'll
give an example:

  album1 = Album.load(:id=>1, :artist_id=>2)
  album2 = Album.load(:id=>3, :artist_id=>2)
  Album.many_to_one :artist
  Album.one_to_many :tracks
  Album.eager(:band, :tracks).all
  # The key_hash provided to the :eager_loader proc would be:
  {:id=>{1=>[album1], 3=>[album2]}, :artist_id=>{2=>[album1, album2]}}

sequel/CHANGELOG

@@ -1,5 +1,7 @@
 === HEAD
 
+* Add :eager_loader association option, to specify code to be run when eager loading (jeremyevans)
+
 * Make :many_to_one associations support :dataset, :order, :limit association options, as well as block arguments (jeremyevans)
 
 * Add :dataset association option, which overrides the default base dataset to use (jeremyevans)

sequel/lib/sequel_model/associations.rb

@@ -95,13 +95,18 @@ def all_association_reflections
   #     For many_to_one associations, this is ignored unless this association is
   #     being eagerly loaded, as it doesn't save queries unless multiple objects
   #     can be loaded at once.
+  #   - :eager_block - If given, use the block instead of the default block when
+  #     eagerly loading.  To not use a block when eager loading (when one is used normally),
+  #     set to nil.
   #   - :eager_graph - The associations to eagerly load via EagerLoading#eager_graph when loading the associated object(s).
   #     For many_to_one associations, this is ignored unless this association is
   #     being eagerly loaded, as it doesn't save queries unless multiple objects
   #     can be loaded at once.
-  #   - :eager_block - If given, use the block instead of the default block when
-  #     eagerly loading.  To not use a block when eager loading (when one is used normally),
-  #     set to nil.
+  #   - :eager_loader - A proc to use to implement eager loading, overriding the default.  Takes three arguments,
+  #     a key hash (used solely to enhance performance), an array of records,
+  #     and a hash of dependent associations.  The associated records should
+  #     be queried from the database and the associations cache for each
+  #     record should be populated for this to work correctly.
   #   - :graph_block - The block to pass to join_table when eagerly loading
   #     the association via eager_graph.
   #   - :graph_conditions - The additional conditions to use on the SQL join when eagerly loading
@@ -248,6 +253,11 @@ def association_dataset_method_name(name)
     :"#{name}_dataset"
   end
 
+  # Name symbol for _dataset association method
+  def association_eager_dataset_method_name(name)
+    :"#{name}_eager_dataset"
+  end
+
   # Name symbol for _helper internal association method
   def association_helper_method_name(name)
     :"#{name}_helper"
@@ -288,9 +298,11 @@ def def_add_method(opts)
   def def_association_dataset_methods(opts)
     name = opts[:name]
     dataset_method = association_dataset_method_name(name)
+    eager_dataset_method = association_eager_dataset_method_name(name)
     helper_method = association_helper_method_name(name)
     dataset = opts[:dataset]
     dataset_helper = opts[:block]
+    eager_block = opts[:eager_block]
     order = opts[:order]
     eager = opts[:eager]
     eager_graph = opts[:eager_graph]
@@ -317,6 +329,18 @@ def def_association_dataset_methods(opts)
       ds
     end
 
+    # define a method returning the association dataset suitable for eager_loading
+    meta_def(eager_dataset_method) do |ds, select, associations|
+      ds = ds.select(*select)
+      ds = ds.order(*order) if order
+      ds = ds.limit(*limit) if limit
+      ds = ds.eager(eager) if eager
+      ds = ds.eager_graph(eager_graph) if eager_graph
+      ds = ds.eager(associations) unless associations.blank?
+      ds = eager_block.call(ds) if eager_block
+      ds
+    end
+
     class_def(name) do |*reload|
       if @associations.include?(name) and !reload[0]
         @associations[name]
@@ -336,17 +360,27 @@ def def_association_dataset_methods(opts)
   # Adds many_to_many association instance methods
   def def_many_to_many(opts)
     name = opts[:name]
+    model = self
     left = (opts[:left_key] ||= default_remote_key)
     right = (opts[:right_key] ||= default_foreign_key(opts))
     opts[:class_name] ||= name.to_s.singularize.camelize
     join_table = (opts[:join_table] ||= default_join_table_name(opts))
-    opts[:left_key_alias] ||= :x_foreign_key_x
-    opts[:left_key_select] ||= left.qualify(join_table).as(opts[:left_key_alias])
+    left_key_alias = opts[:left_key_alias] ||= :x_foreign_key_x
+    left_key_select = opts[:left_key_select] ||= left.qualify(join_table).as(opts[:left_key_alias])
     opts[:graph_join_table_conditions] = opts[:graph_join_table_conditions] ? opts[:graph_join_table_conditions].to_a : []
     opts[:graph_join_table_join_type] ||= opts[:graph_join_type]
     opts[:dataset] ||= proc{opts.associated_class.inner_join(join_table, [[right, opts.associated_primary_key], [left, pk]])}
     database = db
 
+    opts[:eager_loader] ||= proc do |key_hash, records, associations|
+      h = key_hash[model.primary_key]
+      records.each{|object| object.associations[name] = []}
+      model.send(association_eager_dataset_method_name(name), opts.associated_class.inner_join(join_table, [[right, opts.associated_primary_key], [left, h.keys]]), Array(opts.select) + Array(left_key_select), associations).all do |assoc_record|
+        next unless objects = h[assoc_record.values.delete(left_key_alias)]
+        objects.each{|object| object.associations[name].push(assoc_record)}
+      end
+    end
+
     def_association_dataset_methods(opts)
 
     return if opts[:read_only]
@@ -374,12 +408,27 @@ def def_many_to_many(opts)
   # Adds many_to_one association instance methods
   def def_many_to_one(opts)
     name = opts[:name]
+    model = self
     key = (opts[:key] ||= default_foreign_key(opts))
     opts[:class_name] ||= name.to_s.camelize
     opts[:dataset] ||= proc do
       klass = opts.associated_class
       klass.filter(opts.associated_primary_key.qualify(klass.table_name)=>send(key))
     end
+    opts[:eager_loader] ||= proc do |key_hash, records, associations|
+      h = key_hash[key]
+      keys = h.keys
+      # Skip eager loading if no objects have a foreign key for this association
+      unless keys.empty?
+        # Default the cached association to nil, so any object that doesn't have it
+        # populated will have cached the negative lookup.
+        records.each{|object| object.associations[name] = nil}
+        model.send(association_eager_dataset_method_name(name), opts.associated_class.filter(opts.associated_primary_key.qualify(opts.associated_class.table_name)=>keys),  opts.select, associations).all do |assoc_record|
+          next unless objects = h[assoc_record.pk]
+          objects.each{|object| object.associations[name] = assoc_record}
+        end
+      end
+    end
 
     def_association_dataset_methods(opts)
 
@@ -401,12 +450,25 @@ def def_many_to_one(opts)
   # Adds one_to_many association instance methods
   def def_one_to_many(opts)
     name = opts[:name]
+    model = self
     key = (opts[:key] ||= default_remote_key)
     opts[:class_name] ||= name.to_s.singularize.camelize
     opts[:dataset] ||= proc do
       klass = opts.associated_class
       klass.filter(key.qualify(klass.table_name) => pk)
     end
+    opts[:eager_loader] ||= proc do |key_hash, records, associations|
+      h = key_hash[model.primary_key]
+      records.each{|object| object.associations[name] = []}
+      reciprocal = opts.reciprocal
+      model.send(association_eager_dataset_method_name(name), opts.associated_class.filter(key.qualify(opts.associated_class.table_name)=>h.keys), opts.select, associations).all do |assoc_record|
+        next unless objects = h[assoc_record[key]]
+        objects.each do |object| 
+          object.associations[name].push(assoc_record)
+          assoc_record.associations[reciprocal] = object if reciprocal
+        end
+      end
+    end
 
     def_association_dataset_methods(opts)
 

sequel/lib/sequel_model/eager_loading.rb

@@ -362,64 +362,7 @@ def eager_load(a)
       end
     end
 
-    # Iterate through eager associations and assign instance variables
-    # for the association for all model objects
-    reflections.each do |reflection|
-      assoc_class = reflection.associated_class
-      assoc_name = reflection[:name]
-      # Proc for setting cascaded eager loading
-      assoc_block = Proc.new do |d|
-        d = d.order(*reflection[:order]) if reflection[:order]
-        d = d.limit(*reflection[:limit]) if reflection[:limit]
-        d = d.eager_graph(reflection[:eager_graph]) if reflection[:eager_graph]
-        d = d.eager(reflection[:eager]) if reflection[:eager]
-        d = d.eager(eager_assoc[assoc_name]) if eager_assoc[assoc_name]
-        d = reflection[:eager_block].call(d) if reflection[:eager_block]
-        d
-      end
-      case rtype = reflection[:type]
-        when :many_to_one
-          key = reflection[:key]
-          h = key_hash[key]
-          keys = h.keys
-          # No records have the foreign key set for this association, so skip it
-          next unless keys.length > 0
-          # Set the instance variable to null by default, so records that
-          # don't have a associated records will cache the negative lookup.
-          a.each do |object|
-            object.associations[assoc_name] = nil
-          end
-          assoc_block.call(assoc_class.select(*reflection.select).filter(assoc_class.primary_key.qualify(assoc_class.table_name)=>keys)).all do |assoc_object|
-            next unless objects = h[assoc_object.pk]
-            objects.each{|object| object.associations[assoc_name] = assoc_object}
-          end
-        when :one_to_many, :many_to_many
-          h = key_hash[model.primary_key]
-          ds = if rtype == :one_to_many
-            fkey = reflection[:key]
-            reciprocal = reflection.reciprocal
-            assoc_class.select(*reflection.select).filter(fkey.qualify(assoc_class.table_name)=>h.keys)
-          else
-            fkey = reflection[:left_key_alias]
-            assoc_class.select(*(Array(reflection.select)+Array(reflection[:left_key_select]))).inner_join(reflection[:join_table], [[reflection[:right_key], reflection.associated_primary_key], [reflection[:left_key], h.keys]])
-          end
-          h.values.each do |object_array|
-            object_array.each{|object| object.associations[assoc_name] = []}
-          end
-          assoc_block.call(ds).all do |assoc_object|
-            fk = if rtype == :many_to_many
-              assoc_object.values.delete(fkey)
-            else
-              assoc_object[fkey]
-            end
-            next unless objects = h[fk]
-            objects.each do |object|
-              object.associations[assoc_name].push(assoc_object)
-              assoc_object.associations[reciprocal] = object if reciprocal
-            end
-          end
-      end
-    end
+    reflections.each{|r| r[:eager_loader].call(key_hash, a, eager_assoc[r[:name]])}
   end
 
   # Build associations from the graph if #eager_graph was used, 

sequel/spec/eager_loading_spec.rb

@@ -403,6 +403,44 @@ def fetch_rows(sql)
     EagerAlbum.eager(:genre_names).all
     MODEL_DB.sqls.should == ['SELECT * FROM albums', "SELECT id, ag.album_id AS x_foreign_key_x FROM genres INNER JOIN ag ON ((ag.genre_id = genres.id) AND (ag.album_id IN (1)))"]
   end
+
+  it "should use the :eager_loader association option when eager loading" do
+    EagerAlbum.many_to_one :special_band, :eager_loader=>(proc do |key_hash, records, assocs| 
+      item = EagerBand.filter(:album_id=>records.collect{|r| [r.pk, r.pk*2]}.flatten).order(:name).first
+      records.each{|r| r.associations[:special_band] = item}
+    end)
+    EagerAlbum.one_to_many :special_tracks, :eager_loader=>(proc do |key_hash, records, assocs| 
+      items = EagerTrack.filter(:album_id=>records.collect{|r| [r.pk, r.pk*2]}.flatten).all
+      records.each{|r| r.associations[:special_tracks] = items}
+    end)
+    EagerAlbum.many_to_many :special_genres, :eager_loader=>(proc do |key_hash, records, assocs| 
+      items = EagerGenre.inner_join(:ag, [:genre_id]).filter(:album_id=>records.collect{|r| r.pk}).all
+      records.each{|r| r.associations[:special_genres] = items}
+    end)
+    a = EagerAlbum.eager(:special_genres, :special_tracks, :special_band).all
+    a.should be_a_kind_of(Array)
+    a.size.should == 1
+    a.first.should be_a_kind_of(EagerAlbum)
+    a.first.values.should == {:id => 1, :band_id => 2}
+    MODEL_DB.sqls.length.should == 4
+    MODEL_DB.sqls[0].should == 'SELECT * FROM albums'
+    MODEL_DB.sqls[1..-1].should(include('SELECT * FROM bands WHERE (album_id IN (1, 2)) ORDER BY name LIMIT 1'))
+    MODEL_DB.sqls[1..-1].should(include('SELECT * FROM tracks WHERE (album_id IN (1, 2))'))
+    MODEL_DB.sqls[1..-1].should(include('SELECT * FROM genres INNER JOIN ag USING (genre_id) WHERE (album_id IN (1))'))
+    a = a.first
+    a.special_band.should be_a_kind_of(EagerBand)
+    a.special_band.values.should == {:id => 2}
+    a.special_tracks.should be_a_kind_of(Array)
+    a.special_tracks.size.should == 1
+    a.special_tracks.first.should be_a_kind_of(EagerTrack)
+    a.special_tracks.first.values.should == {:id => 3, :album_id=>1}
+    a.special_genres.should be_a_kind_of(Array)
+    a.special_genres.size.should == 1
+    a.special_genres.first.should be_a_kind_of(EagerGenre)
+    a.special_genres.first.values.should == {:id => 4}
+    MODEL_DB.sqls.length.should == 4
+  end
+
 end
 
 describe Sequel::Model, "#eager_graph" do