Merge branch 'database-query-timeout-fix-environment-page' into 'master'

Resolve DB Query Timeout on  Environment page

See merge request gitlab-org/gitlab!75767

app/models/deployment.rb

@@ -8,6 +8,7 @@ class Deployment < ApplicationRecord
   include Importable
   include Gitlab::Utils::StrongMemoize
   include FastDestroyAll
+  include FromUnion
 
   StatusUpdateError = Class.new(StandardError)
   StatusSyncError = Class.new(StandardError)

app/models/preloaders/environments/deployment_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  module Environments
+    # This class is to batch-load deployments of multiple environments.
+    # The deployments to batch-load are fetched using UNION of N selects in a single query instead of default scoping with `IN (environment_id1, environment_id2 ...)`.
+    # See https://gitlab.com/gitlab-org/gitlab/-/issues/345672#note_761852224 for more details.
+    class DeploymentPreloader
+      attr_reader :environments
+
+      def initialize(environments)
+        @environments = environments
+      end
+
+      def execute_with_union(association_name, association_attributes)
+        load_deployment_association(association_name, association_attributes)
+      end
+
+      private
+
+      def load_deployment_association(association_name, association_attributes)
+        return unless environments.present?
+
+        union_arg = environments.inject([]) do |result, environment|
+          result << environment.association(association_name).scope
+        end
+
+        union_sql = Deployment.from_union(union_arg).to_sql
+
+        deployments = Deployment
+                        .from("(#{union_sql}) #{::Deployment.table_name}")
+                        .preload(association_attributes)
+
+        deployments_by_environment_id = deployments.index_by(&:environment_id)
+
+        environments.each do |environment|
+          environment.association(association_name).target = deployments_by_environment_id[environment.id]
+          environment.association(association_name).loaded!
+        end
+      end
+    end
+  end
+end

app/serializers/environment_serializer.rb

@@ -52,7 +52,17 @@ class EnvironmentSerializer < BaseSerializer
   end
 
   def batch_load(resource)
+    if ::Feature.enabled?(:custom_preloader_for_deployments, default_enabled: :yaml)
+      resource = resource.preload(environment_associations.except(:last_deployment, :upcoming_deployment))
+
+      Preloaders::Environments::DeploymentPreloader.new(resource)
+        .execute_with_union(:last_deployment, deployment_associations)
+
+      Preloaders::Environments::DeploymentPreloader.new(resource)
+        .execute_with_union(:upcoming_deployment, deployment_associations)
+    else
       resource = resource.preload(environment_associations)
+    end
 
     resource.all.to_a.tap do |environments|
       environments.each do |environment|

config/feature_flags/development/custom_preloader_for_deployments.yml

+---
+name: custom_preloader_for_deployments
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/75767
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/348289
+milestone: '14.7'
+type: development
+group: group::release
+default_enabled: false

spec/models/preloaders/environments/deployment_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::Environments::DeploymentPreloader do
+  let_it_be(:user) { create(:user) }
+  let_it_be(:project, reload: true) { create(:project, :repository) }
+
+  let_it_be(:pipeline) { create(:ci_pipeline, user: user, project: project, sha: project.commit.sha) }
+  let_it_be(:ci_build_a) { create(:ci_build, user: user, project: project, pipeline: pipeline) }
+  let_it_be(:ci_build_b) { create(:ci_build, user: user, project: project, pipeline: pipeline) }
+  let_it_be(:ci_build_c) { create(:ci_build, user: user, project: project, pipeline: pipeline) }
+
+  let_it_be(:environment_a) { create(:environment, project: project, state: :available) }
+  let_it_be(:environment_b) { create(:environment, project: project, state: :available) }
+
+  before do
+    create(:deployment, :success, project: project, environment: environment_a, deployable: ci_build_a)
+    create(:deployment, :success, project: project, environment: environment_a, deployable: ci_build_b)
+    create(:deployment, :success, project: project, environment: environment_b, deployable: ci_build_c)
+  end
+
+  def preload_association(association_name)
+    described_class.new(project.environments)
+      .execute_with_union(association_name, deployment_associations)
+  end
+
+  def deployment_associations
+    {
+      user: [],
+      deployable: {
+        pipeline: {
+          manual_actions: []
+        }
+      }
+    }
+  end
+
+  it 'does not trigger N+1 queries' do
+    control = ActiveRecord::QueryRecorder.new { preload_association(:last_deployment) }
+
+    ci_build_d = create(:ci_build, user: user, project: project, pipeline: pipeline)
+    create(:deployment, :success, project: project, environment: environment_b, deployable: ci_build_d)
+
+    expect { preload_association(:last_deployment) }.not_to exceed_query_limit(control)
+  end
+
+  it 'batch loads the dependent associations' do
+    preload_association(:last_deployment)
+
+    expect do
+      project.environments.first.last_deployment.deployable.pipeline.manual_actions
+    end.not_to exceed_query_limit(0)
+  end
+
+  # Example query scoped with IN clause for `last_deployment` association preload:
+  # SELECT DISTINCT ON (environment_id) deployments.* FROM "deployments" WHERE "deployments"."status" IN (1, 2, 3, 4, 6) AND "deployments"."environment_id" IN (35, 34, 33) ORDER BY environment_id, deployments.id DESC
+  it 'avoids scoping with IN clause during preload' do
+    control = ActiveRecord::QueryRecorder.new { preload_association(:last_deployment) }
+
+    default_preload_query = control.occurrences_by_line_method.first[1][:occurrences].any? { |i| i.include?('"deployments"."environment_id" IN') }
+
+    expect(default_preload_query).to be(false)
+  end
+end

spec/serializers/environment_serializer_spec.rb

@@ -185,6 +185,42 @@ RSpec.describe EnvironmentSerializer do
     end
   end
 
+  context 'batching loading' do
+    let(:resource) { Environment.all }
+
+    before do
+      create(:environment, name: 'staging/review-1')
+      create_environment_with_associations(project)
+    end
+
+    it 'uses the custom preloader service' do
+      expect_next_instance_of(Preloaders::Environments::DeploymentPreloader) do |preloader|
+        expect(preloader).to receive(:execute_with_union).with(:last_deployment, hash_including(:deployable)).and_call_original
+      end
+
+      expect_next_instance_of(Preloaders::Environments::DeploymentPreloader) do |preloader|
+        expect(preloader).to receive(:execute_with_union).with(:upcoming_deployment, hash_including(:deployable)).and_call_original
+      end
+
+      json
+    end
+
+    # Including for test coverage pipeline failure, remove along with feature flag.
+    context 'when custom preload feature is disabled' do
+      before do
+        Feature.disable(:custom_preloader_for_deployments)
+      end
+
+      it 'avoids N+1 database queries' do
+        control_count = ActiveRecord::QueryRecorder.new { json }.count
+
+        create_environment_with_associations(project)
+
+        expect { json }.not_to exceed_query_limit(control_count)
+      end
+    end
+  end
+
   def create_environment_with_associations(project)
     create(:environment, project: project).tap do |environment|
       create(:deployment, :success, environment: environment, project: project)