From 99725d52198321561774b81c6f633f4983a276db Mon Sep 17 00:00:00 2001
From: Nick LaMuro <nicklamuro@gmail.com>
Date: Mon, 21 May 2018 17:38:44 -0500
Subject: [PATCH] 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       |  3 +-
 app/models/relationship.rb                    |  7 ++-
 spec/models/mixins/relationship_mixin_spec.rb | 56 +++++++++++++++++++
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/app/models/mixins/relationship_mixin.rb b/app/models/mixins/relationship_mixin.rb
index 8f4ddce0f8cd..6ecca5ace459 100644
--- a/app/models/mixins/relationship_mixin.rb
+++ b/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
diff --git a/app/models/relationship.rb b/app/models/relationship.rb
index 08bd479049c4..ba926322aee2 100644
--- a/app/models/relationship.rb
+++ b/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)
diff --git a/spec/models/mixins/relationship_mixin_spec.rb b/spec/models/mixins/relationship_mixin_spec.rb
index 41dff5cab217..a18007a8628c 100644
--- a/spec/models/mixins/relationship_mixin_spec.rb
+++ b/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