Support preloading nested through associations (using the default multi-query strategy)

2 files changed, 216 insertions, 122 deletions

diff --git a/activerecord/lib/active_record/association_preload.rb b/activerecord/lib/active_record/association_preload.rb
index e6b367790b..664b0a7d59 100644
--- a/activerecord/lib/active_record/association_preload.rb
+++ b/activerecord/lib/active_record/association_preload.rb
@@ -202,93 +202,108 @@ module ActiveRecord
         set_association_collection_records(id_to_record_map, reflection.name, associated_records, 'the_parent_record_id')
       end
 
-      def preload_has_one_association(records, reflection, preload_options={})
-        return if records.first.send("loaded_#{reflection.name}?")
-        id_to_record_map, ids = construct_id_map(records, reflection.options[:primary_key])
-        options = reflection.options
-        records.each {|record| record.send("set_#{reflection.name}_target", nil)}
-        if options[:through]
-          through_records = preload_through_records(records, reflection, options[:through])
-
-          unless through_records.empty?
-            through_reflection = reflections[options[:through]]
-            through_primary_key = through_reflection.primary_key_name
-            source = reflection.source_reflection.name
-            through_records.first.class.preload_associations(through_records, source)
-            if through_reflection.macro == :belongs_to
-              id_to_record_map    = construct_id_map(records, through_primary_key).first
-              through_primary_key = through_reflection.klass.primary_key
-            end
-
-            through_records.each do |through_record|
-              add_preloaded_record_to_collection(id_to_record_map[through_record[through_primary_key].to_s],
-                                                 reflection.name, through_record.send(source))
-            end
-          end
+      def preload_has_one_or_has_many_association(records, reflection, preload_options={})
+        if reflection.macro == :has_many
+          return if records.first.send(reflection.name).loaded?
+          records.each { |record| record.send(reflection.name).loaded }
         else
-          set_association_single_records(id_to_record_map, reflection.name, find_associated_records(ids, reflection, preload_options), reflection.primary_key_name)
+          return if records.first.send("loaded_#{reflection.name}?")
+          records.each {|record| record.send("set_#{reflection.name}_target", nil)}
         end
-      end
-
-      def preload_has_many_association(records, reflection, preload_options={})
-        return if records.first.send(reflection.name).loaded?
+        
         options = reflection.options
-
-        primary_key_name = reflection.through_reflection_primary_key_name
-        id_to_record_map, ids = construct_id_map(records, primary_key_name || reflection.options[:primary_key])
-        records.each {|record| record.send(reflection.name).loaded}
-
+        
         if options[:through]
-          through_records = preload_through_records(records, reflection, options[:through])
-          unless through_records.empty?
+          records_with_through_records = preload_through_records(records, reflection, options[:through])
+          all_through_records = records_with_through_records.map(&:last).flatten
+
+          unless all_through_records.empty?
             source = reflection.source_reflection.name
-            through_records.first.class.preload_associations(through_records, source, options)
-            through_records.each do |through_record|
-              through_record_id = through_record[reflection.through_reflection_primary_key].to_s
-              add_preloaded_records_to_collection(id_to_record_map[through_record_id], reflection.name, through_record.send(source))
+            all_through_records.first.class.preload_associations(all_through_records, source, options)
+            
+            records_with_through_records.each do |record, through_records|
+              source_records = through_records.map(&source).flatten.compact
+              
+              case reflection.macro
+                when :has_many, :has_and_belongs_to_many
+                  add_preloaded_records_to_collection([record], reflection.name, source_records)
+                when :has_one, :belongs_to
+                  add_preloaded_record_to_collection([record], reflection.name, source_records.first)
+              end
             end
           end
-
         else
-          set_association_collection_records(id_to_record_map, reflection.name, find_associated_records(ids, reflection, preload_options),
-                                             reflection.primary_key_name)
+          id_to_record_map, ids = construct_id_map(records, reflection.options[:primary_key])
+          associated_records = find_associated_records(ids, reflection, preload_options)
+          
+          if reflection.macro == :has_many
+            set_association_collection_records(
+              id_to_record_map, reflection.name,
+              associated_records, reflection.primary_key_name
+            )
+          else
+            set_association_single_records(
+              id_to_record_map, reflection.name,
+              associated_records, reflection.primary_key_name
+            )
+          end
         end
       end
+      
+      alias_method :preload_has_one_association, :preload_has_one_or_has_many_association
+      alias_method :preload_has_many_association, :preload_has_one_or_has_many_association
 
       def preload_through_records(records, reflection, through_association)
         through_reflection = reflections[through_association]
 
-        through_records = []
+        # If the same through record is loaded twice, we want to return exactly the same
+        # object in the result, rather than two separate instances representing the same
+        # record. This is so that we can preload the source association for each record,
+        # and always be able to access the preloaded association regardless of where we
+        # refer to the record.
+        # 
+        # Suffices to say, if AR had an identity map built in then this would be unnecessary.
+        identity_map = {}
+        
+        options = {}
+        
         if reflection.options[:source_type]
           interface = reflection.source_reflection.options[:foreign_type]
-          preload_options = {:conditions => ["#{connection.quote_column_name interface} = ?", reflection.options[:source_type]]}
-
+          options[:conditions] = ["#{connection.quote_column_name interface} = ?", reflection.options[:source_type]]
           records.compact!
-          records.first.class.preload_associations(records, through_association, preload_options)
+        else
+          if reflection.options[:conditions]
+            options[:include]    = reflection.options[:include] ||
+                                   reflection.options[:source] 
+            options[:conditions] = reflection.options[:conditions]
+          end
+          
+          options[:order] = reflection.options[:order]
+        end
+        
+        records.first.class.preload_associations(records, through_association, options)
 
-          # Dont cache the association - we would only be caching a subset
-          records.each do |record|
+        records.map do |record|
+          if reflection.options[:source_type]
+            # Dont cache the association - we would only be caching a subset
             proxy = record.send(through_association)
-
+            
             if proxy.respond_to?(:target)
-              through_records.concat Array.wrap(proxy.target)
+              through_records = proxy.target
               proxy.reset
             else # this is a has_one :through reflection
-              through_records << proxy if proxy
+              through_records = proxy
             end
+          else
+            through_records = record.send(through_association)
           end
-        else
-          options = {}
-          options[:include] = reflection.options[:include] || reflection.options[:source] if reflection.options[:conditions]
-          options[:order] = reflection.options[:order]
-          options[:conditions] = reflection.options[:conditions]
-          records.first.class.preload_associations(records, through_association, options)
-
-          records.each do |record|
-            through_records.concat Array.wrap(record.send(through_association))
+          
+          through_records = Array.wrap(through_records).map do |through_record|
+            identity_map[through_record] ||= through_record
           end
+          
+          [record, through_records]
         end
-        through_records
       end
 
       def preload_belongs_to_association(records, reflection, preload_options={})
diff --git a/activerecord/test/cases/associations/nested_has_many_through_associations_test.rb b/activerecord/test/cases/associations/nested_has_many_through_associations_test.rb
index eea1c4e54c..32b03bf076 100644
--- a/activerecord/test/cases/associations/nested_has_many_through_associations_test.rb
+++ b/activerecord/test/cases/associations/nested_has_many_through_associations_test.rb
@@ -54,70 +54,98 @@ class NestedHasManyThroughAssociationsTest < ActiveRecord::TestCase
   # Source: has_many through
   # Through: has_many
   def test_has_many_through_has_many_with_has_many_through_source_reflection
-    author = authors(:david)
-    assert_equal [tags(:general), tags(:general)], author.tags
+    general = tags(:general)
+    
+    assert_equal [general, general], authors(:david).tags
     
     # Only David has a Post tagged with General
     authors = Author.joins(:tags).where('tags.id' => tags(:general).id)
     assert_equal [authors(:david)], authors.uniq
     
-    authors = Author.includes(:tags)
-    assert_equal [tags(:general), tags(:general)], authors.first.tags
-    
     # This ensures that the polymorphism of taggings is being observed correctly
     authors = Author.joins(:tags).where('taggings.taggable_type' => 'FakeModel')
     assert authors.empty?
+    
+    assert_queries(5) do
+      authors = Author.includes(:tags).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [general, general], authors.first.tags
+    end
   end
 
   # has_many through
   # Source: has_many
   # Through: has_many through
   def test_has_many_through_has_many_through_with_has_many_source_reflection
+    luke, david = subscribers(:first), subscribers(:second)
+    
     author = authors(:david)
-    assert_equal [subscribers(:first), subscribers(:second), subscribers(:second)], author.subscribers
+    assert_equal [luke, david, david], author.subscribers
     
     # All authors with subscribers where one of the subscribers' nick is 'alterself'
     authors = Author.joins(:subscribers).where('subscribers.nick' => 'alterself')
     assert_equal [authors(:david)], authors
     
-    # TODO: Make this work
-    # authors = Author.includes(:subscribers)
-    # assert_equal [subscribers(:first), subscribers(:second), subscribers(:second)], authors.first.subscribers
+    assert_queries(4) do
+      authors = Author.includes(:subscribers).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [luke, david, david], authors.first.subscribers.sort_by(&:nick)
+    end
+    
+    # TODO: Add eager loading test using LEFT OUTER JOIN
   end
   
   # has_many through
   # Source: has_one through
   # Through: has_one
   def test_has_many_through_has_one_with_has_one_through_source_reflection
-    assert_equal [member_types(:founding)], members(:groucho).nested_member_types
+    founding = member_types(:founding)
+  
+    assert_equal [founding], members(:groucho).nested_member_types
     
-    members = Member.joins(:nested_member_types).where('member_types.id' => member_types(:founding).id)
+    members = Member.joins(:nested_member_types).where('member_types.id' => founding.id)
     assert_equal [members(:groucho)], members
     
-    members = Member.includes(:nested_member_types)
-    assert_equal [member_types(:founding)], members.first.nested_member_types
+    assert_queries(4) do
+      members = Member.includes(:nested_member_types).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [founding], members.first.nested_member_types
+    end
   end
   
   # has_many through
   # Source: has_one
   # Through: has_one through
   def test_has_many_through_has_one_through_with_has_one_source_reflection
-    assert_equal [sponsors(:moustache_club_sponsor_for_groucho)], members(:groucho).nested_sponsors
+    mustache = sponsors(:moustache_club_sponsor_for_groucho)
     
-    members = Member.joins(:nested_sponsors).where('sponsors.id' => sponsors(:moustache_club_sponsor_for_groucho).id)
+    assert_equal [mustache], members(:groucho).nested_sponsors
+    
+    members = Member.joins(:nested_sponsors).where('sponsors.id' => mustache.id)
     assert_equal [members(:groucho)], members
     
-    # TODO: Make this work
-    # members = Member.includes(:nested_sponsors)
-    # assert_equal [sponsors(:moustache_club_sponsor_for_groucho)], members.first.nested_sponsors
+    assert_queries(4) do
+      members = Member.includes(:nested_sponsors).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [mustache], members.first.nested_sponsors
+    end
   end
   
   # has_many through
   # Source: has_many through
   # Through: has_one
   def test_has_many_through_has_one_with_has_many_through_source_reflection
-    assert_equal [member_details(:groucho), member_details(:some_other_guy)],
-                 members(:groucho).organization_member_details
+    groucho_details, other_details = member_details(:groucho), member_details(:some_other_guy)
+  
+    assert_equal [groucho_details, other_details], members(:groucho).organization_member_details
     
     members = Member.joins(:organization_member_details).
                      where('member_details.id' => member_details(:groucho).id)
@@ -127,127 +155,178 @@ class NestedHasManyThroughAssociationsTest < ActiveRecord::TestCase
                      where('member_details.id' => 9)
     assert members.empty?
     
-    members = Member.includes(:organization_member_details)
-    assert_equal [member_details(:groucho), member_details(:some_other_guy)],
-                 members.first.organization_member_details
+    assert_queries(4) do
+      members = Member.includes(:organization_member_details).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [groucho_details, other_details], members.first.organization_member_details
+    end
   end
   
   # has_many through
   # Source: has_many
   # Through: has_one through
   def test_has_many_through_has_one_through_with_has_many_source_reflection
-    assert_equal [member_details(:groucho), member_details(:some_other_guy)],
-                 members(:groucho).organization_member_details_2
+    groucho_details, other_details = member_details(:groucho), member_details(:some_other_guy)
+  
+    assert_equal [groucho_details, other_details], members(:groucho).organization_member_details_2
     
     members = Member.joins(:organization_member_details_2).
-                     where('member_details.id' => member_details(:groucho).id)
+                     where('member_details.id' => groucho_details.id)
     assert_equal [members(:groucho), members(:some_other_guy)], members
     
     members = Member.joins(:organization_member_details_2).
                      where('member_details.id' => 9)
     assert members.empty?
     
-    # TODO: Make this work
-    # members = Member.includes(:organization_member_details_2)
-    # assert_equal [member_details(:groucho), member_details(:some_other_guy)],
-    #              members.first.organization_member_details_2
+    assert_queries(4) do
+      members = Member.includes(:organization_member_details_2).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [groucho_details, other_details], members.first.organization_member_details_2
+    end
   end
   
   # has_many through
   # Source: has_and_belongs_to_many
   # Through: has_many
   def test_has_many_through_has_many_with_has_and_belongs_to_many_source_reflection
-    assert_equal [categories(:general), categories(:cooking)], authors(:bob).post_categories
+    general, cooking = categories(:general), categories(:cooking)
+  
+    assert_equal [general, cooking], authors(:bob).post_categories
     
-    authors = Author.joins(:post_categories).where('categories.id' => categories(:cooking).id)
+    authors = Author.joins(:post_categories).where('categories.id' => cooking.id)
     assert_equal [authors(:bob)], authors
     
-    authors = Author.includes(:post_categories)
-    assert_equal [categories(:general), categories(:cooking)], authors[2].post_categories
+    assert_queries(3) do
+      authors = Author.includes(:post_categories).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [general, cooking], authors[2].post_categories
+    end
   end
   
   # has_many through
   # Source: has_many
   # Through: has_and_belongs_to_many
   def test_has_many_through_has_and_belongs_to_many_with_has_many_source_reflection
-    assert_equal [comments(:greetings), comments(:more_greetings)], categories(:technology).post_comments
+    greetings, more = comments(:greetings), comments(:more_greetings)
+  
+    assert_equal [greetings, more], categories(:technology).post_comments
     
-    categories = Category.joins(:post_comments).where('comments.id' => comments(:more_greetings).id)
+    categories = Category.joins(:post_comments).where('comments.id' => more.id)
     assert_equal [categories(:general), categories(:technology)], categories
     
-    # TODO: Make this work
-    # categories = Category.includes(:post_comments)
-    # assert_equal [comments(:greetings), comments(:more_greetings)], categories[1].post_comments
+    assert_queries(3) do
+      categories = Category.includes(:post_comments).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [greetings, more], categories[1].post_comments
+    end
   end
   
   # has_many through
   # Source: has_many through a habtm
   # Through: has_many through
   def test_has_many_through_has_many_with_has_many_through_habtm_source_reflection
-    assert_equal [comments(:greetings), comments(:more_greetings)], authors(:bob).category_post_comments
+    greetings, more = comments(:greetings), comments(:more_greetings)
+  
+    assert_equal [greetings, more], authors(:bob).category_post_comments
     
     authors = Author.joins(:category_post_comments).where('comments.id' => comments(:does_it_hurt).id)
     assert_equal [authors(:david), authors(:mary)], authors
     
-    comments = Author.joins(:category_post_comments)
-    assert_equal [comments(:greetings), comments(:more_greetings)], comments[2].category_post_comments
+    assert_queries(5) do
+      authors = Author.includes(:category_post_comments).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [greetings, more], authors[2].category_post_comments
+    end
   end
   
   # has_many through
   # Source: belongs_to
   # Through: has_many through
   def test_has_many_through_has_many_through_with_belongs_to_source_reflection
-    author = authors(:david)
-    assert_equal [tags(:general), tags(:general)], author.tagging_tags
+    general = tags(:general)
+    
+    assert_equal [general, general], authors(:david).tagging_tags
     
     authors = Author.joins(:tagging_tags).where('tags.id' => tags(:general).id)
     assert_equal [authors(:david)], authors.uniq
     
-    # TODO: Make this work
-    # authors = Author.includes(:tagging_tags)
-    # assert_equal [tags(:general), tags(:general)], authors.first.tagging_tags
+    assert_queries(5) do
+      authors = Author.includes(:tagging_tags).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [general, general], authors.first.tagging_tags
+    end
   end
   
   # has_many through
   # Source: has_many through
   # Through: belongs_to
   def test_has_many_through_belongs_to_with_has_many_through_source_reflection
-    assert_equal [taggings(:welcome_general), taggings(:thinking_general)],
-                 categorizations(:david_welcome_general).post_taggings
+    welcome_general, thinking_general = taggings(:welcome_general), taggings(:thinking_general)
+  
+    assert_equal [welcome_general, thinking_general], categorizations(:david_welcome_general).post_taggings
     
-    categorizations = Categorization.joins(:post_taggings).where('taggings.id' => taggings(:welcome_general).id)
+    categorizations = Categorization.joins(:post_taggings).where('taggings.id' => welcome_general.id)
     assert_equal [categorizations(:david_welcome_general)], categorizations
     
-    categorizations = Categorization.includes(:post_taggings)
-    assert_equal [taggings(:welcome_general), taggings(:thinking_general)],
-                 categorizations.first.post_taggings
+    assert_queries(4) do
+      categorizations = Categorization.includes(:post_taggings).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal [welcome_general, thinking_general], categorizations.first.post_taggings
+    end
   end
   
   # has_one through
   # Source: has_one through
   # Through: has_one
   def test_has_one_through_has_one_with_has_one_through_source_reflection
-    assert_equal member_types(:founding), members(:groucho).nested_member_type
+    founding = member_types(:founding)
+  
+    assert_equal founding, members(:groucho).nested_member_type
     
-    members = Member.joins(:nested_member_type).where('member_types.id' => member_types(:founding).id)
+    members = Member.joins(:nested_member_type).where('member_types.id' => founding.id)
     assert_equal [members(:groucho)], members
     
-    members = Member.includes(:nested_member_type)
-    assert_equal member_types(:founding), members.first.nested_member_type
+    assert_queries(4) do
+      members = Member.includes(:nested_member_type).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal founding, members.first.nested_member_type
+    end
   end
   
   # has_one through
   # Source: belongs_to
   # Through: has_one through
   def test_has_one_through_has_one_through_with_belongs_to_source_reflection
-    assert_equal categories(:general), members(:groucho).club_category
+    general = categories(:general)
+  
+    assert_equal general, members(:groucho).club_category
     
     members = Member.joins(:club_category).where('categories.id' => categories(:technology).id)
     assert_equal [members(:blarpy_winkup)], members
     
-    # TODO: Make this work
-    # members = Member.includes(:club_category)
-    # assert_equal categories(:general), members.first.club_category
+    assert_queries(4) do
+      members = Member.includes(:club_category).to_a
+    end
+    
+    assert_no_queries do
+      assert_equal general, members.first.club_category
+    end
   end
 
   def test_distinct_has_many_through_a_has_many_through_association_on_source_reflection