When preloading group by reflection rather than by class [#125 state:resolved]

This avoids extra queries when several subclasses inherit the association
from their parent class, while still coping with subclasses redefining
associations.

Signed-off-by: Pratik Naik <pratiknaik@gmail.com>

2 files changed, 12 insertions, 6 deletions

diff --git a/activerecord/lib/active_record/association_preload.rb b/activerecord/lib/active_record/association_preload.rb
index 660d7647c7..a3d1f12b03 100644
--- a/activerecord/lib/active_record/association_preload.rb
+++ b/activerecord/lib/active_record/association_preload.rb
@@ -31,12 +31,12 @@ module ActiveRecord
       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
-
-        # Not all records have the same class, so group then preload.
-        records.group_by(&:class).each do |klass, records|
-          reflection = klass.reflections[association]
+        class_to_reflection = {}
+        # Not all records have the same class, so group then preload
+        # group on the reflection itself so that if various subclass share the same association then we do not split them
+        # unncessarily
+        records.group_by {|record| class_to_reflection[record.class] ||= record.class.reflections[association]}.each do |reflection, records|
+          raise ConfigurationError, "Association named '#{ association }' was not found; perhaps you misspelled it?" unless reflection
           send("preload_#{reflection.macro}_association", records, reflection, preload_options)
         end
       end
diff --git a/activerecord/test/cases/associations/eager_test.rb b/activerecord/test/cases/associations/eager_test.rb
index fb8deedf95..3a3358e39b 100644
--- a/activerecord/test/cases/associations/eager_test.rb
+++ b/activerecord/test/cases/associations/eager_test.rb
@@ -603,4 +603,10 @@ class EagerAssociationTest < ActiveRecord::TestCase
       assert_equal 3, authors(:david).posts_with_comments.count(:conditions => "length(comments.body) > 15")
     end
   end
+
+  def test_load_with_sti_sharing_association
+    assert_queries(2) do #should not do 1 query per subclass
+      Comment.find :all, :include => :post
+    end
+  end
 end