Consolidate Protected Environment batch loading logic

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

@@ -16,19 +16,11 @@ module Preloaders
       def execute
         return if environments.empty?
 
-        project = environments.first.project
-        project_id = project.id
-        group_ids = project.ancestors_upto_ids
+        associated_protected_environments =
+          ProtectedEnvironment.for_environments(environments).preload(:deploy_access_levels)
 
-        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)
+        project_protected_environments = associated_protected_environments.select(&:project_level?).index_by(&:name)
+        group_protected_environments = associated_protected_environments.select(&:group_level?).index_by(&:name)
 
         environments.each do |environment|
           protected_environments ||= []

ee/app/models/protected_environment.rb

@@ -59,12 +59,21 @@ class ProtectedEnvironment < ApplicationRecord
 
       key = "protected_environment:for_environment:#{environment.id}"
 
-      ::Gitlab::SafeRequestStore.fetch(key) do
-        from_union([
-          where(project: environment.project_id, name: environment.name),
-          where(group: environment.project.ancestors_upto_ids, name: environment.tier)
-        ])
-      end
+      ::Gitlab::SafeRequestStore.fetch(key) { for_environments([environment]) }
+    end
+
+    def for_environments(environments)
+      raise ArgumentError, 'Environments must be in the same project' if environments.map(&:project_id).uniq.size > 1
+
+      project_id = environments.first.project_id
+      group_ids = environments.first.project.ancestors_upto_ids
+      names = environments.map(&:name)
+      tiers = environments.map(&:tier)
+
+      from_union([
+        where(project: project_id, name: names),
+        where(group: group_ids, name: tiers)
+      ])
     end
   end
 
@@ -81,14 +90,6 @@ class ProtectedEnvironment < ApplicationRecord
     end
   end
 
-  private
-
-  def valid_tier_name
-    unless Environment.tiers[name]
-      errors.add(:name, "must be one of environment tiers: #{Environment.tiers.keys.join(', ')}.")
-    end
-  end
-
   def project_level?
     project_id.present?
   end
@@ -96,4 +97,12 @@ class ProtectedEnvironment < ApplicationRecord
   def group_level?
     group_id.present?
   end
+
+  private
+
+  def valid_tier_name
+    unless Environment.tiers[name]
+      errors.add(:name, "must be one of environment tiers: #{Environment.tiers.keys.join(', ')}.")
+    end
+  end
 end

ee/spec/models/protected_environment_spec.rb

@@ -179,6 +179,38 @@ RSpec.describe ProtectedEnvironment do
     end
   end
 
+  describe '#project_level?' do
+    subject { protected_environment.project_level? }
+
+    context 'for a project-level protected environment' do
+      let(:protected_environment) { create(:protected_environment, :project_level) }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'for a group-level protected environment' do
+      let(:protected_environment) { create(:protected_environment, :group_level) }
+
+      it { is_expected.to be_falsey }
+    end
+  end
+
+  describe '#group_level?' do
+    subject { protected_environment.group_level? }
+
+    context 'for a group-level protected environment' do
+      let(:protected_environment) { create(:protected_environment, :group_level) }
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'for a project-level protected environment' do
+      let(:protected_environment) { create(:protected_environment, :project_level) }
+
+      it { is_expected.to be_falsey }
+    end
+  end
+
   describe '.sorted_by_name' do
     subject(:protected_environments) { described_class.sorted_by_name }
 
@@ -294,7 +326,7 @@ RSpec.describe ProtectedEnvironment do
 
     subject { described_class.for_environment(environment) }
 
-    it { is_expected.to eq([protected_environment]) }
+    it { is_expected.to match_array([protected_environment]) }
 
     it 'caches result', :request_store do
       described_class.for_environment(environment).to_a
@@ -303,31 +335,54 @@ RSpec.describe ProtectedEnvironment do
         .not_to exceed_query_limit(0)
     end
 
-    context 'when environment is a different name' do
-      let!(:environment) { create(:environment, name: 'staging', project: project) }
+    context 'when environment does not exist' do
+      let!(:environment) { }
 
-      it { is_expected.to be_empty }
+      it 'raises an error' do
+        expect { subject }.to raise_error(ArgumentError)
+      end
     end
 
-    context 'when environment exists in a different project' do
-      let!(:environment) { create(:environment, name: 'production', project: create(:project)) }
+    it 'calls .for_environments with the environment' do
+      expect(described_class).to receive(:for_environments).with([environment]).and_call_original
+
+      described_class.for_environment(environment)
+    end
+  end
+
+  describe '.for_environments' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:project, reload: true) { create(:project, group: group) }
+
+    let!(:environments) { [create(:environment, name: 'production', project: project)] }
+    let!(:protected_environment) { create(:protected_environment, name: 'production', project: project) }
+
+    subject { described_class.for_environments(environments) }
+
+    it { is_expected.to match_array([protected_environment]) }
+
+    it 'raises an error if environments belong to more than one project' do
+      expect { described_class.for_environments([create(:environment), create(:environment)]) }
+        .to raise_error(ArgumentError, 'Environments must be in the same project')
+    end
+
+    context 'when environment is a different name' do
+      let!(:environments) { [create(:environment, name: 'staging', project: project)] }
 
       it { is_expected.to be_empty }
     end
 
-    context 'when environment does not exist' do
-      let!(:environment) { }
+    context 'when environment exists in a different project' do
+      let!(:environments) { [create(:environment, name: 'production', project: create(:project))] }
 
-      it 'raises an error' do
-        expect { subject }.to raise_error(ArgumentError)
-      end
+      it { is_expected.to be_empty }
     end
 
     context 'with group-level protected environment' do
       let!(:group_protected_environment) { create(:protected_environment, :production, :group_level, group: group) }
 
       context 'with project-level production environment' do
-        let!(:environment) { create(:environment, :production, project: project) }
+        let!(:environments) { [create(:environment, :production, project: project)] }
 
         it 'has multiple protections' do
           is_expected.to contain_exactly(protected_environment, group_protected_environment)
@@ -337,19 +392,39 @@ RSpec.describe ProtectedEnvironment do
           let!(:protected_environment) { }
 
           it 'has only group-level protection' do
-            is_expected.to eq([group_protected_environment])
+            is_expected.to match_array([group_protected_environment])
           end
         end
       end
 
       context 'with staging environment' do
-        let(:environment) { create(:environment, :staging, project: project) }
+        let(:environments) { [create(:environment, :staging, project: project)] }
 
         it 'does not have any protections' do
           is_expected.to be_empty
         end
       end
     end
+
+    context 'with multiple environments' do
+      let!(:protected_environment) {}
+      let!(:environments) do
+        [
+          create(:environment, name: 'production', project: project),
+          create(:environment, name: 'canary', project: project),
+          create(:environment, name: 'dev', project: project)
+        ]
+      end
+
+      let!(:protected_environments) do
+        [
+          create(:protected_environment, name: 'production', project: project),
+          create(:protected_environment, name: 'canary', project: project)
+        ]
+      end
+
+      it { is_expected.to match_array(protected_environments) }
+    end
   end
 
   def create_deploy_access_level(protected_environment, **opts)