Add MiqPreloader.polymorphic_preload_for_child_classes

This is a specialized preloader for classes with polymorphic
relationships that allows for targeted preloading for those specific
class types on said polymorphic relationships.  This allows for taking
advantage of those associations by allowing scopes to be applied to the
specific relationships and subsequent calls to the relationship can have
those query definitions applied.

lib/miq_preloader.rb

@@ -31,4 +31,88 @@ def self.preload_and_scope(records, association_name)
       target_klass.where(join_key.key.to_sym => records.select(join_key.foreign_key.to_sym))
     end
   end
+
+  # Allows having a polymorphic preloader, but then having class specific
+  # preloaders fire for the loaded polymorphic classes.
+  #
+  # @param records [ActiveRecord::Relation] collection of activerecord objects to preload into
+  # @param associations [Symbol|String|Array|Hash] association(s) to load (see .includes for examples)
+  # @param class_preloaders [Hash] keys are Classes, and values are associations for that polymorphic type
+  #
+  # Values for class_preloaders can either be an Array of two (args for sub
+  # preloader relationships), or an arel statement, which is the arel scope to
+  # execute for that specific class only (no sub relationships preloaded.
+  #
+  # Example:
+  #
+  #   irb> tree = ExtManagementSystem.last.fulltree_rels_arranged(:except_type => "VmOrTemplate")
+  #   irb> records = Relationship.flatten_arranged_rels(tree)
+  #   irb> hosts_scope = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+  #   irb> preloaders_per_class = { EmsCluster => [:hosts, hosts_scope], Host => hosts_scope }
+  #   irb> MiqPreloader.polymorphic_preload_for_child_classes(records, nil, preloaders_per_class)
+  #
+  # Note:  Class's .base_class are favored over their specific class
+  #
+  def self.polymorphic_preload_for_child_classes(records, associations, class_preloaders = {})
+    preloader = ActiveRecord::Associations::Preloader.new
+    preloader.extend(Module.new {
+      attr_accessor :class_specific_preloaders
+
+      # DIRTY HACK... but hey... at least I am just isolating it to this method
+      # right...
+      #
+      # Everyone else: "No Nick... just no..."
+
+      # Updated form from ar_virtual.rb, and merged with the code originally in
+      # ActiveRecord.  If the code in ar_virtual.rb is changed, this should
+      # probably also be updated.
+      def preloaders_for_one(association, records, scope)
+        klass_map = records.compact.group_by(&:class)
+
+        loaders = klass_map.keys.group_by { |klass| klass.virtual_includes(association) }.flat_map do |virtuals, klasses|
+          subset = klasses.flat_map { |klass| klass_map[klass] }
+          preload(subset, virtuals)
+        end
+
+        records_with_association = klass_map.select { |k, rs| k.reflect_on_association(association) }.flat_map { |k, rs| rs }
+        if records_with_association.any?
+          # This injects the original code from preloaders_for_one from
+          # ActiveRecord so we can add our own `if` in the middle of it.  The
+          # positive part of the `if` is the only portion of this that has
+          # changed, and the code copied is within the `loaders.concat`.
+          loaders.concat(grouped_records(association, records_with_association).flat_map do |reflection, klasses|
+            klasses.map do |rhs_klass, rs|
+              base_klass = rhs_klass.base_class if rhs_klass.respond_to?(:base_class)
+
+              # Start of new code (1)
+              class_preloader = (class_specific_preloaders[base_klass] || class_specific_preloaders[rhs_klass])
+              loader_klass    = preloader_for(reflection, rs, rhs_klass)
+
+              loader = if class_preloader.kind_of?(ActiveRecord::Relation)
+                         loader_klass.new(rhs_klass, rs, reflection, class_preloader)
+                       else
+                         loader_klass.new(rhs_klass, rs, reflection, scope)
+                       end
+              # End of new code (1)
+
+              loader.run self
+
+              # Start of new code (2)
+              if class_preloader.kind_of?(Array) && class_preloader.count == 2
+                [loader, MiqPreloader.preload(loader.preloaded_records, class_preloader[0], class_preloader[1])]
+              else
+                loader
+              end
+              # End of new code (2)
+            end
+          end)
+        end
+
+        loaders
+      end
+    })
+    preloader.class_specific_preloaders = class_preloaders || {}
+
+    preloader.preload(records, associations, nil)
+  end
 end

spec/lib/miq_preloader_spec.rb

@@ -112,4 +112,63 @@ def preload_and_scope(*args)
       MiqPreloader.preload_and_scope(*args)
     end
   end
+
+  describe ".polymorphic_preload_for_child_classes" do
+    it "preloads polymorphic relationships that are defined" do
+      ems         = FactoryGirl.create(:ems_infra)
+      clusters    = FactoryGirl.create_list(:ems_cluster, 2,
+                                            :ext_management_system => ems)
+      host_group1 = FactoryGirl.create_list(:host, 2,
+                                            :ext_management_system => ems,
+                                            :ems_cluster           => clusters.first)
+      host_group2 = FactoryGirl.create_list(:host, 2,
+                                            :ext_management_system => ems,
+                                            :ems_cluster           => clusters.last)
+
+      ems_rel      = ems.init_relationship
+      cluster_rels = clusters.map { |cluster| cluster.init_relationship(ems_rel) }
+      host_rels1   = host_group1.map { |host| [host, host.init_relationship(cluster_rels.first)] }
+      host_rels2   = host_group2.map { |host| [host, host.init_relationship(cluster_rels.last)] }
+
+      (host_rels1 + host_rels2).each do |(host, host_rel)|
+        FactoryGirl.create_list(:vm, 2, :ext_management_system => ems, :host => host).each do |vm|
+          vm.init_relationship(host_rel)
+        end
+      end
+
+      tree    = ExtManagementSystem.last.fulltree_rels_arranged(:except_type => "VmOrTemplate")
+      records = Relationship.flatten_arranged_rels(tree)
+
+      hosts_scope   = Host.select(Host.arel_table[Arel.star], :v_total_vms)
+      class_loaders = { EmsCluster => [:hosts, hosts_scope], Host => hosts_scope }
+
+      # 4 queries are expected here:
+      #
+      # - 1 for ExtManagementSystem (root)
+      # - 1 for EmsClusters in the tree
+      # - 1 for Hosts in the tree
+      # - 1 for Hosts from relation in EmsClusters
+      #
+      # Since all the hosts in this case are also part of the tree, there are
+      # "duplicate hosts loaded", but that was the nature of this prior to the
+      # change anyway, so this is not new.  This does make it soe that any
+      # hosts accessed through a EmsCluster are preloaded, however, instead of
+      # an N+1.
+      #
+      # In some cases, a ems_metatdata tree might not include all of the
+      # hosts of a EMS, but some still exist as part of the cluster.  This also
+      # makes sure both cases are covered, and the minimal amount of queries
+      # are still executed.
+      #
+      # rubocop:disable Style/BlockDelimiters, Style/MethodCallWithArgsParentheses
+      expect {
+        MiqPreloader.polymorphic_preload_for_child_classes(records, :resource, class_loaders)
+        records.select   { |rel|  rel.resource if rel.resource_type == "Host" }
+               .each     { |rel|  rel.resource.v_total_vms }
+        records.select   { |rel|  rel.resource if rel.resource_type == "EmsCluster" }
+               .flat_map { |rel|  rel.resource.hosts }.each(&:v_total_vms)
+      }.to match_query_limit_of(4)
+      # rubocop:enable Style/BlockDelimiters, Style/MethodCallWithArgsParentheses
+    end
+  end
 end