1 files changed, 42 insertions, 24 deletions

diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index f55b4824e5..83c69586e6 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -3,6 +3,7 @@ module ActiveRecord
     class Preloader
       class Association #:nodoc:
         attr_reader :owners, :reflection, :preload_scope, :model, :klass
+        attr_reader :preloaded_records
 
         def initialize(klass, owners, reflection, preload_scope)
           @klass         = klass
@@ -12,15 +13,14 @@ module ActiveRecord
           @model         = owners.first && owners.first.class
           @scope         = nil
           @owners_by_key = nil
+          @preloaded_records = []
         end
 
-        def run
-          unless owners.first.association(reflection.name).loaded?
-            preload
-          end
+        def run(preloader)
+          preload(preloader)
         end
 
-        def preload
+        def preload(preloader)
           raise NotImplementedError
         end
 
@@ -29,6 +29,10 @@ module ActiveRecord
         end
 
         def records_for(ids)
+          query_scope(ids)
+        end
+
+        def query_scope(ids)
           scope.where(association_key.in(ids))
         end
 
@@ -52,12 +56,9 @@ module ActiveRecord
           raise NotImplementedError
         end
 
-        # We're converting to a string here because postgres will return the aliased association
-        # key in a habtm as a string (for whatever reason)
         def owners_by_key
           @owners_by_key ||= owners.group_by do |owner|
-            key = owner[owner_key_name]
-            key && key.to_s
+            owner[owner_key_name]
           end
         end
 
@@ -67,38 +68,47 @@ module ActiveRecord
 
         private
 
-        def associated_records_by_owner
+        def associated_records_by_owner(preloader)
           owners_map = owners_by_key
           owner_keys = owners_map.keys.compact
 
-          if klass.nil? || owner_keys.empty?
-            records = []
-          else
+          # Each record may have multiple owners, and vice-versa
+          records_by_owner = owners.each_with_object({}) do |owner,h|
+            h[owner] = []
+          end
+
+          if owner_keys.any?
             # Some databases impose a limit on the number of ids in a list (in Oracle it's 1000)
             # Make several smaller queries if necessary or make one query if the adapter supports it
             sliced  = owner_keys.each_slice(klass.connection.in_clause_length || owner_keys.size)
-            records = sliced.map { |slice| records_for(slice).to_a }.flatten
-          end
-
-          # Each record may have multiple owners, and vice-versa
-          records_by_owner = Hash[owners.map { |owner| [owner, []] }]
-          records.each do |record|
-            owner_key = record[association_key_name].to_s
 
-            owners_map[owner_key].each do |owner|
-              records_by_owner[owner] << record
+            records = load_slices sliced
+            records.each do |record, owner_key|
+              owners_map[owner_key].each do |owner|
+                records_by_owner[owner] << record
+              end
             end
           end
+
           records_by_owner
         end
 
+        def load_slices(slices)
+          @preloaded_records = slices.flat_map { |slice|
+            records_for(slice)
+          }
+
+          @preloaded_records.map { |record|
+            [record, record[association_key_name]]
+          }
+        end
+
         def reflection_scope
           @reflection_scope ||= reflection.scope ? klass.unscoped.instance_exec(nil, &reflection.scope) : klass.unscoped
         end
 
         def build_scope
           scope = klass.unscoped
-          scope.default_scoped = true
 
           values         = reflection_scope.values
           reflection_binds = reflection_scope.bind_values
@@ -112,11 +122,19 @@ module ActiveRecord
           scope.select!   preload_values[:select] || values[:select] || table[Arel.star]
           scope.includes! preload_values[:includes] || values[:includes]
 
+          if preload_values.key? :order
+            scope.order! preload_values[:order]
+          else
+            if values.key? :order
+              scope.order! values[:order]
+            end
+          end
+
           if options[:as]
             scope.where!(klass.table_name => { reflection.type => model.base_class.sti_name })
           end
 
-          scope
+          klass.default_scoped.merge(scope)
         end
       end
     end