Add MiqPreloader.polymorphic_preload_for_child_classes

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
+      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
+    end
+  end
 end