gitlab-org--gitlab-foss/lib/gitlab/background_migration/populate_finding_uuid_for_v...

129 lines
3.8 KiB
Ruby

# frozen_string_literal: true
module Gitlab
module BackgroundMigration
# This class populates the `finding_uuid` attribute for
# the existing `vulnerability_feedback` records.
class PopulateFindingUuidForVulnerabilityFeedback
REPORT_TYPES = {
sast: 0,
dependency_scanning: 1,
container_scanning: 2,
dast: 3,
secret_detection: 4,
coverage_fuzzing: 5,
api_fuzzing: 6
}.freeze
class VulnerabilityFeedback < ActiveRecord::Base # rubocop:disable Style/Documentation
include EachBatch
self.table_name = 'vulnerability_feedback'
enum category: REPORT_TYPES
scope :in_range, -> (start, stop) { where(id: start..stop) }
scope :without_uuid, -> { where(finding_uuid: nil) }
def self.load_vulnerability_findings
all.to_a.tap { |collection| collection.each(&:vulnerability_finding) }
end
def set_finding_uuid
return unless vulnerability_finding.present? && vulnerability_finding.primary_identifier.present?
update_column(:finding_uuid, calculated_uuid)
rescue StandardError => error
Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error)
end
def vulnerability_finding
BatchLoader.for(finding_key).batch(replace_methods: false) do |finding_keys, loader|
project_ids = finding_keys.map { |key| key[:project_id] }
categories = finding_keys.map { |key| key[:category] }
fingerprints = finding_keys.map { |key| key[:project_fingerprint] }
findings = Finding.with_primary_identifier.where(
project_id: project_ids.uniq,
report_type: categories.uniq,
project_fingerprint: fingerprints.uniq
).to_a
finding_keys.each do |finding_key|
loader.call(
finding_key,
findings.find { |f| finding_key == f.finding_key }
)
end
end
end
private
def calculated_uuid
Gitlab::UUID.v5(uuid_components)
end
def uuid_components
[
category,
vulnerability_finding.primary_identifier.fingerprint,
vulnerability_finding.location_fingerprint,
project_id
].join('-')
end
def finding_key
{
project_id: project_id,
category: category,
project_fingerprint: project_fingerprint
}
end
end
class Finding < ActiveRecord::Base # rubocop:disable Style/Documentation
include ShaAttribute
self.table_name = 'vulnerability_occurrences'
sha_attribute :project_fingerprint
sha_attribute :location_fingerprint
belongs_to :primary_identifier, class_name: 'Gitlab::BackgroundMigration::PopulateFindingUuidForVulnerabilityFeedback::Identifier'
enum report_type: REPORT_TYPES
scope :with_primary_identifier, -> { includes(:primary_identifier) }
def finding_key
{
project_id: project_id,
category: report_type,
project_fingerprint: project_fingerprint
}
end
end
class Identifier < ActiveRecord::Base # rubocop:disable Style/Documentation
self.table_name = 'vulnerability_identifiers'
end
def perform(*range)
feedback = VulnerabilityFeedback.without_uuid.in_range(*range).load_vulnerability_findings
feedback.each(&:set_finding_uuid)
log_info(feedback.count)
end
def log_info(feedback_count)
::Gitlab::BackgroundMigration::Logger.info(
migrator: self.class.name,
message: '`finding_uuid` attributes has been set',
count: feedback_count
)
end
end
end
end