Add GraphQL aggregate to prevent N+1 query on DAST profiles

https://gitlab.com/gitlab-org/gitlab/-/issues/324382

ee/app/graphql/ee/gitlab_schema.rb

@@ -9,6 +9,7 @@ module EE
       lazy_resolve ::Gitlab::Graphql::Aggregations::Issues::LazyBlockAggregate, :block_aggregate
       lazy_resolve ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate, :execute
       lazy_resolve ::Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate, :execute
+      lazy_resolve ::Gitlab::Graphql::Aggregations::SecurityOrchestrationPolicies::LazyDastProfileAggregate, :execute
     end
   end
 end

ee/app/graphql/types/dast_scanner_profile_type.rb

@@ -41,12 +41,18 @@ module Types
           description: 'Relative web path to the edit page of a scanner profile.'
 
     field :referenced_in_security_policies, [GraphQL::STRING_TYPE], null: true,
-          complexity: 10,
           calls_gitaly: true,
           description: 'List of security policy names that are referencing given project.'
 
     def edit_path
       Rails.application.routes.url_helpers.edit_project_security_configuration_dast_scans_dast_scanner_profile_path(object.project, object)
     end
+
+    def referenced_in_security_policies
+      ::Gitlab::Graphql::Aggregations::SecurityOrchestrationPolicies::LazyDastProfileAggregate.new(
+        context,
+        object
+      )
+    end
   end
 end

ee/app/graphql/types/dast_site_profile_type.rb

@@ -51,7 +51,6 @@ module Types
           description: 'Normalized URL of the target to be scanned.'
 
     field :referenced_in_security_policies, [GraphQL::STRING_TYPE], null: true,
-          complexity: 10,
           calls_gitaly: true,
           description: 'List of security policy names that are referencing given project.'
 
@@ -84,5 +83,12 @@ module Types
     def normalized_target_url
       DastSiteValidation.get_normalized_url_base(object.dast_site.url)
     end
+
+    def referenced_in_security_policies
+      ::Gitlab::Graphql::Aggregations::SecurityOrchestrationPolicies::LazyDastProfileAggregate.new(
+        context,
+        object
+      )
+    end
   end
 end

ee/app/models/security/orchestration_policy_configuration.rb

@@ -23,6 +23,8 @@ module Security
     validates :project, presence: true, uniqueness: true
     validates :security_policy_management_project, presence: true
 
+    scope :for_project, -> (project_id) { where(project_id: project_id) }
+
     def enabled?
       ::Feature.enabled?(:security_orchestration_policies_configuration, project)
     end
@@ -92,10 +94,15 @@ module Security
     end
 
     def policy_hash
-      blob_data = policy_repo.blob_data_at(default_branch_or_main, POLICY_PATH)
-      return if blob_data.blank?
+      return if policy_blob.blank?
+
+      Gitlab::Config::Loader::Yaml.new(policy_blob).load!
+    end
 
-      Gitlab::Config::Loader::Yaml.new(blob_data).load!
+    def policy_blob
+      strong_memoize(:policy_blob) do
+        policy_repo.blob_data_at(default_branch_or_main, POLICY_PATH)
+      end
     end
 
     def applicable_for_branch?(policy, ref)

ee/changelogs/unreleased/324382-add-aggregate-to-limit-calls-to-gitaly-to-fetch-policies.yml

+---
+title: Add GraphQL aggregate to prevent N+1 query on DAST profiles
+merge_request: 61024
+author:
+type: performance

ee/lib/gitlab/graphql/aggregations/security_orchestration_policies/lazy_dast_profile_aggregate.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Aggregations
+      module SecurityOrchestrationPolicies
+        class LazyDastProfileAggregate
+          include ::Gitlab::Graphql::Deferred
+
+          attr_reader :dast_profile, :lazy_state
+
+          def initialize(query_ctx, dast_profile)
+            raise ArgumentError, 'only DastSiteProfile or DastScannerProfile are allowed' if !dast_profile.is_a?(DastSiteProfile) && !dast_profile.is_a?(DastScannerProfile)
+
+            @dast_profile = Gitlab::Graphql::Lazy.force(dast_profile)
+            # Initialize the loading state for this query,
+            # or get the previously-initiated state
+            @lazy_state = query_ctx[:lazy_dast_profile_in_policies_aggregate] ||= {
+              dast_pending_profiles: [],
+              loaded_objects: {}
+            }
+            # Register this ID to be loaded later:
+            @lazy_state[:dast_pending_profiles] << dast_profile
+          end
+
+          # Return the loaded record, hitting the database if needed
+          def execute
+            # Check if the record was already loaded
+            if @lazy_state[:dast_pending_profiles].present?
+              load_records_into_loaded_objects
+            end
+
+            @lazy_state[:loaded_objects][@dast_profile]
+          end
+
+          private
+
+          def load_records_into_loaded_objects
+            # The record hasn't been loaded yet, so
+            # hit the database with all pending IDs to prevent N+1
+            profiles_by_project_id = @lazy_state[:dast_pending_profiles].group_by(&:project_id)
+            policy_configurations = Security::OrchestrationPolicyConfiguration.for_project(profiles_by_project_id.keys).index_by(&:project_id)
+
+            profiles_by_project_id.each do |project_id, dast_pending_profiles|
+              dast_pending_profiles.each do |profile|
+                @lazy_state[:loaded_objects][profile] = active_policy_names_for_profile(policy_configurations[project_id], profile)
+              end
+            end
+
+            @lazy_state[:dast_pending_profiles].clear
+          end
+
+          def active_policy_names_for_profile(policy_configuration, profile)
+            return [] if policy_configuration.blank?
+
+            case profile
+            when DastSiteProfile
+              policy_configuration.active_policy_names_with_dast_site_profile(profile.name)
+            when DastScannerProfile
+              policy_configuration.active_policy_names_with_dast_scanner_profile(profile.name)
+            end
+          end
+        end
+      end
+    end
+  end
+end

ee/spec/graphql/types/dast_scanner_profile_type_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe GitlabSchema.types['DastScannerProfile'] do
+  include RepoHelpers
+
   let_it_be(:dast_scanner_profile) { create(:dast_scanner_profile) }
   let_it_be(:project) { dast_scanner_profile.project }
   let_it_be(:user) { create(:user) }
@@ -40,7 +42,7 @@ RSpec.describe GitlabSchema.types['DastScannerProfile'] do
       %(
         query project($fullPath: ID!) {
           project(fullPath: $fullPath) {
-            dastScannerProfiles(first: 1) {
+            dastScannerProfiles {
               nodes {
                 id
                 globalId
@@ -67,5 +69,53 @@ RSpec.describe GitlabSchema.types['DastScannerProfile'] do
 
       it { is_expected.to eq(dast_scanner_profile.name) }
     end
+
+    context 'when security policies are enabled' do
+      let_it_be(:policies_project) { create(:project, :repository) }
+      let_it_be(:security_orchestration_policy_configuration) { create(:security_orchestration_policy_configuration, project: project, security_policy_management_project: policies_project) }
+
+      let_it_be(:policy_yml) do
+        <<-EOS
+        scan_execution_policy:
+        - name: Run DAST in every pipeline
+          description: This policy enforces to run DAST for every pipeline within the project
+          enabled: true
+          rules:
+          - type: pipeline
+            branches:
+            - "master"
+          actions:
+          - scan: dast
+            site_profile: Site Profile
+            scanner_profile: #{dast_scanner_profile.name}
+          - scan: dast
+            site_profile: Site Profile 2
+            scanner_profile: Scanner Profile 2
+        - name: Run DAST in every pipeline 2
+          description: This policy enforces to run DAST for every pipeline within the project
+          enabled: true
+          rules:
+          - type: pipeline
+            branches:
+            - "master"
+          actions:
+          - scan: dast
+            site_profile: Site Profile 3
+            scanner_profile: Scanner Profile 3
+          - scan: dast
+            site_profile: Site Profile 4
+            scanner_profile: Scanner Profile 4
+        EOS
+      end
+
+      before do
+        create_list(:dast_scanner_profile, 30, project: project)
+        create_file_in_repo(policies_project, 'master', 'master', Security::OrchestrationPolicyConfiguration::POLICY_PATH, policy_yml)
+      end
+
+      it 'only calls Gitaly twice when multiple profiles are present', :request_store do
+        expect { response }.to change { Gitlab::GitalyClient.get_request_count }.by(2)
+      end
+    end
   end
 end

ee/spec/graphql/types/dast_site_profile_type_spec.rb

@@ -4,6 +4,7 @@ require 'spec_helper'
 
 RSpec.describe GitlabSchema.types['DastSiteProfile'] do
   include GraphqlHelpers
+  include RepoHelpers
 
   let_it_be(:project) { create(:project) }
   let_it_be(:user) { create(:user, developer_projects: [project]) }
@@ -19,7 +20,7 @@ RSpec.describe GitlabSchema.types['DastSiteProfile'] do
   specify { expect(described_class).to expose_permissions_using(Types::PermissionTypes::DastSiteProfile) }
 
   it { expect(described_class).to have_graphql_fields(fields) }
-  it { expect(described_class).to have_graphql_field(:referenced_in_security_policies, calls_gitaly?: true, complexity: 10) }
+  it { expect(described_class).to have_graphql_field(:referenced_in_security_policies, calls_gitaly?: true) }
 
   describe 'id field' do
     it 'is the global id' do
@@ -136,8 +137,89 @@ RSpec.describe GitlabSchema.types['DastSiteProfile'] do
   end
 
   describe 'referencedInSecurityPolicies field' do
-    it 'is the policies' do
-      expect(resolve_field(:referenced_in_security_policies, object, current_user: user)).to eq(object.referenced_in_security_policies)
+    it 'is the lazy aggregate that is resolved to policies', :aggregate_failures do
+      field_value = resolve_field(:referenced_in_security_policies, object, current_user: user)
+
+      expect(field_value).to be_a(GraphQL::Execution::Lazy)
+      expect(field_value.value).to eq(object.referenced_in_security_policies)
+    end
+  end
+
+  describe 'dast_site_profiles' do
+    subject(:response) do
+      GitlabSchema.execute(
+        query,
+        context: {
+          current_user: user
+        },
+        variables: {
+          fullPath: project.full_path
+        }
+      ).as_json
+    end
+
+    let(:query) do
+      %(
+        query project($fullPath: ID!) {
+          project(fullPath: $fullPath) {
+            dastSiteProfiles {
+              nodes {
+                id
+                profileName
+                referencedInSecurityPolicies
+              }
+            }
+          }
+        }
+      )
+    end
+
+    context 'when security policies are enabled' do
+      let_it_be(:policies_project) { create(:project, :repository) }
+      let_it_be(:security_orchestration_policy_configuration) { create(:security_orchestration_policy_configuration, project: project, security_policy_management_project: policies_project) }
+
+      let_it_be(:policy_yml) do
+        <<-EOS
+        scan_execution_policy:
+        - name: Run DAST in every pipeline
+          description: This policy enforces to run DAST for every pipeline within the project
+          enabled: true
+          rules:
+          - type: pipeline
+            branches:
+            - "master"
+          actions:
+          - scan: dast
+            site_profile: Site Profile
+            scanner_profile: Scanner Profile
+          - scan: dast
+            site_profile: Site Profile 2
+            scanner_profile: Scanner Profile 2
+        - name: Run DAST in every pipeline 2
+          description: This policy enforces to run DAST for every pipeline within the project
+          enabled: true
+          rules:
+          - type: pipeline
+            branches:
+            - "master"
+          actions:
+          - scan: dast
+            site_profile: Site Profile 3
+            scanner_profile: Scanner Profile 3
+          - scan: dast
+            site_profile: Site Profile 4
+            scanner_profile: Scanner Profile 4
+        EOS
+      end
+
+      before do
+        create_list(:dast_site_profile, 30, project: project)
+        create_file_in_repo(policies_project, 'master', 'master', Security::OrchestrationPolicyConfiguration::POLICY_PATH, policy_yml)
+      end
+
+      it 'only calls Gitaly twice when multiple profiles are present', :request_store do
+        expect { response }.to change { Gitlab::GitalyClient.get_request_count }.by(2)
+      end
     end
   end
 end

ee/spec/lib/gitlab/graphql/aggregations/security_orchestration_policies/lazy_dast_profile_aggregate_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Aggregations::SecurityOrchestrationPolicies::LazyDastProfileAggregate do
+  let(:query_ctx) do
+    {}
+  end
+
+  let_it_be(:dast_site_profile) { create(:dast_site_profile) }
+  let_it_be(:other_dast_site_profile) { create(:dast_site_profile, project: dast_site_profile.project) }
+
+  let_it_be(:dast_scanner_profile) { create(:dast_scanner_profile, project: dast_site_profile.project) }
+  let_it_be(:other_dast_scanner_profile) { create(:dast_scanner_profile, project: dast_site_profile.project) }
+
+  let(:dast_profile) { dast_site_profile }
+  let(:other_dast_profile) { other_dast_site_profile }
+
+  describe '#initialize' do
+    it 'adds the dast_profile to the lazy state' do
+      subject = described_class.new(query_ctx, dast_profile)
+
+      expect(subject.lazy_state[:dast_pending_profiles]).to eq [dast_profile]
+      expect(subject.dast_profile).to eq dast_profile
+    end
+
+    it 'uses lazy_dast_profile_in_policies_aggregate to collect aggregates' do
+      subject = described_class.new({ lazy_dast_profile_in_policies_aggregate: { dast_pending_profiles: [other_dast_profile], loaded_objects: {} } }, dast_profile)
+
+      expect(subject.lazy_state[:dast_pending_profiles]).to match_array [other_dast_profile, dast_profile]
+      expect(subject.dast_profile).to eq dast_profile
+    end
+
+    it 'raises ArgumentError when is not DastSiteProfile or DastScannerProfile' do
+      expect { described_class.new(query_ctx, Project.new) }.to raise_error(ArgumentError, 'only DastSiteProfile or DastScannerProfile are allowed')
+    end
+  end
+
+  describe '#execute' do
+    subject { described_class.new(query_ctx, dast_profile) }
+
+    before do
+      subject.instance_variable_set(:@lazy_state, fake_state)
+    end
+
+    context 'if the record has already been loaded' do
+      let(:fake_state) do
+        { dast_pending_profiles: [], loaded_objects: { dast_profile => ['Dast Profile Name'] } }
+      end
+
+      it 'does not make the query again' do
+        expect(::Security::OrchestrationPolicyConfiguration).not_to receive(:for_project)
+
+        subject.execute
+      end
+    end
+
+    context 'if the record has not been loaded' do
+      let(:fake_state) do
+        { dast_pending_profiles: Set.new([dast_profile, other_dast_profile]), loaded_objects: {} }
+      end
+
+      let(:fake_policy_configuration) do
+        instance_double(::Security::OrchestrationPolicyConfiguration,
+          project_id: dast_profile.project_id,
+          active_policy_names_with_dast_site_profile: ['Dast Site Name'],
+          active_policy_names_with_dast_scanner_profile: ['Dast Scanner Name']
+        )
+      end
+
+      before do
+        allow(::Security::OrchestrationPolicyConfiguration).to receive(:for_project).and_return([fake_policy_configuration])
+      end
+
+      context 'when Dast Site profile is provided' do
+        it 'makes the query' do
+          expect(subject.execute).to eq(['Dast Site Name'])
+        end
+      end
+
+      context 'when Dast Scanner profile is provided' do
+        let(:dast_profile) { dast_scanner_profile }
+        let(:other_dast_profile) { other_dast_scanner_profile }
+
+        it 'makes the query' do
+          expect(subject.execute).to eq(['Dast Scanner Name'])
+        end
+      end
+
+      it 'clears the pending IDs' do
+        subject.execute
+
+        expect(subject.lazy_state[:dast_pending_profiles]).to be_empty
+      end
+    end
+  end
+end

ee/spec/models/security/orchestration_policy_configuration_spec.rb

@@ -4,7 +4,8 @@ require 'spec_helper'
 
 RSpec.describe Security::OrchestrationPolicyConfiguration do
   let_it_be(:security_policy_management_project) { create(:project, :repository) }
-  let_it_be(:security_orchestration_policy_configuration) do
+
+  let!(:security_orchestration_policy_configuration) do
     create(:security_orchestration_policy_configuration, security_policy_management_project: security_policy_management_project)
   end
 
@@ -26,6 +27,18 @@ RSpec.describe Security::OrchestrationPolicyConfiguration do
     it { is_expected.to validate_uniqueness_of(:project) }
   end
 
+  describe '.for_project' do
+    let!(:security_orchestration_policy_configuration_1) { create(:security_orchestration_policy_configuration) }
+    let!(:security_orchestration_policy_configuration_2) { create(:security_orchestration_policy_configuration) }
+    let!(:security_orchestration_policy_configuration_3) { create(:security_orchestration_policy_configuration) }
+
+    subject { described_class.for_project([security_orchestration_policy_configuration_2.project, security_orchestration_policy_configuration_3.project]) }
+
+    it 'returns configuration for given projects' do
+      is_expected.to contain_exactly(security_orchestration_policy_configuration_2, security_orchestration_policy_configuration_3)
+    end
+  end
+
   describe '#enabled?' do
     subject { security_orchestration_policy_configuration.enabled? }