Allow DAST profile build associations to be cloned

- Moves association persistence to after commit hook
- Removes allow wrapper around test

app/services/ci/retry_build_service.rb

@@ -70,7 +70,9 @@ module Ci
     def check_assignable_runners!(build); end
 
     def clone_build(build)
-      project.builds.new(build_attributes(build))
+      project.builds.new(build_attributes(build)).tap do |new_build|
+        yield(new_build) if block_given?
+      end
     end
 
     def build_attributes(build)

config/feature_flags/development/dast_sharded_cloned_ci_builds.yml

+---
+name: dast_sharded_cloned_ci_builds
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/78164
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/351980
+milestone: '14.8'
+type: development
+group: group::dynamic analysis
+default_enabled: true

config/sidekiq_queues.yml

@@ -35,8 +35,12 @@
   - 1
 - - analytics_usage_trends_counter_job
   - 1
+- - app_sec_dast_scanner_profiles_builds_consistency
+  - 1
 - - app_sec_dast_scans_consistency
   - 1
+- - app_sec_dast_site_profiles_builds_consistency
+  - 1
 - - approval_rules_external_approval_rule_payload
   - 1
 - - approve_blocked_pending_approval_users

ee/app/services/app_sec/dast/builds/associate_service.rb

+# frozen_string_literal: true
+
+module AppSec
+  module Dast
+    module Builds
+      class AssociateService
+        def initialize(params)
+          @params = params
+        end
+
+        def execute
+          responses = [associate_site_profile, associate_scanner_profile]
+
+          responses.each { |response| return response if response.error? }
+
+          ServiceResponse.success
+        end
+
+        private
+
+        attr_reader :params
+
+        def associate_site_profile
+          return ServiceResponse.success unless params[:dast_site_profile_id]
+
+          association = ::Dast::SiteProfilesBuild.new(
+            ci_build_id: params[:ci_build_id],
+            dast_site_profile_id: params[:dast_site_profile_id]
+          )
+
+          save(association).tap do |response|
+            if response.error?
+              SiteProfilesBuilds::ConsistencyWorker.perform_async(params[:ci_build_id], params[:dast_site_profile_id])
+            end
+          end
+        end
+
+        def associate_scanner_profile
+          return ServiceResponse.success unless params[:dast_scanner_profile_id]
+
+          association = ::Dast::ScannerProfilesBuild.new(
+            ci_build_id: params[:ci_build_id],
+            dast_scanner_profile_id: params[:dast_scanner_profile_id]
+          )
+
+          save(association).tap do |response|
+            if response.error?
+              ScannerProfilesBuilds::ConsistencyWorker.perform_async(params[:ci_build_id], params[:dast_scanner_profile_id])
+            end
+          end
+        end
+
+        def save(association)
+          return ServiceResponse.success if association.save
+
+          ServiceResponse.error(message: association.errors.full_messages)
+        end
+      end
+    end
+  end
+end

ee/app/services/ee/ci/retry_build_service.rb

@@ -16,12 +16,31 @@ module EE
 
         override :extra_accessors
         def extra_accessors
+          return %i[secrets].freeze if ::Feature.enabled?(:dast_sharded_cloned_ci_builds, default_enabled: :yaml)
+
           %i[dast_site_profile dast_scanner_profile secrets].freeze
         end
       end
 
       private
 
+      override :clone_build
+      def clone_build(build)
+        super do |new_build|
+          if ::Feature.enabled?(:dast_sharded_cloned_ci_builds, default_enabled: :yaml)
+            new_build.run_after_commit do
+              response = AppSec::Dast::Builds::AssociateService.new(
+                ci_build_id: new_build.id,
+                dast_site_profile_id: build.dast_site_profile&.id,
+                dast_scanner_profile_id: build.dast_scanner_profile&.id
+              ).execute
+
+              new_build.reset.drop! if response.error?
+            end
+          end
+        end
+      end
+
       override :check_access!
       def check_access!(build)
         super

ee/app/workers/all_queues.yml

@@ -876,6 +876,15 @@
   :weight: 1
   :idempotent: true
   :tags: []
+- :name: app_sec_dast_scanner_profiles_builds_consistency
+  :worker_name: AppSec::Dast::ScannerProfilesBuilds::ConsistencyWorker
+  :feature_category: :dynamic_application_security_testing
+  :has_external_dependencies:
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: app_sec_dast_scans_consistency
   :worker_name: AppSec::Dast::Scans::ConsistencyWorker
   :feature_category: :dynamic_application_security_testing
@@ -885,6 +894,15 @@
   :weight: 1
   :idempotent: true
   :tags: []
+- :name: app_sec_dast_site_profiles_builds_consistency
+  :worker_name: AppSec::Dast::SiteProfilesBuilds::ConsistencyWorker
+  :feature_category: :dynamic_application_security_testing
+  :has_external_dependencies:
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: approval_rules_external_approval_rule_payload
   :worker_name: ApprovalRules::ExternalApprovalRulePayloadWorker
   :feature_category: :source_code_management

ee/app/workers/app_sec/dast/scanner_profiles_builds/consistency_worker.rb

+# frozen_string_literal: true
+
+module AppSec
+  module Dast
+    module ScannerProfilesBuilds
+      class ConsistencyWorker
+        include ApplicationWorker
+
+        data_consistency :always
+
+        deduplicate :until_executed
+        idempotent!
+
+        feature_category :dynamic_application_security_testing
+
+        def perform(ci_pipeline_id, dast_scanner_profile_id)
+          ::Dast::ScannerProfilesBuild.create!(ci_build_id: ci_pipeline_id, dast_scanner_profile_id: dast_scanner_profile_id)
+        rescue ActiveRecord::RecordNotUnique
+          # assume record is already associated
+        end
+      end
+    end
+  end
+end

ee/app/workers/app_sec/dast/site_profiles_builds/consistency_worker.rb

+# frozen_string_literal: true
+
+module AppSec
+  module Dast
+    module SiteProfilesBuilds
+      class ConsistencyWorker
+        include ApplicationWorker
+
+        data_consistency :always
+
+        deduplicate :until_executed
+        idempotent!
+
+        feature_category :dynamic_application_security_testing
+
+        def perform(ci_pipeline_id, dast_site_profile_id)
+          ::Dast::SiteProfilesBuild.create!(ci_build_id: ci_pipeline_id, dast_site_profile_id: dast_site_profile_id)
+        rescue ActiveRecord::RecordNotUnique
+          # assume record is already associated
+        end
+      end
+    end
+  end
+end

ee/spec/services/app_sec/dast/builds/associate_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe AppSec::Dast::Builds::AssociateService do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:pipeline) { create(:ci_pipeline, project: project) }
+  let_it_be(:build) { create(:ci_build, pipeline: pipeline) }
+  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(:params) do
+    { ci_build_id: build.id, dast_site_profile_id: dast_site_profile.id, dast_scanner_profile_id: dast_scanner_profile.id }
+  end
+
+  describe '#execute' do
+    subject(:execute) do
+      described_class.new(params).execute
+    end
+
+    context 'params' do
+      context 'when no keys are supplied' do
+        let(:params) { {} }
+
+        it 'returns a success response' do
+          expect(execute).to be_success
+        end
+      end
+
+      context 'when the ci_build_id key missing' do
+        let(:params) do
+          { dast_site_profile_id: dast_site_profile.id, dast_scanner_profile_id: dast_scanner_profile.id }
+        end
+
+        it 'returns an error response' do
+          expect(execute).to have_attributes(status: :error, message: ['Ci build must exist', 'Ci build can\'t be blank'])
+        end
+      end
+    end
+
+    context 'success' do
+      it 'returns a success response' do
+        expect(execute).to be_success
+      end
+
+      it 'associates the site profile' do
+        execute
+
+        expect(build.reload.dast_site_profile).to eq(dast_site_profile)
+      end
+
+      it 'associates the scanner profile' do
+        execute
+
+        expect(build.reload.dast_scanner_profile).to eq(dast_scanner_profile)
+      end
+
+      it 'does not call any consistency workers' do
+        expect(AppSec::Dast::SiteProfilesBuilds::ConsistencyWorker).not_to receive(:perform_async)
+        expect(AppSec::Dast::ScannerProfilesBuilds::ConsistencyWorker).not_to receive(:perform_async)
+
+        execute
+      end
+    end
+
+    context 'error' do
+      shared_examples 'an error' do
+        it 'returns an error response' do
+          expect(execute).to be_error
+        end
+      end
+
+      shared_examples 'it attempts to maintain site profile association consistency' do
+        it 'calls the site profile consistency worker' do
+          expect(AppSec::Dast::SiteProfilesBuilds::ConsistencyWorker).to receive(:perform_async).with(build.id, dast_site_profile.id).and_call_original
+
+          execute
+        end
+      end
+
+      shared_examples 'it attempts to maintain scanner profile association consistency' do
+        it 'calls the scanner profile consistency worker' do
+          expect(AppSec::Dast::ScannerProfilesBuilds::ConsistencyWorker).to receive(:perform_async).with(build.id, dast_scanner_profile.id).and_call_original
+
+          execute
+        end
+      end
+
+      context 'when saving a SiteProfilesBuild fails' do
+        before do
+          stub_save_failure(::Dast::SiteProfilesBuild)
+        end
+
+        it_behaves_like 'an error'
+        it_behaves_like 'it attempts to maintain site profile association consistency'
+      end
+
+      context 'when saving a ScannerProfilesBuild fails' do
+        before do
+          stub_save_failure(::Dast::ScannerProfilesBuild)
+        end
+
+        it_behaves_like 'an error'
+        it_behaves_like 'it attempts to maintain scanner profile association consistency'
+      end
+
+      context 'when saving both associations fails' do
+        before do
+          stub_save_failure(::Dast::SiteProfilesBuild)
+          stub_save_failure(::Dast::ScannerProfilesBuild)
+        end
+
+        it_behaves_like 'an error'
+        it_behaves_like 'it attempts to maintain site profile association consistency'
+        it_behaves_like 'it attempts to maintain scanner profile association consistency'
+      end
+
+      def stub_save_failure(klass)
+        allow_next_instance_of(klass) do |instance|
+          allow(instance).to receive(:save).and_return(false)
+        end
+      end
+    end
+  end
+end

ee/spec/services/ci/retry_build_service_spec.rb

@@ -20,27 +20,48 @@ RSpec.describe Ci::RetryBuildService do
   describe '#clone!' do
     context 'when user has ability to execute build' do
       let_it_be(:namespace) { create(:namespace) }
+      let_it_be(:project) { create(:project, namespace: namespace, creator: user) }
 
-      let(:project) { create(:project, namespace: namespace, creator: user) }
-
-      let(:new_build) do
-        travel_to(1.second.from_now) do
-          service.clone!(build)
-        end
-      end
+      let(:new_build) { service.clone!(build)}
 
       context 'dast' do
-        let(:dast_site_profile) { create(:dast_site_profile, project: project) }
-        let(:dast_scanner_profile) { create(:dast_scanner_profile, project: 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) }
 
         before do
           build.update!(dast_site_profile: dast_site_profile, dast_scanner_profile: dast_scanner_profile)
         end
 
-        it 'clones the profile associations', :aggregate_failures do
-          ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification_within_transaction(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/350051') do
+        context 'failure' do
+          it 'drops the build' do
+            allow_next_instance_of(AppSec::Dast::Builds::AssociateService) do |instance|
+              allow(instance).to receive(:execute).and_return(ServiceResponse.error(message: 'oops'))
+            end
+
+            expect(new_build.reload).to be_failed
+          end
+        end
+
+        context 'success' do
+          it 'clones the profile associations', :aggregate_failures do
+            new_build.reload
+
             expect(new_build.dast_site_profile).to eq(dast_site_profile)
             expect(new_build.dast_scanner_profile).to eq(dast_scanner_profile)
+            expect(new_build).not_to be_failed
+          end
+        end
+
+        context 'when dast_sharded_cloned_ci_builds is disabled' do
+          before do
+            stub_feature_flags(dast_sharded_cloned_ci_builds: false)
+          end
+
+          it 'clones the profile associations', :aggregate_failures do
+            ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification_within_transaction(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/350051') do
+              expect(new_build.dast_site_profile).to eq(dast_site_profile)
+              expect(new_build.dast_scanner_profile).to eq(dast_scanner_profile)
+            end
           end
         end
       end
@@ -127,11 +148,7 @@ RSpec.describe Ci::RetryBuildService do
   end
 
   describe '#execute' do
-    let(:new_build) do
-      travel_to(1.second.from_now) do
-        service.execute(build)
-      end
-    end
+    let(:new_build) { service.execute(build) }
 
     context 'when the CI quota is exceeded' do
       let_it_be(:namespace) { create(:namespace, :with_used_build_minutes_limit) }

ee/spec/workers/app_sec/dast/scanner_profiles_builds/consistency_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe AppSec::Dast::ScannerProfilesBuilds::ConsistencyWorker do
+  let(:worker) { described_class.new }
+
+  describe '#perform' do
+    let_it_be(:project) { create(:project) }
+    let_it_be(:build) { create(:ci_build, project: project) }
+    let_it_be(:profile) { create(:dast_scanner_profile, project: project) }
+
+    let(:job_args) { [build.id, profile.id] }
+
+    it 'ensures cross database association is created', :aggregate_failures do
+      expect { worker.perform(*job_args) }.to change { Dast::ScannerProfilesBuild.count }.by(1)
+
+      expect(Dast::ScannerProfilesBuild.where(ci_build_id: build.id, dast_scanner_profile_id: profile.id)).to exist
+    end
+
+    it_behaves_like 'an idempotent worker'
+  end
+end

ee/spec/workers/app_sec/dast/site_profiles_builds/consistency_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe AppSec::Dast::SiteProfilesBuilds::ConsistencyWorker do
+  let(:worker) { described_class.new }
+
+  describe '#perform' do
+    let_it_be(:project) { create(:project) }
+    let_it_be(:build) { create(:ci_build, project: project) }
+    let_it_be(:profile) { create(:dast_site_profile, project: project) }
+
+    let(:job_args) { [build.id, profile.id] }
+
+    it 'ensures cross database association is created', :aggregate_failures do
+      expect { worker.perform(*job_args) }.to change { Dast::SiteProfilesBuild.count }.by(1)
+
+      expect(Dast::SiteProfilesBuild.where(ci_build_id: build.id, dast_site_profile_id: profile.id)).to exist
+    end
+
+    it_behaves_like 'an idempotent worker'
+  end
+end

spec/services/ci/retry_build_service_spec.rb

@@ -60,7 +60,8 @@ RSpec.describe Ci::RetryBuildService do
        artifacts_file artifacts_metadata artifacts_size commands
        resource resource_group_id processed security_scans author
        pipeline_id report_results pending_state pages_deployments
-       queuing_entry runtime_metadata trace_metadata].freeze
+       queuing_entry runtime_metadata trace_metadata
+       dast_site_profile dast_scanner_profile].freeze
 
   shared_examples 'build duplication' do
     let_it_be(:another_pipeline) { create(:ci_empty_pipeline, project: project) }