Fix N+1 issue when associating DAST profiles and CI Builds

This Merge Request fix the N+1 issue when are multiple builds with the dast_configuration keyword.

The N+1 issue was happening when associating profiles and finding profiles as well.

Changelog: fixed
MR: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/75448
EE: true

ee/app/services/app_sec/dast/profiles/create_associations_service.rb

+# frozen_string_literal: true
+
+module AppSec
+  module Dast
+    module Profiles
+      class CreateAssociationsService < BaseProjectService
+        include ::Gitlab::Ci::Pipeline::Chain::Helpers
+
+        def execute
+          return ServiceResponse.error(message: _('Insufficient permissions for dast_configuration keyword')) unless allowed?
+
+          dast_site_profiles = find_dast_site_profiles
+          dast_scanner_profiles = find_dast_scanner_profiles
+
+          dast_site_profiles_builds, dast_scanner_profiles_builds = prepare_batch_inserts(dast_scanner_profiles, dast_site_profiles)
+
+          return ServiceResponse.error(message: errors) unless errors.empty?
+
+          insert_builds(dast_site_profiles_builds, dast_scanner_profiles_builds)
+
+          ServiceResponse.success
+        end
+
+        private
+
+        def allowed?
+          can?(current_user, :create_on_demand_dast_scan, project)
+        end
+
+        def has_permission?(profile, name)
+          if can?(current_user, :read_on_demand_dast_scan, profile)
+            true
+          else
+            errors.push(_('DAST profile not found: %{name}') % { name: name })
+            false
+          end
+        end
+
+        def builds
+          @builds ||= params[:builds] || []
+        end
+
+        def errors
+          @errors ||= []
+        end
+
+        def prepare_batch_inserts(dast_scanner_profiles, dast_site_profiles)
+          dast_site_profiles_builds = []
+          dast_scanner_profiles_builds = []
+
+          builds.each do |build|
+            next unless build.is_a?(::Ci::Build)
+
+            if (site_profile_name = build.options.dig(:dast_configuration, :site_profile))
+              dast_site_profile = dast_site_profiles.find { |dsp| dsp.name == site_profile_name }
+
+              dast_site_profiles_builds.append({ ci_build_id: build.id, dast_site_profile_id: dast_site_profile.id }) if has_permission?(dast_site_profile, site_profile_name)
+            end
+
+            if (scanner_profile_name = build.options.dig(:dast_configuration, :scanner_profile))
+              dast_scanner_profile = dast_scanner_profiles.find { |dsp| dsp.name == scanner_profile_name }
+
+              dast_scanner_profiles_builds.append({ ci_build_id: build.id, dast_scanner_profile_id: dast_scanner_profile.id }) if has_permission?(dast_scanner_profile, scanner_profile_name)
+            end
+          end
+
+          [dast_site_profiles_builds, dast_scanner_profiles_builds]
+        end
+
+        def find(key, with:)
+          names = builds.map { |build| build.options.dig(:dast_configuration, key) }.compact
+
+          with.new(project_id: project.id, name: names).execute
+        end
+
+        def find_dast_site_profiles
+          find(:site_profile, with: DastSiteProfilesFinder)
+        end
+
+        def find_dast_scanner_profiles
+          find(:scanner_profile, with: DastScannerProfilesFinder)
+        end
+
+        def insert_builds(dast_site_profiles_builds, dast_scanner_profiles_builds)
+          ::Dast::SiteProfilesBuild.insert_all(dast_site_profiles_builds, unique_by: 'ci_build_id') unless dast_site_profiles_builds.empty?
+          ::Dast::ScannerProfilesBuild.insert_all(dast_scanner_profiles_builds, unique_by: 'ci_build_id') unless dast_scanner_profiles_builds.empty?
+        end
+      end
+    end
+  end
+end

ee/lib/ee/gitlab/ci/pipeline/chain/create_cross_database_associations.rb

@@ -26,44 +26,21 @@ module EE
               dast_stage = pipeline.stages.find { |stage| stage.name == ::AppSec::Dast::ScanConfigs::BuildService::STAGE_NAME }
               return unless dast_stage
 
-              # we use dast_stage.statuses to avoid extra sql queries
-              dast_stage.statuses.each do |status|
-                next unless status.is_a?(::Ci::Build)
-
-                associate_dast_profiles(dast_stage, status)
-              end
-            end
-
-            def associate_dast_profiles(stage, build)
-              response = find_dast_profiles(build)
+              response = AppSec::Dast::Profiles::CreateAssociationsService.new(
+                project: project,
+                current_user: current_user,
+                params: {
+                  builds: dast_stage.statuses, # we use dast_stage.statuses to avoid extra sql queries
+                  project: command.project
+                }
+              ).execute
 
               error(response.errors.join(', '), config_error: true) if response.error?
-              return if response.error? || response.payload.blank?
-
-              dast_site_profile = response.payload[:dast_site_profile]
-              Dast::SiteProfilesBuild.create!(ci_build: build, dast_site_profile: dast_site_profile) if dast_site_profile
-
-              dast_scanner_profile = response.payload[:dast_scanner_profile]
-              Dast::ScannerProfilesBuild.create!(ci_build: build, dast_scanner_profile: dast_scanner_profile) if dast_scanner_profile
             rescue ActiveRecord::ActiveRecordError => e
               ::Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e, extra: { pipeline_id: pipeline.id })
 
               error('Failed to associate DAST profiles')
             end
-
-            def find_dast_profiles(build)
-              dast_configuration = build.options[:dast_configuration]
-              return ServiceResponse.success unless dast_configuration
-
-              AppSec::Dast::Profiles::BuildConfigService.new(
-                project: build.project,
-                current_user: build.user,
-                params: {
-                  dast_site_profile: dast_configuration[:site_profile],
-                  dast_scanner_profile: dast_configuration[:scanner_profile]
-                }
-              ).execute
-            end
           end
         end
       end

ee/spec/finders/dast_scanner_profiles_finder_spec.rb

@@ -43,6 +43,14 @@ RSpec.describe DastScannerProfilesFinder do
       end
     end
 
+    context 'filter by names' do
+      let(:params) { { name: [dast_scanner_profile1.name, dast_scanner_profile2.name] } }
+
+      it 'returns the matching dast_scanner_profiles' do
+        expect(subject).to contain_exactly(dast_scanner_profile1, dast_scanner_profile2)
+      end
+    end
+
     context 'when DastScannerProfile id is for a different project' do
       let(:params) { { ids: [dast_scanner_profile1.id], project_ids: [dast_scanner_profile2.project.id] } }
 

ee/spec/finders/dast_site_profiles_finder_spec.rb

@@ -57,6 +57,14 @@ RSpec.describe DastSiteProfilesFinder do
       end
     end
 
+    context 'filtering by names' do
+      let(:params) { { name: [dast_site_profile1.name, dast_site_profile2.name] } }
+
+      it 'returns a single dast_site_profile' do
+        expect(subject).to contain_exactly(dast_site_profile1, dast_site_profile2)
+      end
+    end
+
     context 'when the dast_site_profile1 does not exist' do
       let(:params) { { id: 0 } }
 

ee/spec/services/app_sec/dast/profiles/create_associations_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe AppSec::Dast::Profiles::CreateAssociationsService do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:user) { create(:user, developer_projects: [project]) }
+  let_it_be(:outsider) { create(:user) }
+  let_it_be(:dast_site_profile) { create(:dast_site_profile, project: project) }
+  let_it_be(:dast_scanner_profile) { create(:dast_scanner_profile, project: project) }
+
+  let(:dast_site_profile_name) { dast_site_profile.name }
+  let(:dast_scanner_profile_name) { dast_scanner_profile.name }
+
+  let!(:pipeline) { create(:ci_pipeline, project: project, user: user) }
+  let!(:stage) { create(:ci_stage_entity, project: project, pipeline: pipeline, name: :dast) }
+
+  let!(:dast_build) do
+    create(:ci_build, project: project, user: user, pipeline: pipeline, stage_id: stage.id,
+                             options: { dast_configuration: { site_profile: dast_site_profile_name,
+                                                              scanner_profile: dast_scanner_profile_name } })
+  end
+
+  let(:params) { { builds: [dast_build] } }
+
+  subject { described_class.new(project: project, current_user: user, params: params).execute }
+
+  describe '#execute' do
+    context 'when the feature is licensed' do
+      before do
+        stub_licensed_features(security_on_demand_scans: true)
+        subject
+      end
+
+      context 'when the user cannot create dast scans' do
+        let_it_be(:user) { outsider }
+
+        it_behaves_like 'an error occurred during the dast profile association' do
+          let(:error_message) { 'Insufficient permissions for dast_configuration keyword' }
+        end
+      end
+
+      context 'dast_site_profile' do
+        let(:profile) { dast_site_profile }
+
+        it_behaves_like 'it attempts to associate the profile', :dast_site_profile_name
+      end
+
+      context 'dast_scanner_profile' do
+        let(:profile) { dast_scanner_profile }
+
+        it_behaves_like 'it attempts to associate the profile', :dast_scanner_profile_name
+      end
+
+      context 'when the user cannot create dast scans' do
+        let_it_be(:user) { outsider }
+
+        it_behaves_like 'an error occurred during the dast profile association' do
+          let(:error_message) { 'Insufficient permissions for dast_configuration keyword' }
+        end
+      end
+
+      context 'when the build has multiple dast_configurations' do
+        let_it_be(:dast_site_profile_2) { create(:dast_site_profile, project: project) }
+        let_it_be(:dast_scanner_profile_2) { create(:dast_scanner_profile, project: project) }
+
+        let(:dast_scanner_profile_2_name) { dast_scanner_profile_2.name }
+        let(:dast_site_profile_2_name) { dast_site_profile_2.name }
+
+        let!(:dast_build_2) do
+          create(:ci_build, project: project, user: user, pipeline: pipeline, stage_id: stage.id,
+                 options: { dast_configuration: { site_profile: dast_site_profile_2_name,
+                                                  scanner_profile: dast_scanner_profile_2_name } })
+        end
+
+        let(:builds) { [dast_build, dast_build_2] }
+        let(:params) { { builds: builds } }
+
+        let(:expected_associations) do
+          {
+            dast_build => {
+              dast_site_profile: dast_site_profile,
+              dast_scanner_profile: dast_scanner_profile
+            },
+            dast_build_2 => {
+              dast_site_profile: dast_site_profile_2,
+              dast_scanner_profile: dast_scanner_profile_2
+            }
+          }
+        end
+
+        it 'associations the associations correctly', :aggregate_failures do
+          expected_associations.each do |build, associations|
+            associations.each do |association_name, association|
+              expect(build.public_send(association_name)).to eq(association)
+            end
+          end
+        end
+      end
+    end
+
+    context 'when not licensed' do
+      before do
+        stub_licensed_features(security_on_demand_scans: false)
+      end
+
+      let(:error_message) { 'Insufficient permissions for dast_configuration keyword' }
+
+      it_behaves_like 'an error occurred during the dast profile association' do
+        let(:error_message) { 'Insufficient permissions for dast_configuration keyword' }
+      end
+    end
+  end
+end

ee/spec/services/ci/create_pipeline_service/dast_configuration_spec.rb

@@ -7,6 +7,8 @@ RSpec.describe Ci::CreatePipelineService do
   let_it_be(:user) { create(:user, developer_projects: [project]) }
   let_it_be(:dast_site_profile) { create(:dast_site_profile, project: project) }
   let_it_be(:dast_scanner_profile) { create(:dast_scanner_profile, project: project) }
+  let_it_be(:another_dast_site_profile) { create(:dast_site_profile, project: project) }
+  let_it_be(:another_dast_scanner_profile) { create(:dast_scanner_profile, project: project) }
 
   let(:dast_variables) do
     dast_site_profile.ci_variables
@@ -97,7 +99,7 @@ RSpec.describe Ci::CreatePipelineService do
 
       shared_examples 'a missing profile' do
         it 'communicates failure' do
-          expect(subject.yaml_errors).to eq("DAST profile not found: #{profile.name}")
+          expect(subject.yaml_errors).to include("DAST profile not found: #{profile.name}")
         end
       end
 
@@ -122,7 +124,7 @@ RSpec.describe Ci::CreatePipelineService do
         before do
           allow(error_tracking).to receive(:track_and_raise_for_dev_exception)
 
-          allow_next_instance_of(AppSec::Dast::Profiles::BuildConfigService) do |instance|
+          allow_next_instance_of(AppSec::Dast::Profiles::CreateAssociationsService) do |instance|
             allow(instance).to receive(:execute).and_raise(exception)
           end
         end
@@ -182,5 +184,55 @@ RSpec.describe Ci::CreatePipelineService do
         service.execute(:push)
       end
     end
+
+    it_behaves_like 'pipelines are created without N+1 SQL queries' do
+      let_it_be(:config1) do
+        <<~YAML
+        include:
+          - template: Security/DAST.gitlab-ci.yml
+        stages:
+          - dast
+        dast:
+          dast_configuration:
+            site_profile: #{dast_site_profile.name}
+            scanner_profile: #{dast_scanner_profile.name}
+        YAML
+      end
+
+      let_it_be(:config2) do
+        <<~YAML
+        stages:
+          - dast
+        dast:
+          stage: dast
+          dast_configuration:
+            site_profile: #{dast_site_profile.name}
+            scanner_profile: #{dast_scanner_profile.name}
+          script:
+            - exit 0
+        dast2:
+          stage: dast
+          dast_configuration:
+            site_profile: #{another_dast_site_profile.name}
+            scanner_profile: #{another_dast_scanner_profile.name}
+          script:
+            - exit 0
+        YAML
+      end
+
+      let(:accepted_n_plus_ones) do
+        1 + # SELECT "ci_instance_variables"
+          1 + # SELECT "ci_builds".* FROM "ci_builds"
+          1 + # INSERT INTO "ci_builds"
+          1 + # INSERT INTO "ci_builds_metadata"
+          1 + # SELECT "taggings".* FROM "taggings"
+          1 + # SELECT "ci_pipelines"."id" FROM
+          1 # SELECT "projects".id
+      end
+
+      def execute_service
+        service.execute(:push)
+      end
+    end
   end
 end

ee/spec/support/shared_examples/services/dast/dast_profile_associations_shared_examples.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.shared_examples 'it attempts to associate the profile' do |dast_profile_name_key|
+  let(:association) { dast_build.public_send(profile.class.underscore.to_sym) }
+  let(:profile_name) { public_send(dast_profile_name_key) }
+
+  context 'when the profile exists' do
+    it 'assigns the association' do
+      expect(association).to eq(profile)
+    end
+  end
+
+  shared_examples 'it has no effect' do
+    it 'does not assign the association' do
+      expect(association).to be_nil
+    end
+  end
+
+  context 'when the profile is not provided' do
+    let(dast_profile_name_key) { nil }
+
+    it_behaves_like 'it has no effect'
+  end
+
+  context 'when the profile does not exist' do
+    let(dast_profile_name_key) { SecureRandom.hex }
+
+    it_behaves_like 'an error occurred during the dast profile association' do
+      let(:error_message) { "DAST profile not found: #{profile_name}" }
+    end
+  end
+end
+
+RSpec.shared_examples 'an error occurred during the dast profile association' do
+  it 'communicates failure', :aggregate_failures do
+    expect(subject).to be_error
+    expect(subject.errors).to include(error_message)
+  end
+end