Introduce preload query strategy for eager :includes. Closes #9640.

git-svn-id: http://svn-commit.rubyonrails.org/rails/trunk@8672 5ecf4fe2-1ee6-0310-87b1-e25e094e27de

1 files changed, 231 insertions, 0 deletions

diff --git a/activerecord/lib/active_record/association_preload.rb b/activerecord/lib/active_record/association_preload.rb
new file mode 100644
index 0000000000..6fa49dfbe8
--- /dev/null
+++ b/activerecord/lib/active_record/association_preload.rb
@@ -0,0 +1,231 @@
+module ActiveRecord
+  module AssociationPreload #:nodoc:
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+
+      # Loads the named associations for the activerecord record (or records) given
+      # preload_options is passed only one level deep: don't pass to the child associations when associations is a Hash
+      protected
+      def preload_associations(records, associations, preload_options={})
+        records = [records].flatten.compact
+        return if records.empty?
+        case associations
+        when Array then associations.each {|association| preload_associations(records, association, preload_options)}
+        when Symbol, String then preload_one_association(records, associations.to_sym, preload_options)
+        when Hash then
+          associations.each do |parent, child|
+            raise "parent must be an association name" unless parent.is_a?(String) || parent.is_a?(Symbol)
+            preload_associations(records, parent, preload_options)
+            reflection = reflections[parent]
+            parents = records.map {|record| record.send(reflection.name)}.flatten
+            unless parents.empty?
+              parents.first.class.preload_associations(parents, child)
+            end
+          end
+        end
+      end
+
+      private
+
+      def preload_one_association(records, association, preload_options={})
+        reflection = reflections[association]
+        raise ConfigurationError, "Association named '#{ association }' was not found; perhaps you misspelled it?" unless reflection
+
+        send(:"preload_#{reflection.macro}_association", records, reflection, preload_options)
+      end
+
+      def add_preloaded_records_to_collection(parent_records, reflection_name, associated_record)
+        parent_records.each do |parent_record|
+          association_proxy = parent_record.send(reflection_name)
+          association_proxy.loaded
+          association_proxy.target.push(*[associated_record].flatten)
+        end
+      end
+
+      def set_association_collection_records(id_to_record_map, reflection_name, associated_records, key)
+        associated_records.each do |associated_record|
+          mapped_records = id_to_record_map[associated_record[key].to_i]
+          add_preloaded_records_to_collection(mapped_records, reflection_name, associated_record)
+        end
+      end
+
+      def set_association_single_records(id_to_record_map, reflection_name, associated_records, key)
+        associated_records.each do |associated_record|
+          mapped_records = id_to_record_map[associated_record[key].to_i]
+          mapped_records.each do |mapped_record|
+            mapped_record.send("set_#{reflection_name}_target", associated_record)
+          end
+        end
+      end
+
+      def construct_id_map(records)
+        id_to_record_map = {}
+        ids = []
+        records.each do |record|
+          ids << record.id
+          mapped_records = (id_to_record_map[record.id] ||= [])
+          mapped_records << record
+        end
+        ids.uniq!
+        return id_to_record_map, ids
+      end
+
+      # FIXME: quoting
+      def preload_has_and_belongs_to_many_association(records, reflection, preload_options={})
+        table_name = reflection.klass.table_name
+        id_to_record_map, ids = construct_id_map(records)
+        records.each {|record| record.send(reflection.name).loaded}
+        options = reflection.options
+
+        conditions = "t0.#{reflection.primary_key_name}  IN (?)"
+        conditions << append_conditions(options, preload_options)
+
+        associated_records = reflection.klass.find(:all, :conditions => [conditions, ids],
+        :include => options[:include],
+        :joins => "INNER JOIN #{options[:join_table]} as t0 ON #{reflection.klass.table_name}.#{reflection.klass.primary_key} = t0.#{reflection.association_foreign_key}",
+        :select => "#{options[:select] || table_name+'.*'}, t0.#{reflection.primary_key_name} as _parent_record_id",
+        :order => options[:order])
+
+        set_association_collection_records(id_to_record_map, reflection.name, associated_records, '_parent_record_id')
+      end
+
+      def preload_has_one_association(records, reflection, preload_options={})
+        id_to_record_map, ids = construct_id_map(records)
+        records.each {|record| record.send("set_#{reflection.name}_target", nil)}
+
+        set_association_single_records(id_to_record_map, reflection.name, find_associated_records(ids, reflection, preload_options),
+                                       reflection.primary_key_name)
+      end
+
+      def preload_has_many_association(records, reflection, preload_options={})
+        id_to_record_map, ids = construct_id_map(records)
+        records.each {|record| record.send(reflection.name).loaded}
+        options = reflection.options
+
+        if options[:through]
+          through_records = preload_through_records(records, reflection, options[:through])
+          through_reflection = reflections[options[:through]]
+          through_primary_key = through_reflection.primary_key_name
+          unless through_records.empty?
+            source = reflection.source_reflection.name
+            through_records.first.class.preload_associations(through_records, source)
+            through_records.compact.each do |through_record|
+              add_preloaded_records_to_collection(id_to_record_map[through_record[through_primary_key].to_i],
+                                                 reflection.name, through_record.send(source))
+            end
+          end
+        else
+          set_association_collection_records(id_to_record_map, reflection.name, find_associated_records(ids, reflection, preload_options),
+                                             reflection.primary_key_name)
+        end
+      end
+
+      def preload_through_records(records, reflection, through_association)
+        through_reflection = reflections[through_association]
+        through_primary_key = through_reflection.primary_key_name
+
+        if reflection.options[:source_type]
+          interface = reflection.source_reflection.options[:foreign_type]
+          preload_options = {:conditions => ["#{interface} = ?", reflection.options[:source_type]]}
+
+          records.first.class.preload_associations(records, through_association, preload_options)
+
+          # Dont cache the association - we would only be caching a subset
+          through_records = []
+          records.compact.each do |record|
+            proxy = record.send(through_association)
+            through_records << proxy.target
+            proxy.reset
+          end
+          through_records = through_records.flatten
+        else
+          records.first.class.preload_associations(records, through_association)
+          through_records = records.compact.map {|record| record.send(through_association)}.flatten
+        end
+      end
+
+      # FIXME: quoting
+      def preload_belongs_to_association(records, reflection, preload_options={})
+        options = reflection.options
+        primary_key_name = reflection.primary_key_name
+
+        if options[:polymorphic]
+          polymorph_type = options[:foreign_type]
+          klasses_and_ids = {}
+
+          # Construct a mapping from klass to a list of ids to load and a mapping of those ids back to their parent_records
+          records.each do |record|
+            klass = record.send(polymorph_type)
+            klass_id = record.send(primary_key_name)
+
+            id_map = klasses_and_ids[klass] ||= {}
+            id_list_for_klass_id = (id_map[klass_id] ||= [])
+            id_list_for_klass_id << record
+          end
+          klasses_and_ids = klasses_and_ids.to_a
+        else
+          id_map = {}
+          records.each do |record|
+            mapped_records = (id_map[record.send(primary_key_name)] ||= [])
+            mapped_records << record
+          end
+          klasses_and_ids = [[reflection.klass.name, id_map]]
+        end
+
+        klasses_and_ids.each do |klass_and_id|
+          klass_name, id_map = *klass_and_id
+          klass = klass_name.constantize
+
+          table_name = klass.table_name
+          conditions = "#{table_name}.#{primary_key} IN (?)"
+          conditions << append_conditions(options, preload_options)
+          associated_records = klass.find(:all, :conditions => [conditions, id_map.keys.uniq],
+                                          :include => options[:include],
+                                          :select => options[:select],
+                                          :joins => options[:joins],
+                                          :order => options[:order])
+          set_association_single_records(id_map, reflection.name, associated_records, 'id')
+        end
+      end
+
+      # FIXME: quoting
+      def find_associated_records(ids, reflection, preload_options)
+        options = reflection.options
+        table_name = reflection.klass.table_name
+
+        if interface = reflection.options[:as]
+          conditions = "#{reflection.klass.table_name}.#{interface}_id IN (?) and #{reflection.klass.table_name}.#{interface}_type = '#{self.base_class.name.demodulize}'"
+        else
+          foreign_key = reflection.primary_key_name
+          conditions = "#{reflection.klass.table_name}.#{foreign_key} IN (?)"
+        end
+
+        conditions << append_conditions(options, preload_options)
+
+        reflection.klass.find(:all,
+                              :select => (options[:select] || "#{table_name}.*"),
+                              :include => options[:include],
+                              :conditions => [conditions, ids],
+                              :joins => options[:joins],
+                              :group => options[:group],
+                              :order => options[:order])
+      end
+
+
+      def interpolate_sql_for_preload(sql)
+        instance_eval("%@#{sql.gsub('@', '\@')}@")
+      end
+
+      def append_conditions(options, preload_options)
+        sql = ""
+        sql << " AND (#{interpolate_sql_for_preload(sanitize_sql(options[:conditions]))})" if options[:conditions]
+        sql << " AND (#{sanitize_sql preload_options[:conditions]})" if preload_options[:conditions]
+        sql
+      end
+
+    end
+  end
+end