From 38ef36c22f84a0e823008ab4e3eadba42543297b Mon Sep 17 00:00:00 2001
From: Jamis Buck <jamis.buck@mongodb.com>
Date: Wed, 22 Jan 2025 07:56:27 -0700
Subject: [PATCH] MONGOID-5839 Fix eager-loading from STI subclasses (#5934)

If the root of a query is an STI subclass (e.g. `Subclass.all`)
AND the query tries to eager load (`includes`) another association,
the eager load was failing because it was looking for inverse
associations using the STI subclass name, instead of the class at
the root of the hierarchy.
---
 lib/mongoid/association/eager_loadable.rb |  3 +++
 lib/mongoid/traversable.rb                | 12 +++++++++++
 spec/mongoid/association/eager_spec.rb    | 26 +++++++++++++++++++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/lib/mongoid/association/eager_loadable.rb b/lib/mongoid/association/eager_loadable.rb
index 5ae3b67bfd..84cfd14996 100644
--- a/lib/mongoid/association/eager_loadable.rb
+++ b/lib/mongoid/association/eager_loadable.rb
@@ -31,6 +31,9 @@ def preload(associations, docs)
         docs_map = {}
         queue = [ klass.to_s ]
 
+        # account for single-collection inheritance
+        queue.push(klass.root_class.to_s) if klass != klass.root_class
+
         while klass = queue.shift
           if as = assoc_map.delete(klass)
             as.each do |assoc|
diff --git a/lib/mongoid/traversable.rb b/lib/mongoid/traversable.rb
index 0098d4c4d2..79db072606 100644
--- a/lib/mongoid/traversable.rb
+++ b/lib/mongoid/traversable.rb
@@ -323,6 +323,18 @@ def hereditary?
         !!(Mongoid::Document > superclass)
       end
 
+      # Returns the root class of the STI tree that the current
+      # class participates in. If the class is not an STI subclass, this
+      # returns the class itself.
+      #
+      # @return [ Mongoid::Document ] the root of the STI tree
+      def root_class
+        root = self
+        root = root.superclass while root.hereditary?
+
+        root
+      end
+
       # When inheriting, we want to copy the fields from the parent class and
       # set the on the child to start, mimicking the behavior of the old
       # class_inheritable_accessor that was deprecated in Rails edge.
diff --git a/spec/mongoid/association/eager_spec.rb b/spec/mongoid/association/eager_spec.rb
index a194fc7411..2171f63040 100644
--- a/spec/mongoid/association/eager_spec.rb
+++ b/spec/mongoid/association/eager_spec.rb
@@ -14,14 +14,36 @@
       Mongoid::Contextual::Mongo.new(criteria)
     end
 
+    let(:association_host) { Account }
+
     let(:inclusions) do
       includes.map do |key|
-        Account.reflect_on_association(key)
+        association_host.reflect_on_association(key)
       end
     end
 
     let(:doc) { criteria.first }
 
+    context 'when root is an STI subclass' do
+      # Driver has_one Vehicle
+      # Vehicle belongs_to Driver
+      # Truck is a Vehicle
+
+      before do
+        Driver.create!(vehicle: Truck.new)
+      end
+
+      let(:criteria) { Truck.all }
+      let(:includes) { %i[ driver ] }
+      let(:association_host) { Truck }
+
+      it 'preloads the driver' do
+        expect(doc.ivar(:driver)).to be false
+        context.preload(inclusions, [ doc ])
+        expect(doc.ivar(:driver)).to be == Driver.first
+      end
+    end
+
     context "when belongs_to" do
 
       let!(:account) do
@@ -42,7 +64,7 @@
       it "preloads the parent" do
         expect(doc.ivar(:person)).to be false
         context.preload(inclusions, [doc])
-        expect(doc.ivar(:person)).to eq(doc.person)
+        expect(doc.ivar(:person)).to be == person
       end
     end