Resolve N+1 query issue for scanners in StoreSecurityReportsWorker [RUN ALL RSPEC] [RUN AS-IF-FOSS]

ee/app/services/security/store_report_service.rb

@@ -8,6 +8,8 @@ module Security
 
     attr_reader :pipeline, :report, :project
 
+    BATCH_SIZE = 1000
+
     def initialize(pipeline, report)
       @pipeline = pipeline
       @report = report
@@ -39,6 +41,8 @@ module Security
         .where(uuid: finding_uuids) # rubocop: disable CodeReuse/ActiveRecord
         .to_h { |vf| [vf.uuid, vf] }
 
+      update_vulnerability_scanners!(@report.findings) if Feature.enabled?(:optimize_sql_query_for_security_report, project)
+
       @report.findings.map do |finding|
         create_vulnerability_finding(vulnerability_findings_by_uuid, finding)&.id
       end.compact.uniq
@@ -63,7 +67,7 @@ module Security
       vulnerability_finding = vulnerability_findings_by_uuid[finding.uuid] ||
         create_new_vulnerability_finding(finding, vulnerability_params.merge(entity_params))
 
-      update_vulnerability_scanner(finding)
+      update_vulnerability_scanner(finding) unless Feature.enabled?(:optimize_sql_query_for_security_report, project)
 
       update_vulnerability_finding(vulnerability_finding, vulnerability_params)
       reset_remediations_for(vulnerability_finding, finding)
@@ -121,6 +125,50 @@ module Security
       scanner.update!(finding.scanner.to_hash)
     end
 
+    def vulnerability_scanner_attributes_keys
+      strong_memoize(:vulnerability_scanner_attributes_keys) do
+        Vulnerabilities::Scanner.new.attributes.keys
+      end
+    end
+
+    def valid_vulnerability_scanner_record?(record)
+      return false if (record.keys - vulnerability_scanner_attributes_keys).present?
+
+      record.values.all? {|value| value.present?}
+    end
+
+    def create_vulnerability_scanner_records(findings)
+      findings.map do |finding|
+        scanner = scanners_objects[finding.scanner.key]
+
+        next nil if scanner.nil?
+
+        scanner_attr = scanner.attributes.with_indifferent_access.except(:id)
+          .merge(finding.scanner.to_hash)
+
+        scanner_attr.compact!
+
+        scanner_attr
+      end
+    end
+
+    def update_vulnerability_scanners!(report_findings)
+      report_findings.in_groups_of(BATCH_SIZE, false) do |findings|
+        records = create_vulnerability_scanner_records(findings)
+        records.compact!
+        records.uniq!
+        records.each { |record| record.merge!({ created_at: Time.current, updated_at: Time.current }) }
+        records.filter! { |record| valid_vulnerability_scanner_record?(record) }
+
+        Vulnerabilities::Scanner.insert_all(records) if records.present?
+      end
+    rescue StandardError => e
+      Gitlab::ErrorTracking.track_exception(e)
+    ensure
+      clear_memoization(:scanners_objects)
+      clear_memoization(:existing_scanner_objects)
+    end
+
     def update_vulnerability_finding(vulnerability_finding, update_params)
       vulnerability_finding.update!(update_params)
     end

ee/config/feature_flags/development/optimize_sql_query_for_security_report.yml

+---
+name: optimize_sql_query_for_security_report
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/57426
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/323059
+milestone: '13.11'
+type: development
+group: group::static analysis
+default_enabled: false

ee/spec/services/security/store_report_service_spec.rb

@@ -28,46 +28,84 @@ RSpec.describe Security::StoreReportService, '#execute' do
       allow(pipeline).to receive(:user).and_return(user)
     end
 
-    using RSpec::Parameterized::TableSyntax
+    context 'for different security reports' do
+      using RSpec::Parameterized::TableSyntax
+
+      where(:case_name, :trait, :scanners, :identifiers, :findings, :finding_identifiers, :finding_pipelines, :remediations, :signatures, :optimize_sql_query_for_security_report_ff) do
+        'with SAST report'                | :sast                            | 1 | 6  | 5  | 7  | 5  | 0 | 2 | false
+        'with exceeding identifiers'      | :with_exceeding_identifiers      | 1 | 20 | 1  | 20 | 1  | 0 | 0 | false
+        'with Dependency Scanning report' | :dependency_scanning_remediation | 1 | 3  | 2  | 3  | 2  | 1 | 0 | false
+        'with Container Scanning report'  | :container_scanning              | 1 | 8  | 8  | 8  | 8  | 0 | 0 | false
+        'with SAST report'                | :sast                            | 1 | 6  | 5  | 7  | 5  | 0 | 2 | true
+        'with exceeding identifiers'      | :with_exceeding_identifiers      | 1 | 20 | 1  | 20 | 1  | 0 | 0 | true
+        'with Dependency Scanning report' | :dependency_scanning_remediation | 1 | 3  | 2  | 3  | 2  | 1 | 0 | true
+        'with Container Scanning report'  | :container_scanning              | 1 | 8  | 8  | 8  | 8  | 0 | 0 | true
+      end
 
-    where(:case_name, :trait, :scanners, :identifiers, :findings, :finding_identifiers, :finding_pipelines, :remediations, :signatures) do
-      'with SAST report'                | :sast                            | 1 | 6  | 5  | 7  | 5  | 0 | 2
-      'with exceeding identifiers'      | :with_exceeding_identifiers      | 1 | 20 | 1  | 20 | 1  | 0 | 0
-      'with Dependency Scanning report' | :dependency_scanning_remediation | 1 | 3  | 2  | 3  | 2  | 1 | 0
-      'with Container Scanning report'  | :container_scanning              | 1 | 8  | 8  | 8  | 8  | 0 | 0
-    end
+      with_them do
+        before do
+          stub_feature_flags(optimize_sql_query_for_security_report: optimize_sql_query_for_security_report_ff)
+        end
 
-    with_them do
-      it 'inserts all scanners' do
-        expect { subject }.to change { Vulnerabilities::Scanner.count }.by(scanners)
-      end
+        it 'inserts all scanners' do
+          expect { subject }.to change { Vulnerabilities::Scanner.count }.by(scanners)
+        end
 
-      it 'inserts all identifiers' do
-        expect { subject }.to change { Vulnerabilities::Identifier.count }.by(identifiers)
-      end
+        it 'inserts all identifiers' do
+          expect { subject }.to change { Vulnerabilities::Identifier.count }.by(identifiers)
+        end
 
-      it 'inserts all findings' do
-        expect { subject }.to change { Vulnerabilities::Finding.count }.by(findings)
-      end
+        it 'inserts all findings' do
+          expect { subject }.to change { Vulnerabilities::Finding.count }.by(findings)
+        end
 
-      it 'inserts all finding identifiers (join model)' do
-        expect { subject }.to change { Vulnerabilities::FindingIdentifier.count }.by(finding_identifiers)
-      end
+        it 'inserts all finding identifiers (join model)' do
+          expect { subject }.to change { Vulnerabilities::FindingIdentifier.count }.by(finding_identifiers)
+        end
 
-      it 'inserts all finding pipelines (join model)' do
-        expect { subject }.to change { Vulnerabilities::FindingPipeline.count }.by(finding_pipelines)
-      end
+        it 'inserts all finding pipelines (join model)' do
+          expect { subject }.to change { Vulnerabilities::FindingPipeline.count }.by(finding_pipelines)
+        end
+
+        it 'inserts all remediations' do
+          expect { subject }.to change { project.vulnerability_remediations.count }.by(remediations)
+        end
 
-      it 'inserts all remediations' do
-        expect { subject }.to change { project.vulnerability_remediations.count }.by(remediations)
+        it 'inserts all vulnerabilities' do
+          expect { subject }.to change { Vulnerability.count }.by(findings)
+        end
+
+        it 'inserts all signatures' do
+          expect { subject }.to change { Vulnerabilities::FindingSignature.count }.by(signatures)
+        end
       end
+    end
+
+    context 'when there is an exception' do
+      let(:trait) { :sast }
+
+      subject { described_class.new(pipeline, report) }
 
-      it 'inserts all vulnerabilities' do
-        expect { subject }.to change { Vulnerability.count }.by(findings)
+      it 'does not insert any scanner' do
+        allow(Vulnerabilities::Scanner).to receive(:insert_all).with(anything).and_raise(StandardError)
+        expect { subject.send(:update_vulnerability_scanners!, report.findings) }.to change { Vulnerabilities::Scanner.count }.by(0)
       end
+    end
+
+    context 'when N+1 database queries have been removed' do
+      let(:trait) { :sast }
+      let(:bandit_scanner) { build(:ci_reports_security_scanner, external_id: 'bandit', name: 'Bandit') }
+
+      subject { described_class.new(pipeline, report) }
+
+      it "avoids N+1 database queries for updating vulnerability scanners", :use_sql_query_cache do
+        report.add_scanner(bandit_scanner)
+
+        control_count = ActiveRecord::QueryRecorder.new(skip_cached: false) { subject.send(:update_vulnerability_scanners!, report.findings) }.count
+
+        5.times { report.add_finding(build(:ci_reports_security_finding, scanner: bandit_scanner)) }
 
-      it 'inserts all signatures' do
-        expect { subject }.to change { Vulnerabilities::FindingSignature.count }.by(signatures)
+        expect {  described_class.new(pipeline, report).send(:update_vulnerability_scanners!, report.findings) }.not_to exceed_query_limit(control_count)
       end
     end