Unify CommitStatus preloaders

We have a few places where we preload relations for commit status
objects. Ci::Build, Ci::Bridge and GenericCommitStatus have different
relations, so we have conditions where we preload them.

In this commit, we are unifying this logic.

app/graphql/types/ci/stage_type.rb

@@ -55,18 +55,10 @@ module Types
       def jobs_for_pipeline(pipeline, stage_ids, include_needs)
         jobs = pipeline.statuses.latest.where(stage_id: stage_ids)
 
-        common_relations = [:project]
-        common_relations << :needs if include_needs
+        preloaded_relations = [:project, :metadata, :job_artifacts, :downstream_pipeline]
+        preloaded_relations << :needs if include_needs
 
-        preloaders = {
-          ::Ci::Build => [:metadata, :job_artifacts],
-          ::Ci::Bridge => [:metadata, :downstream_pipeline],
-          ::GenericCommitStatus => []
-        }
-
-        preloaders.each do |klass, relations|
-          ActiveRecord::Associations::Preloader.new.preload(jobs.select { |job| job.is_a?(klass) }, relations + common_relations)
-        end
+        Preloaders::CommitStatusPreloader.new(jobs).execute(preloaded_relations)
 
         jobs.group_by(&:stage_id)
       end

app/models/ci/pipeline.rb

@@ -589,13 +589,11 @@ module Ci
     end
 
     def cancel_running(retries: 1)
-      commit_status_relations = [:project, :pipeline]
-      ci_build_relations = [:deployment, :taggings]
+      preloaded_relations = [:project, :pipeline, :deployment, :taggings]
 
       retry_lock(cancelable_statuses, retries, name: 'ci_pipeline_cancel_running') do |cancelables|
         cancelables.find_in_batches do |batch|
-          ActiveRecord::Associations::Preloader.new.preload(batch, commit_status_relations)
-          ActiveRecord::Associations::Preloader.new.preload(batch.select { |job| job.is_a?(Ci::Build) }, ci_build_relations)
+          Preloaders::CommitStatusPreloader.new(batch).execute(preloaded_relations)
 
           batch.each do |job|
             yield(job) if block_given?

app/models/preloaders/commit_status_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  class CommitStatusPreloader
+    CLASSES = [::Ci::Build, ::Ci::Bridge, ::GenericCommitStatus].freeze
+
+    def initialize(statuses)
+      @statuses = statuses
+    end
+
+    def execute(relations)
+      preloader = ActiveRecord::Associations::Preloader.new
+
+      CLASSES.each do |klass|
+        preloader.preload(objects(klass), associations(klass, relations))
+      end
+    end
+
+    private
+
+    def objects(klass)
+      @statuses.select { |job| job.is_a?(klass) }
+    end
+
+    def associations(klass, relations)
+      klass.reflections.keys.map(&:to_sym) & relations.map(&:to_sym)
+    end
+  end
+end

app/presenters/ci/legacy_stage_presenter.rb

@@ -15,18 +15,9 @@ module Ci
     private
 
     def preload_statuses(statuses)
-      loaded_statuses = statuses.load
-      statuses.tap do |statuses|
-        # rubocop: disable CodeReuse/ActiveRecord
-        ActiveRecord::Associations::Preloader.new.preload(preloadable_statuses(loaded_statuses), %w[tags job_artifacts_archive metadata])
-        # rubocop: enable CodeReuse/ActiveRecord
-      end
-    end
+      Preloaders::CommitStatusPreloader.new(statuses).execute(Ci::StagePresenter::PRELOADED_RELATIONS)
 
-    def preloadable_statuses(statuses)
-      statuses.reject do |status|
-        status.instance_of?(::GenericCommitStatus) || status.instance_of?(::Ci::Bridge)
-      end
+      statuses
     end
   end
 end

app/presenters/ci/stage_presenter.rb

@@ -4,6 +4,8 @@ module Ci
   class StagePresenter < Gitlab::View::Presenter::Delegated
     presents :stage
 
+    PRELOADED_RELATIONS = [:pipeline, :metadata, :tags, :job_artifacts_archive, :downstream_pipeline].freeze
+
     def latest_ordered_statuses
       preload_statuses(stage.statuses.latest_ordered)
     end
@@ -15,21 +17,7 @@ module Ci
     private
 
     def preload_statuses(statuses)
-      common_relations = [:pipeline]
-
-      preloaders = {
-        ::Ci::Build => [:metadata, :tags, :job_artifacts_archive],
-        ::Ci::Bridge => [:metadata, :downstream_pipeline],
-        ::GenericCommitStatus => []
-      }
-
-      # rubocop: disable CodeReuse/ActiveRecord
-      preloaders.each do |klass, relations|
-        ActiveRecord::Associations::Preloader
-          .new
-          .preload(statuses.select { |job| job.is_a?(klass) }, relations + common_relations)
-      end
-      # rubocop: enable CodeReuse/ActiveRecord
+      Preloaders::CommitStatusPreloader.new(statuses).execute(PRELOADED_RELATIONS)
 
       statuses
     end

app/services/ci/drop_pipeline_service.rb

@@ -2,8 +2,7 @@
 
 module Ci
   class DropPipelineService
-    PRELOADED_COMMIT_STATUS_RELATIONS = [:project, :pipeline].freeze
-    PRELOADED_CI_BUILD_RELATIONS = [:metadata, :deployment, :taggings].freeze
+    PRELOADED_RELATIONS = [:project, :pipeline, :metadata, :deployment, :taggings].freeze
 
     # execute service asynchronously for each cancelable pipeline
     def execute_async_for_all(pipelines, failure_reason, context_user)
@@ -30,11 +29,8 @@ module Ci
 
     private
 
-    # rubocop: disable CodeReuse/ActiveRecord
     def preload_associations_for_drop(commit_status_batch)
-      ActiveRecord::Associations::Preloader.new.preload(commit_status_batch, PRELOADED_COMMIT_STATUS_RELATIONS)
-      ActiveRecord::Associations::Preloader.new.preload(commit_status_batch.select { |job| job.is_a?(Ci::Build) }, PRELOADED_CI_BUILD_RELATIONS)
+      Preloaders::CommitStatusPreloader.new(commit_status_batch).execute(PRELOADED_RELATIONS)
     end
-    # rubocop: enable CodeReuse/ActiveRecord
   end
 end

spec/models/preloaders/commit_status_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::CommitStatusPreloader do
+  let_it_be(:pipeline) { create(:ci_pipeline) }
+
+  let_it_be(:build1) { create(:ci_build, :tags, pipeline: pipeline) }
+  let_it_be(:build2) { create(:ci_build, :tags, pipeline: pipeline) }
+  let_it_be(:bridge1) { create(:ci_bridge, pipeline: pipeline) }
+  let_it_be(:bridge2) { create(:ci_bridge, pipeline: pipeline) }
+  let_it_be(:generic_commit_status1) { create(:generic_commit_status, pipeline: pipeline) }
+  let_it_be(:generic_commit_status2) { create(:generic_commit_status, pipeline: pipeline) }
+
+  describe '#execute' do
+    let(:relations) { %i[pipeline metadata tags job_artifacts_archive downstream_pipeline] }
+    let(:statuses) { CommitStatus.where(commit_id: pipeline.id).all }
+
+    subject(:execute) { described_class.new(statuses).execute(relations) }
+
+    it 'prevents N+1 for specified relations', :use_sql_query_cache do
+      execute
+
+      control_count = ActiveRecord::QueryRecorder.new(skip_cached: false) do
+        call_each_relation(statuses.sample(3))
+      end
+
+      expect do
+        call_each_relation(statuses)
+      end.to issue_same_number_of_queries_as(control_count)
+    end
+
+    private
+
+    def call_each_relation(statuses)
+      statuses.each do |status|
+        relations.each { |relation| status.public_send(relation) if status.respond_to?(relation) }
+      end
+    end
+  end
+end