Use .polymorphic_preload_for_child_classes in .fulltree_arranged

Makes use of MiqPreloader.polymorphic_preload_for_child_classes in
RelationshipMixin.fulltree_arranged to allow the caller to preload
specific resource relationships and sub relationships.

app/models/mixins/relationship_mixin.rb

@@ -515,7 +515,8 @@ def fulltree_ids_arranged(*args)
 
   # Returns the records in the tree from the root arranged in a tree
   def fulltree_arranged(*args)
-    Relationship.arranged_rels_to_resources(fulltree_rels_arranged(*args))
+    class_specific_preloaders = args.last.delete(:class_specific_preloaders) if args.last.kind_of?(Hash)
+    Relationship.arranged_rels_to_resources(fulltree_rels_arranged(*args), true, class_specific_preloaders)
   end
 
   # Returns a list of all unique child types

app/models/relationship.rb

@@ -82,8 +82,11 @@ def self.flatten_arranged_rels(relationships)
     end
   end
 
-  def self.arranged_rels_to_resources(relationships, initial = true)
-    MiqPreloader.preload(flatten_arranged_rels(relationships), :resource) if initial
+  def self.arranged_rels_to_resources(relationships, initial = true, class_specific_preloaders = nil)
+    if initial
+      record_set = flatten_arranged_rels(relationships)
+      MiqPreloader.polymorphic_preload_for_child_classes(record_set, :resource, class_specific_preloaders)
+    end
 
     relationships.each_with_object({}) do |(rel, children), h|
       h[rel.resource] = arranged_rels_to_resources(children, false)

spec/models/mixins/relationship_mixin_spec.rb

@@ -747,6 +747,62 @@
         }
       )
     end
+
+    context "with a EMS based tree" do
+      # Note:  This shared_context overwrites the `let(:vms)` at the top of
+      # this file.
+      include_context "simple ems_metadata tree" do
+        before { init_full_tree }
+      end
+
+      it "builds the tree normally" do
+        nodes = ems.fulltree_arranged
+        expect(nodes).to eq(
+          ems => {
+            clusters[0] => {
+              hosts[0] => {
+                vms[0] => {},
+                vms[1] => {}
+              },
+              hosts[1] => {
+                vms[2] => {},
+                vms[3] => {}
+              }
+            },
+            clusters[1] => {
+              hosts[2] => {
+                vms[4] => {},
+                vms[5] => {}
+              },
+              hosts[3] => {
+                vms[6] => {},
+                vms[7] => {}
+              }
+            }
+          }
+        )
+      end
+
+      it "can preload certain relationships to avoid N+1s" do
+        hosts_scope   = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+        fulltree_opts = {
+          :except_type               => "VmOrTemplate",
+          :class_specific_preloaders => {
+            EmsCluster => [:hosts, hosts_scope],
+            Host       => hosts_scope
+          }
+        }
+        tree = ems.fulltree_arranged(fulltree_opts)
+
+        expect {
+          # get the v_total_vms through the EmsCluster records
+          tree.values.first.keys.flat_map(&:hosts).each(&:v_total_vms)
+
+          # get the v_total_vms through the Host records in the tree
+          tree.values.first.values.flat_map(&:keys).each(&:v_total_vms)
+        }.to match_query_limit_of(0)
+      end
+    end
   end
 
   describe "#fulltree_ids_arranged" do