Merge branch 'ali/preload-protected-environments-in-env-serializer' into 'master'

Fix remaining N+1 queries in EnvironmentSerializer

See merge request gitlab-org/gitlab!82746

ee/app/models/ee/environment.rb

@@ -110,6 +110,11 @@ module EE
       end
     end
 
+    # Use ProtectedEnvironmentPreloader to fetch associated protected environments in batch
+    def associated_protected_environments=(protected_environments)
+      strong_memoize(:associated_protected_environments) { protected_environments }
+    end
+
     private
 
     def protected_environment_accesses(user)

ee/app/models/preloaders/environments/protected_environment_preloader.rb

+# frozen_string_literal: true
+
+module Preloaders
+  module Environments
+    class ProtectedEnvironmentPreloader
+      attr_reader :environments
+
+      def initialize(environments)
+        @environments = environments
+
+        if environments.map(&:project_id).uniq.size > 1
+          raise 'This preloader supports only environments in the same project'
+        end
+      end
+
+      def execute
+        return if environments.empty?
+
+        project = environments.first.project
+        project_id = project.id
+        group_ids = project.ancestors_upto_ids
+
+        names = environments.map(&:name)
+        tiers = environments.map(&:tier)
+
+        project_protected_environments = ProtectedEnvironment.preload(:deploy_access_levels)
+                                                             .where(project_id: project_id, name: names)
+                                                             .index_by(&:name)
+        group_protected_environments = ProtectedEnvironment.preload(:deploy_access_levels)
+                                                           .where(group_id: group_ids, name: tiers)
+                                                           .index_by(&:name)
+
+        environments.each do |environment|
+          protected_environments ||= []
+          protected_environments << project_protected_environments[environment.name]
+          protected_environments << group_protected_environments[environment.tier]
+          environment.associated_protected_environments = protected_environments.flatten.compact
+        end
+      end
+    end
+  end
+end

ee/app/serializers/ee/environment_serializer.rb

@@ -21,5 +21,18 @@ module EE
     def project_associations
       super.deep_merge(protected_environments: [])
     end
+
+    override :batch_load
+    def batch_load(resource)
+      environments = super
+
+      ::Preloaders::Environments::ProtectedEnvironmentPreloader.new(environments).execute
+
+      environments.each do |environment|
+        # JobEntity loads environment for permission checks in #cancelable?, #retryable?, #playable?
+        environment.last_deployment&.deployable&.persisted_environment = environment
+        environment.upcoming_deployment&.deployable&.persisted_environment = environment
+      end
+    end
   end
 end

ee/spec/models/preloaders/environments/protected_environment_preloader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Preloaders::Environments::ProtectedEnvironmentPreloader, :aggregate_failures do
+  before do
+    stub_licensed_features(protected_environments: true)
+  end
+
+  describe '#initialize' do
+    it 'raises an error if environments belong to more than one project' do
+      expect { described_class.new([create(:environment), create(:environment)]) }
+        .to raise_error('This preloader supports only environments in the same project')
+    end
+  end
+
+  describe '#execute' do
+    let_it_be(:root_group) { create(:group) }
+    let_it_be(:subgroup) { create(:group, parent: root_group) }
+    let_it_be(:project, reload: true) { create(:project, :repository, group: subgroup) }
+    let_it_be(:production, refind: true) { create(:environment, name: 'production', project: project) }
+    let_it_be(:staging, refind: true) { create(:environment, name: 'staging', project: project) }
+    let_it_be(:production_operator) { create(:user, developer_projects: [project]) }
+    let_it_be(:staging_operator) { create(:user, developer_projects: [project]) }
+
+    subject { described_class.new([production, staging]).execute }
+
+    shared_examples 'preloads and associates environments' do
+      it 'preloads protected environments' do
+        subject
+
+        expect { production.protected? }.not_to exceed_query_limit(0)
+        expect { staging.protected? }.not_to exceed_query_limit(0).with_threshold(1) # 1 project load is expected
+
+        expect(production.protected?).to be_truthy
+        expect(staging.protected?).to be_truthy
+      end
+
+      it 'preloads deploy access levels' do
+        subject
+
+        expect { production.protected_by?(production_operator) }.not_to exceed_query_limit(0)
+        expect { staging.protected_by?(staging_operator) }
+          .not_to exceed_query_limit(0)
+          .with_threshold(1) # 1 project load is expected
+      end
+
+      it 'associates protected environments to the correct environment' do
+        subject
+
+        expect(production.protected_by?(production_operator)).to be_truthy
+        expect(staging.protected_by?(staging_operator)).to be_truthy
+        expect(production.protected_from?(staging_operator)).to be_truthy
+        expect(staging.protected_from?(production_operator)).to be_truthy
+      end
+    end
+
+    context 'with project-level protected environments' do
+      before(:all) do
+        create(:protected_environment, project: project, name: production.name,
+               authorize_user_to_deploy: production_operator)
+        create(:protected_environment, project: project, name: staging.name,
+               authorize_user_to_deploy: staging_operator)
+      end
+
+      include_examples 'preloads and associates environments'
+    end
+
+    context 'with group-level protected environments' do
+      before(:all) do
+        create(:protected_environment, :group_level, group: root_group, name: production.name,
+               authorize_user_to_deploy: production_operator)
+        create(:protected_environment, :group_level, group: subgroup, name: staging.name,
+               authorize_user_to_deploy: staging_operator)
+      end
+
+      include_examples 'preloads and associates environments'
+    end
+  end
+end

spec/models/preloaders/environments/deployment_preloader_spec.rb

@@ -6,14 +6,14 @@ 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) }
 
+  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, environment: environment_a.name) }
+  let_it_be(:ci_build_b) { create(:ci_build, user: user, project: project, pipeline: pipeline, environment: environment_a.name) }
+  let_it_be(:ci_build_c) { create(:ci_build, user: user, project: project, pipeline: pipeline, environment: environment_b.name) }
+
   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)

spec/serializers/environment_serializer_spec.rb

@@ -50,7 +50,7 @@ RSpec.describe EnvironmentSerializer do
 
     context 'when there is a single environment' do
       before do
-        create(:environment, name: 'staging')
+        create(:environment, project: project, name: 'staging')
       end
 
       it 'represents one standalone environment' do
@@ -63,8 +63,8 @@ RSpec.describe EnvironmentSerializer do
 
     context 'when there are multiple environments in folder' do
       before do
-        create(:environment, name: 'staging/my-review-1')
-        create(:environment, name: 'staging/my-review-2')
+        create(:environment, project: project, name: 'staging/my-review-1')
+        create(:environment, project: project, name: 'staging/my-review-2')
       end
 
       it 'represents one item that is a folder' do
@@ -78,10 +78,10 @@ RSpec.describe EnvironmentSerializer do
 
     context 'when there are multiple folders and standalone environments' do
       before do
-        create(:environment, name: 'staging/my-review-1')
-        create(:environment, name: 'staging/my-review-2')
-        create(:environment, name: 'production/my-review-3')
-        create(:environment, name: 'testing')
+        create(:environment, project: project, name: 'staging/my-review-1')
+        create(:environment, project: project, name: 'staging/my-review-2')
+        create(:environment, project: project, name: 'production/my-review-3')
+        create(:environment, project: project, name: 'testing')
       end
 
       it 'represents multiple items grouped within folders' do
@@ -124,7 +124,7 @@ RSpec.describe EnvironmentSerializer do
     context 'when resource is paginatable relation' do
       context 'when there is a single environment object in relation' do
         before do
-          create(:environment)
+          create(:environment, project: project)
         end
 
         it 'serializes environments' do
@@ -134,7 +134,7 @@ RSpec.describe EnvironmentSerializer do
 
       context 'when multiple environment objects are serialized' do
         before do
-          create_list(:environment, 3)
+          create_list(:environment, 3, project: project)
         end
 
         it 'serializes appropriate number of objects' do
@@ -159,10 +159,10 @@ RSpec.describe EnvironmentSerializer do
         end
 
         before do
-          create(:environment, name: 'staging/review-1')
-          create(:environment, name: 'staging/review-2')
-          create(:environment, name: 'production/deploy-3')
-          create(:environment, name: 'testing')
+          create(:environment, project: project, name: 'staging/review-1')
+          create(:environment, project: project, name: 'staging/review-2')
+          create(:environment, project: project, name: 'production/deploy-3')
+          create(:environment, project: project, name: 'testing')
         end
 
         it 'paginates grouped items including ordering' do
@@ -189,7 +189,7 @@ RSpec.describe EnvironmentSerializer do
     let(:resource) { Environment.all }
 
     before do
-      create(:environment, name: 'staging/review-1')
+      create(:environment, project: project, name: 'staging/review-1')
       create_environment_with_associations(project)
     end
 

spec/support/shared_examples/serializers/environment_serializer_shared_examples.rb

 # frozen_string_literal: true
 RSpec.shared_examples 'avoid N+1 on environments serialization' do |ee: false|
-  # Investigating in https://gitlab.com/gitlab-org/gitlab/-/issues/353209
-  let(:query_threshold) { 1 + (ee ? 4 : 0) }
-
   it 'avoids N+1 database queries with grouping', :request_store do
     create_environment_with_associations(project)
 
@@ -11,9 +8,7 @@ RSpec.shared_examples 'avoid N+1 on environments serialization' do |ee: false|
     create_environment_with_associations(project)
     create_environment_with_associations(project)
 
-    expect { serialize(grouping: true) }
-      .not_to exceed_query_limit(control.count)
-      .with_threshold(query_threshold)
+    expect { serialize(grouping: true) }.not_to exceed_query_limit(control.count)
   end
 
   it 'avoids N+1 database queries without grouping', :request_store do
@@ -24,9 +19,7 @@ RSpec.shared_examples 'avoid N+1 on environments serialization' do |ee: false|
     create_environment_with_associations(project)
     create_environment_with_associations(project)
 
-    expect { serialize(grouping: false) }
-      .not_to exceed_query_limit(control.count)
-      .with_threshold(query_threshold)
+    expect { serialize(grouping: false) }.not_to exceed_query_limit(control.count)
   end
 
   it 'does not preload for environments that does not exist in the page', :request_store do