Merge branch...

Merge branch 'schmil-Optimise-SAST-Vulnerability-Auto-Resolving-with-In-Operator-Optimisation' into 'master' 

Optimise Vulnerability Auto-Resolving with In-Operator Optimisation

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/154553



Merged-by: Alex Pooley <apooley@gitlab.com>
Approved-by: Alex Pooley <apooley@gitlab.com>
Reviewed-by: Alex Pooley <apooley@gitlab.com>
Co-authored-by: Schmil Monderer <smonderer@gitlab.com>

ee/app/models/vulnerabilities/identifier.rb

@@ -29,9 +29,6 @@ class Identifier < ApplicationRecord
     scope :by_projects, ->(values) { where(project_id: values) }
     scope :with_fingerprint, ->(fingerprints) { where(fingerprint: fingerprints) }
     scope :with_external_type, ->(external_type) { where('LOWER(external_type) = LOWER(?)', external_type) }
-    scope :select_primary_finding_vulnerability_ids, -> {
-      joins(:primary_findings).select('vulnerability_occurrences.vulnerability_id AS vulnerability_id')
-    }
 
     def cve?
       external_type.casecmp?('cve')

ee/app/workers/vulnerabilities/mark_dropped_as_resolved_worker.rb

@@ -5,6 +5,8 @@ module Vulnerabilities
   class MarkDroppedAsResolvedWorker
     include ApplicationWorker
 
+    VULNERABILITY_IDENTIFIERS_BATCH_SIZE = 250
+
     data_consistency :delayed
     idempotent!
     deduplicate :until_executing, including_scheduled: true
@@ -14,50 +16,54 @@ class MarkDroppedAsResolvedWorker
     loggable_arguments 1
 
     def perform(_, dropped_identifier_ids)
-      vulnerability_ids = vulnerability_ids_for(dropped_identifier_ids)
-      vulnerabilities = resolvable_vulnerabilities(vulnerability_ids)
+      resolvable_vulnerabilities(dropped_identifier_ids) do |vulnerabilities|
+        next unless vulnerabilities.present?
 
-      resolve_vulnerabilities(vulnerabilities)
-    end
+        # rubocop:disable CodeReuse/ActiveRecord -- `update_all` changes the result of the query, so grab the ids first to do the system notes
+        vulnerability_ids = vulnerabilities.pluck(:id)
+        # rubocop:enable CodeReuse/ActiveRecord
 
-    private
+        current_time = Time.zone.now
 
-    def vulnerability_ids_for(identifier_ids)
-      ::Vulnerabilities::Identifier.id_in(identifier_ids)
-      .order(:id) # rubocop:disable CodeReuse/ActiveRecord -- unusual order call is very specific to this query
-      .select_primary_finding_vulnerability_ids
-      .map(&:vulnerability_id)
-    end
-
-    def resolvable_vulnerabilities(vulnerability_ids)
-      return [] unless vulnerability_ids.present?
-
-      ::Vulnerability.with_states(:detected)
-      .with_resolution(true)
-      .by_ids(vulnerability_ids)
-    end
+        state_transitions = build_state_transitions(vulnerabilities, current_time)
 
-    def resolve_vulnerabilities(vulnerabilities)
-      return unless vulnerabilities.present?
+        ::Vulnerability.transaction do
+          vulnerabilities.update_all(
+            resolved_by_id: Users::Internal.security_bot.id,
+            resolved_at: current_time,
+            updated_at: current_time,
+            state: :resolved)
 
-      current_time = Time.zone.now
+          Vulnerabilities::StateTransition.bulk_insert!(state_transitions)
+        end
 
-      state_transitions = build_state_transitions(vulnerabilities, current_time)
-      # rubocop:disable CodeReuse/ActiveRecord -- `update_all` changes the result of the query, preventing the system note update
-      vuln_ids_to_be_updated = vulnerabilities.pluck(:id)
-      # rubocop:enable CodeReuse/ActiveRecord
+        create_system_notes(Vulnerability.by_ids(vulnerability_ids))
+      end
+    end
 
-      ::Vulnerability.transaction do
-        vulnerabilities.update_all(
-          resolved_by_id: Users::Internal.security_bot.id,
-          resolved_at: current_time,
-          updated_at: current_time,
-          state: :resolved)
+    private
 
-        Vulnerabilities::StateTransition.bulk_insert!(state_transitions)
+    def resolvable_vulnerabilities(identifier_ids)
+      # rubocop:disable CodeReuse/ActiveRecord -- unusual order, pluck and where calls is very specific to this query
+      order = Gitlab::Pagination::Keyset::Order.build([
+        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+          attribute_name: 'vulnerability_id',
+          order_expression: Vulnerabilities::Finding.arel_table[:vulnerability_id].asc,
+          nullable: :not_nullable
+        )
+      ])
+
+      identifier_ids.each do |identifier_id|
+        scope = Vulnerabilities::Finding.where(primary_identifier_id: identifier_id).order(order)
+
+        Gitlab::Pagination::Keyset::Iterator.new(scope: scope)
+                                            .each_batch(of: VULNERABILITY_IDENTIFIERS_BATCH_SIZE) do |records|
+          vulnerability_ids = records.map(&:vulnerability_id)
+
+          yield ::Vulnerability.id_in(vulnerability_ids).with_states(:detected).with_resolution(true)
+        end
       end
-
-      create_system_notes(Vulnerability.by_ids(vuln_ids_to_be_updated))
+      # rubocop:enable CodeReuse/ActiveRecord
     end
 
     def build_state_transitions(vulnerabilities, current_time)

ee/spec/models/vulnerabilities/identifier_spec.rb

@@ -62,18 +62,6 @@
     end
   end
 
-  describe '.select_primary_finding_vulnerability_ids' do
-    let!(:identifier) { create(:vulnerabilities_identifier) }
-    let!(:primary_finding) { create(:vulnerabilities_finding, primary_identifier_id: identifier.id) }
-
-    subject { described_class.select_primary_finding_vulnerability_ids.map(&:vulnerability_id) }
-
-    it 'selects the primary finding vulnerability ids' do
-      is_expected.not_to be_empty
-      is_expected.to match_array(identifier.primary_findings.map(&:vulnerability_id))
-    end
-  end
-
   describe '.with_fingerprint' do
     let(:fingerprint) { 'f5724386167705667ae25a1390c0a516020690ba' }