Follow up Protected Environment refactoring

This commit follows up the protected environment
refactoring.

app/models/ci/bridge.rb

@@ -163,7 +163,7 @@ module Ci
     def expanded_environment_name
     end
 
-    def instantized_environment
+    def persisted_environment
     end
 
     def execute_hooks

app/models/ci/build.rb

@@ -90,16 +90,6 @@ module Ci
       end
     end
 
-    # Initializing an object instead of fetching `persisted_environment` for avoiding unnecessary queries.
-    # We're planning to introduce a direct relationship between build and environment
-    # in https://gitlab.com/gitlab-org/gitlab/-/issues/326445 to let us to preload
-    # in batch.
-    def instantized_environment
-      return unless has_environment?
-
-      ::Environment.new(project: self.project, name: self.expanded_environment_name)
-    end
-
     serialize :options # rubocop:disable Cop/ActiveRecordSerialize
     serialize :yaml_variables, Gitlab::Serializer::Ci::Variables # rubocop:disable Cop/ActiveRecordSerialize
 

app/models/ci/processable.rb

@@ -120,7 +120,7 @@ module Ci
       raise NotImplementedError
     end
 
-    def instantized_environment
+    def persisted_environment
       raise NotImplementedError
     end
 

app/services/base_container_service.rb

@@ -18,4 +18,12 @@ class BaseContainerService
     @current_user = current_user
     @params = params.dup
   end
+
+  def project_container?
+    container.is_a?(::Project)
+  end
+
+  def group_container?
+    container.is_a?(::Group)
+  end
 end

ee/app/controllers/projects/protected_environments_controller.rb

@@ -6,7 +6,7 @@ class Projects::ProtectedEnvironmentsController < Projects::ApplicationControlle
   feature_category :continuous_delivery
 
   def create
-    protected_environment = ::ProtectedEnvironments::CreateService.new(@project, current_user, protected_environment_params).execute
+    protected_environment = ::ProtectedEnvironments::CreateService.new(container: @project, current_user: current_user, params: protected_environment_params).execute
 
     if protected_environment.persisted?
       flash[:notice] = s_('ProtectedEnvironment|Your environment has been protected.')
@@ -18,7 +18,7 @@ class Projects::ProtectedEnvironmentsController < Projects::ApplicationControlle
   end
 
   def update
-    result = ::ProtectedEnvironments::UpdateService.new(@project, current_user, protected_environment_params).execute(@protected_environment)
+    result = ::ProtectedEnvironments::UpdateService.new(container: @project, current_user: current_user, params: protected_environment_params).execute(@protected_environment)
 
     if result
       render json: @protected_environment, status: :ok, include: :deploy_access_levels
@@ -28,7 +28,7 @@ class Projects::ProtectedEnvironmentsController < Projects::ApplicationControlle
   end
 
   def destroy
-    result = ::ProtectedEnvironments::DestroyService.new(@project, current_user).execute(@protected_environment)
+    result = ::ProtectedEnvironments::DestroyService.new(container: @project, current_user: current_user).execute(@protected_environment)
 
     if result
       flash[:notice] = s_('ProtectedEnvironment|Your environment has been unprotected')
@@ -40,7 +40,7 @@ class Projects::ProtectedEnvironmentsController < Projects::ApplicationControlle
   end
 
   def search
-    unprotected_environment_names = ::ProtectedEnvironments::SearchService.new(@project, current_user).execute(search_params[:query])
+    unprotected_environment_names = ::ProtectedEnvironments::SearchService.new(container: @project, current_user: current_user).execute(search_params[:query])
 
     render json: unprotected_environment_names, status: :ok
   end

ee/app/models/ee/environment.rb

@@ -82,7 +82,7 @@ module EE
     private
 
     def protected_environment_accesses(user)
-      key = "environment:#{self.project_id}:#{self.name}:for:#{user.id}"
+      key = "environment:#{self.id}:for:#{user.id}"
 
       ::Gitlab::SafeRequestStore.fetch(key) do
         associated_protected_environments.group_by do |pe|

ee/app/models/protected_environment.rb

@@ -28,7 +28,7 @@ class ProtectedEnvironment < ApplicationRecord
     def for_environment(environment)
       raise ArgumentError unless environment.is_a?(::Environment)
 
-      key = "protected_environment:for_environment:#{environment.project_id}:#{environment.name}"
+      key = "protected_environment:for_environment:#{environment.id}"
 
       ::Gitlab::SafeRequestStore.fetch(key) do
         where(project: environment.project_id, name: environment.name)

ee/app/policies/ee/ci/build_policy.rb

@@ -7,11 +7,11 @@ module EE
       prepended do
         # overriding
         condition(:protected_environment) do
-          @subject.instantized_environment&.protected_from?(user)
+          @subject.persisted_environment.try(:protected_from?, user)
         end
 
         condition(:reporter_has_access_to_protected_environment) do
-          @subject.instantized_environment&.protected_by?(user) &&
+          @subject.persisted_environment.try(:protected_by?, user) &&
             can?(:reporter_access, @subject.project)
         end
 

ee/app/services/ee/ci/process_build_service.rb

@@ -6,7 +6,7 @@ module EE
 
       override :enqueue
       def enqueue(build)
-        if build.instantized_environment&.protected_from?(build.user)
+        if build.persisted_environment.try(:protected_from?, build.user)
           return build.drop!(:protected_environment_failure)
         end
 

ee/app/services/protected_environments/base_service.rb

 # frozen_string_literal: true
 module ProtectedEnvironments
-  class BaseService < ::BaseService
+  class BaseService < ::BaseContainerService
     include Gitlab::Utils::StrongMemoize
 
     protected
@@ -25,7 +25,7 @@ module ProtectedEnvironments
       return false unless keys.count == 1
 
       if attribute[:group_id].present?
-        invited_group_ids.include?(attribute[:group_id])
+        qualified_group_ids.include?(attribute[:group_id])
       elsif attribute[:user_id].present?
         qualified_user_ids.include?(attribute[:user_id])
       else
@@ -33,9 +33,13 @@ module ProtectedEnvironments
       end
     end
 
-    def invited_group_ids
-      strong_memoize(:invited_group_ids) do
-        project.invited_groups.pluck_primary_key.to_set
+    def qualified_group_ids
+      strong_memoize(:qualified_group_ids) do
+        if project_container?
+          container.invited_groups.pluck_primary_key.to_set
+        elsif group_container?
+          raise NotImplementedError
+        end
       end
     end
 
@@ -46,10 +50,14 @@ module ProtectedEnvironments
           user_ids
         end
 
-        project.project_authorizations
-          .visible_to_user_and_access_level(user_ids, Gitlab::Access::DEVELOPER)
-          .pluck_user_ids
-          .to_set
+        if project_container?
+          container.project_authorizations
+            .visible_to_user_and_access_level(user_ids, Gitlab::Access::DEVELOPER)
+            .pluck_user_ids
+            .to_set
+        elsif group_container?
+          raise NotImplementedError
+        end
       end
     end
   end

ee/app/services/protected_environments/create_service.rb

@@ -2,7 +2,7 @@
 module ProtectedEnvironments
   class CreateService < ProtectedEnvironments::BaseService
     def execute
-      project.protected_environments.create(sanitized_params)
+      container.protected_environments.create(sanitized_params)
     end
   end
 end

ee/app/services/protected_environments/search_service.rb

@@ -5,9 +5,11 @@ module ProtectedEnvironments
     # Limited to 20 per performance reasons
     # rubocop: disable CodeReuse/ActiveRecord
     def execute(name)
-      project
+      raise NotImplementedError unless project_container?
+
+      container
         .environments
-        .where.not(name: project.protected_environments.select(:name))
+        .where.not(name: container.protected_environments.select(:name))
         .where('environments.name LIKE ?', "#{name}%")
         .order_by_last_deployed_at
         .limit(20)

ee/lib/api/protected_environments.rb

@@ -65,7 +65,7 @@ module API
         # original issue - https://github.com/ruby-grape/grape/issues/1874
         declared_params[:deploy_access_levels_attributes] = declared_params.delete(:deploy_access_levels)
         protected_environment = ::ProtectedEnvironments::CreateService
-                                  .new(user_project, current_user, declared_params).execute
+                                  .new(container: user_project, current_user: current_user, params: declared_params).execute
 
         if protected_environment.persisted?
           present protected_environment, with: ::EE::API::Entities::ProtectedEnvironment, project: user_project
@@ -82,7 +82,7 @@ module API
       end
       delete ':id/protected_environments/:name', requirements: ENVIRONMENT_ENDPOINT_REQUIREMENTS do
         destroy_conditionally!(protected_environment) do
-          ::ProtectedEnvironments::DestroyService.new(user_project, current_user).execute(protected_environment)
+          ::ProtectedEnvironments::DestroyService.new(container: user_project, current_user: current_user).execute(protected_environment)
         end
       end
     end

ee/spec/models/ee/project_group_link_spec.rb

@@ -41,9 +41,9 @@ RSpec.describe ProjectGroupLink do
     context 'protected environments' do
       let!(:protected_environment) do
         ProtectedEnvironments::CreateService.new(
-          project,
-          project.owner,
-          attributes_for(
+          container: project,
+          current_user: project.owner,
+          params: attributes_for(
             :protected_environment,
             deploy_access_levels_attributes: [{ group_id: group.id }, { user_id: user.id }]
           )

ee/spec/serializers/pipeline_serializer_spec.rb

@@ -30,4 +30,28 @@ RSpec.describe PipelineSerializer do
       expect(name).to eq 'bridge'
     end
   end
+
+  describe 'N+1 checks' do
+    let_it_be(:production) { create(:environment, :production, project: project) }
+    let_it_be(:staging) { create(:environment, :staging, project: project) }
+    let_it_be(:protected_production) { create(:protected_environment, project: project, name: production.name) }
+    let_it_be(:protected_staging) { create(:protected_environment, project: project, name: staging.name) }
+
+    context 'with protected environments' do
+      before do
+        stub_licensed_features(protected_environments: true)
+      end
+
+      it 'executes minimal queries to fetch all related protected environments', :request_store do
+        pipeline = create(:ci_pipeline, project: project)
+        create(:ci_build, :manual, pipeline: pipeline, environment: production.name)
+        create(:ci_build, :manual, pipeline: pipeline, environment: staging.name)
+        create(:ci_build, :scheduled, pipeline: pipeline, environment: production.name)
+        create(:ci_build, :scheduled, pipeline: pipeline, environment: staging.name)
+
+        expect { serializer.represent(Ci::Pipeline.all, preload: true) }
+          .not_to exceed_query_limit(2).for_query /SELECT "protected_environments".*/
+      end
+    end
+  end
 end

ee/spec/services/protected_environments/create_service_spec.rb

@@ -11,7 +11,7 @@ RSpec.describe ProtectedEnvironments::CreateService, '#execute' do
                    deploy_access_levels_attributes: [{ access_level: maintainer_access }])
   end
 
-  subject { described_class.new(project, user, params).execute }
+  subject { described_class.new(container: project, current_user: user, params: params).execute }
 
   context 'with valid params' do
     it { is_expected.to be_truthy }

ee/spec/services/protected_environments/destroy_service_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe ProtectedEnvironments::DestroyService, '#execute' do
   let!(:protected_environment) { create(:protected_environment, project: project) }
   let(:deploy_access_level) { protected_environment.deploy_access_levels.first }
 
-  subject { described_class.new(project, user).execute(protected_environment) }
+  subject { described_class.new(container: project, current_user: user).execute(protected_environment) }
 
   context 'when the Protected Environment is deleted' do
     it { is_expected.to be_truthy }

ee/spec/services/protected_environments/search_service_spec.rb

@@ -5,7 +5,7 @@ RSpec.describe ProtectedEnvironments::SearchService, '#execute' do
   let(:project) { create(:project) }
   let(:user) { create(:user) }
 
-  subject { described_class.new(project, user).execute(environment_name) }
+  subject { described_class.new(container: project, current_user: user).execute(environment_name) }
 
   before do
     %w(production staging review/app_1 review/app_2 test canary).each do |environment_name|

ee/spec/services/protected_environments/update_service_spec.rb

@@ -17,7 +17,7 @@ RSpec.describe ProtectedEnvironments::UpdateService, '#execute' do
     }
   end
 
-  subject { described_class.new(project, user, params).execute(protected_environment) }
+  subject { described_class.new(container: project, current_user: user, params: params).execute(protected_environment) }
 
   before do
     deploy_access_level

lib/gitlab/ci/pipeline/preloader.rb

@@ -20,6 +20,7 @@ module Gitlab
               preloader.preload_ref_commits
               preloader.preload_pipeline_warnings
               preloader.preload_stages_warnings
+              preloader.preload_persisted_environments
             end
           end
         end
@@ -54,6 +55,13 @@ module Gitlab
         def preload_stages_warnings
           @pipeline.stages.each { |stage| stage.number_of_warnings }
         end
+
+        # This batch loads the associated environments of multiple actions (builds)
+        # that can't use `preload` due to the indirect relationship.
+        def preload_persisted_environments
+          @pipeline.scheduled_actions.each { |action| action.persisted_environment }
+          @pipeline.manual_actions.each { |action| action.persisted_environment }
+        end
       end
     end
   end

spec/lib/gitlab/ci/pipeline/preloader_spec.rb

@@ -5,9 +5,11 @@ require 'spec_helper'
 RSpec.describe Gitlab::Ci::Pipeline::Preloader do
   let(:stage) { double(:stage) }
   let(:commit) { double(:commit) }
+  let(:scheduled_action) { double(:scheduled_action) }
+  let(:manual_action) { double(:manual_action) }
 
   let(:pipeline) do
-    double(:pipeline, commit: commit, stages: [stage])
+    double(:pipeline, commit: commit, stages: [stage], scheduled_actions: [scheduled_action], manual_actions: [manual_action])
   end
 
   describe '.preload!' do
@@ -33,6 +35,8 @@ RSpec.describe Gitlab::Ci::Pipeline::Preloader do
       expect(pipeline).to receive(:lazy_ref_commit)
       expect(pipeline).to receive(:number_of_warnings)
       expect(stage).to receive(:number_of_warnings)
+      expect(scheduled_action).to receive(:persisted_environment)
+      expect(manual_action).to receive(:persisted_environment)
 
       described_class.preload!([pipeline])
     end
@@ -42,6 +46,8 @@ RSpec.describe Gitlab::Ci::Pipeline::Preloader do
       allow(pipeline).to receive(:lazy_ref_commit)
       allow(pipeline).to receive(:number_of_warnings)
       allow(stage).to receive(:number_of_warnings)
+      allow(scheduled_action).to receive(:persisted_environment)
+      allow(manual_action).to receive(:persisted_environment)
 
       pipelines = [pipeline, pipeline]
 

spec/serializers/pipeline_serializer_spec.rb

@@ -212,18 +212,17 @@ RSpec.describe PipelineSerializer do
       context 'with build environments' do
         let(:ref) { 'feature' }
 
-        it 'verifies number of queries', :request_store do
-          stub_licensed_features(protected_environments: true)
+        let_it_be(:production) { create(:environment, :production, project: project) }
+        let_it_be(:staging) { create(:environment, :staging, project: project) }
 
-          env = create(:environment, project: project)
-          create(:ci_build, :scheduled, project: project, environment: env.name)
-          create(:ci_build, :scheduled, project: project, environment: env.name)
-          create(:ci_build, :scheduled, project: project, environment: env.name)
+        it 'executes one query to fetch all related environments', :request_store do
+          pipeline = create(:ci_pipeline, project: project)
+          create(:ci_build, :manual, pipeline: pipeline, environment: production.name)
+          create(:ci_build, :manual, pipeline: pipeline, environment: staging.name)
+          create(:ci_build, :scheduled, pipeline: pipeline, environment: production.name)
+          create(:ci_build, :scheduled, pipeline: pipeline, environment: staging.name)
 
-          recorded = ActiveRecord::QueryRecorder.new { subject }
-          expected_queries = Gitlab.ee? ? 56 : 52
-          expect(recorded.count).to be_within(1).of(expected_queries)
-          expect(recorded.cached_count).to eq(0)
+          expect { subject }.not_to exceed_query_limit(1).for_query /SELECT "environments".*/
         end
       end