Fix STI in available_records causing new instances of records to be loaded from database

Removes grouping by available_record class names in ActiveRecord::Associations::Preloader::Batch so that records of differing STI classes can be located.

Author: octatone

activerecord/lib/active_record/associations/preloader/association.rb

@@ -205,7 +205,7 @@ def associate_records_from_unscoped(unscoped_records)
           return if preload_scope && !preload_scope.empty_scope?
           return if reflection.collection?
 
-          unscoped_records.each do |record|
+          unscoped_records.select { |r| r[association_key_name].present? }.each do |record|
             owners = owners_by_key[convert_key(record[association_key_name])]
             owners&.each_with_index do |owner, i|
               association = owner.association(reflection.name)

activerecord/lib/active_record/associations/preloader/batch.rb

@@ -6,15 +6,15 @@ class Preloader
       class Batch # :nodoc:
         def initialize(preloaders, available_records:)
           @preloaders = preloaders.reject(&:empty?)
-          @available_records = available_records.flatten.group_by(&:class)
+          @available_records = available_records.flatten
         end
 
         def call
           branches = @preloaders.flat_map(&:branches)
           until branches.empty?
             loaders = branches.flat_map(&:runnable_loaders)
 
-            loaders.each { |loader| loader.associate_records_from_unscoped(@available_records[loader.klass]) }
+            loaders.each { |loader| loader.associate_records_from_unscoped(@available_records) }
 
             if loaders.any?
               future_tables = branches.flat_map do |branch|

activerecord/test/cases/associations_test.rb

@@ -804,6 +804,23 @@ def test_preload_with_available_records
     end
   end
 
+  def test_preload_with_available_records_sti
+    book = Book.create!
+    essay_special = EssaySpecial.create!
+    book.essay = essay_special
+    book.save!
+    book.reload
+
+    assert_not_predicate book.association(:essay), :loaded?
+
+    assert_no_queries do
+      ActiveRecord::Associations::Preloader.new(records: [book], associations: :essay, available_records: [[essay_special]]).call
+    end
+
+    assert_predicate book.association(:essay), :loaded?
+    assert_same essay_special, book.essay
+  end
+
   def test_preload_with_only_some_records_available
     bob_post = posts(:misc_by_bob)
     mary_post = posts(:misc_by_mary)

activerecord/test/models/book.rb

@@ -9,6 +9,8 @@ class Book < ActiveRecord::Base
   has_many :subscriptions
   has_many :subscribers, through: :subscriptions
 
+  has_one :essay
+
   enum status: [:proposed, :written, :published]
   enum last_read: { unread: 0, reading: 2, read: 3, forgotten: nil }
   enum nullable_status: [:single, :married]

activerecord/test/models/essay.rb

@@ -6,3 +6,6 @@ class Essay < ActiveRecord::Base
   belongs_to :category, primary_key: :name
   has_one :owner, primary_key: :name
 end
+
+class EssaySpecial < Essay
+end