Merge branch...

Merge branch 'minac_474280_set_total_number_of_vulnerabilities_for_existing_projects' into 'master' 

Set `vulnerability_count` of project_security_statistics records

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



Merged-by: Gregory Havenga <11164960-ghavenga@users.noreply.gitlab.com>
Approved-by: Buck O'Leary <bucoleary@gitlab.com>
Approved-by: Alex Pooley <apooley@gitlab.com>
Approved-by: Nicolas Dular <ndular@gitlab.com>
Approved-by: Gregory Havenga <11164960-ghavenga@users.noreply.gitlab.com>
Reviewed-by: Nicolas Dular <ndular@gitlab.com>
Co-authored-by: Mehmet Emin INAC <minac@gitlab.com>

db/docs/batched_background_migrations/set_project_vulnerability_count.yml → db/docs/batched_background_migrations/set_total_number_of_vulnerabilities_for_existing_projects.yml

 ---
-migration_job_name: SetProjectVulnerabilityCount
-description: Sets the number of vulnerabilities a project has in the vulnerbaility_count column in the project_statistics table
+migration_job_name: SetTotalNumberOfVulnerabilitiesForExistingProjects
+description: Sets the `vulnerability_count` column of `project_security_statistics` table.
 feature_category: vulnerability_management
-introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/166471
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/167812
 milestone: '17.5'
-queued_migration_version: 20240918122045
-# Replace with the approximate date you think it's best to ensure the completion of this BBM.
+queued_migration_version: 20241001115912
 finalized_by: # version of the migration that finalized this BBM

db/post_migrate/20240918122045_reschedule_setting_project_vulnerability_count.rb

@@ -3,27 +3,7 @@
 class RescheduleSettingProjectVulnerabilityCount < Gitlab::Database::Migration[2.2]
   milestone '17.5'
 
-  restrict_gitlab_migration gitlab_schema: :gitlab_main
-
-  MIGRATION = "SetProjectVulnerabilityCount"
-  DELAY_INTERVAL = 2.minutes
-  BATCH_SIZE = 1000
-  SUB_BATCH_SIZE = 100
-
-  def up
-    delete_batched_background_migration(MIGRATION, :project_settings, :project_id, [])
-
-    queue_batched_background_migration(
-      MIGRATION,
-      :project_settings,
-      :project_id,
-      job_interval: DELAY_INTERVAL,
-      batch_size: BATCH_SIZE,
-      sub_batch_size: SUB_BATCH_SIZE
-    )
-  end
-
-  def down
-    delete_batched_background_migration(MIGRATION, :project_settings, :project_id, [])
+  def change
+    # no-op
   end
 end

db/post_migrate/20241001115912_queue_set_total_number_of_vulnerabilities_for_existing_projects.rb

+# frozen_string_literal: true
+
+class QueueSetTotalNumberOfVulnerabilitiesForExistingProjects < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+
+  restrict_gitlab_migration gitlab_schema: :gitlab_sec
+
+  MIGRATION = "SetTotalNumberOfVulnerabilitiesForExistingProjects"
+  DELAY_INTERVAL = 2.minutes
+  BATCH_SIZE = 1000
+  SUB_BATCH_SIZE = 100
+
+  def up
+    delete_batched_background_migration('SetProjectVulnerabilityCount', :project_settings, :project_id, [])
+
+    queue_batched_background_migration(
+      MIGRATION,
+      :vulnerability_reads,
+      :project_id,
+      batch_class_name: 'LooseIndexScanBatchingStrategy',
+      job_interval: DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      sub_batch_size: SUB_BATCH_SIZE
+    )
+  end
+
+  def down
+    delete_batched_background_migration(MIGRATION, :vulnerability_reads, :project_id, [])
+  end
+end

db/schema_migrations/20241001115912

+7f37b4779c85f0290bc5c27104934e74cf9afe968906cdebfd02f5d915336c67
\ No newline at end of file

ee/lib/ee/gitlab/background_migration/set_project_vulnerability_count.rb

-# frozen_string_literal: true
-
-module EE
-  module Gitlab
-    module BackgroundMigration
-      module SetProjectVulnerabilityCount
-        extend ActiveSupport::Concern
-        extend ::Gitlab::Utils::Override
-
-        prepended do
-          scope_to ->(relation) { relation.where("has_vulnerabilities IS TRUE") }
-          operation_name :set_project_vulnerability_count
-          feature_category :vulnerability_management
-        end
-
-        class ProjectStatistics < ApplicationRecord
-          self.table_name = 'project_statistics'
-        end
-
-        class VulnerabilityRead < ApplicationRecord
-          include EachBatch
-
-          self.table_name = 'vulnerability_reads'
-          self.primary_key = :vulnerability_id
-        end
-
-        override :perform
-
-        def perform
-          each_sub_batch do |sub_batch|
-            sub_batch.each do |project_settings|
-              project_statistics = ProjectStatistics.find_by(project_id: project_settings.project_id)
-              latest_vulnerability_id = nil
-
-              ProjectStatistics.transaction do
-                project_statistics.lock!
-
-                latest_vulnerability_id = VulnerabilityRead.where(project_id: project_settings.project_id).last&.id
-
-                project_statistics.vulnerability_count = 0
-                project_statistics.save!
-              end
-
-              next unless latest_vulnerability_id
-
-              total_count = 0
-
-              VulnerabilityRead.where('project_id = ? AND vulnerability_id <= ?',
-                project_settings.project_id, latest_vulnerability_id).each_batch do |batch|
-                total_count += batch.count
-              end
-
-              ProjectStatistics.connection.exec_query(update_query, 'UPDATE', [total_count, project_statistics.id])
-            end
-          end
-        end
-
-        def update_query
-          <<~SQL
-            UPDATE
-              project_statistics
-            SET
-              vulnerability_count = vulnerability_count + $1
-            WHERE
-              id = $2
-          SQL
-        end
-      end
-    end
-  end
-end

ee/lib/ee/gitlab/background_migration/set_total_number_of_vulnerabilities_for_existing_projects.rb

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module BackgroundMigration
+      module SetTotalNumberOfVulnerabilitiesForExistingProjects
+        extend ActiveSupport::Concern
+        extend ::Gitlab::Utils::Override
+
+        prepended do
+          operation_name :set_vulnerability_count
+        end
+
+        PROJECT_STATISTICS_UPDATE_QUERY = <<~SQL
+          UPDATE
+            project_security_statistics
+          SET
+            vulnerability_count = vulnerability_count + $1
+          WHERE
+            project_id = $2
+        SQL
+
+        class ProjectSecurityStatistics < ::Gitlab::Database::SecApplicationRecord; end
+
+        class VulnerabilityRead < ::Gitlab::Database::SecApplicationRecord
+          include EachBatch
+
+          self.primary_key = :vulnerability_id
+        end
+
+        override :perform
+        def perform
+          distinct_each_batch do |sub_batch|
+            project_ids = sub_batch.pluck(:project_id)
+
+            ensure_statistics(project_ids)
+
+            project_ids.each { |project_id| update_project_statistics(project_id) }
+          end
+        end
+
+        private
+
+        def ensure_statistics(project_ids)
+          project_ids.map { |project_id| { project_id: project_id } }
+                     .then { |attributes| ProjectSecurityStatistics.upsert_all(attributes) }
+        end
+
+        def update_project_statistics(project_id)
+          latest_vulnerability_id = reset_statistics_and_get_latest_vulnerability_id(project_id)
+          total_count = calculate_number_of_vulnerabilities(project_id, latest_vulnerability_id)
+
+          persist_statistics(project_id, total_count)
+        end
+
+        def reset_statistics_and_get_latest_vulnerability_id(project_id)
+          latest_vulnerability_id = nil
+
+          ProjectSecurityStatistics.transaction do
+            statistics = ProjectSecurityStatistics.lock.find(project_id)
+
+            latest_vulnerability_id = VulnerabilityRead.where(project_id: project_id).last.vulnerability_id
+
+            statistics.vulnerability_count = 0
+            statistics.save!
+          end
+
+          latest_vulnerability_id
+        end
+
+        def calculate_number_of_vulnerabilities(project_id, latest_vulnerability_id)
+          total_count = 0
+
+          VulnerabilityRead.where('project_id = ? AND vulnerability_id <= ?',
+            project_id, latest_vulnerability_id).each_batch do |batch|
+            total_count += batch.count
+          end
+
+          total_count
+        end
+
+        def persist_statistics(project_id, total_count)
+          bind_params = [total_count, project_id]
+
+          ProjectSecurityStatistics.connection.exec_query(PROJECT_STATISTICS_UPDATE_QUERY, 'UPDATE', bind_params)
+        end
+      end
+    end
+  end
+end

ee/spec/lib/ee/gitlab/background_migration/set_project_vulnerability_count_spec.rb → ee/spec/lib/ee/gitlab/background_migration/set_total_number_of_vulnerabilities_for_existing_projects_spec.rb

@@ -2,129 +2,104 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::BackgroundMigration::SetProjectVulnerabilityCount, feature_category: :vulnerability_management do
-  let!(:user) { table(:users).create!(email: 'author@example.com', username: 'author', projects_limit: 10) }
+RSpec.describe Gitlab::BackgroundMigration::SetTotalNumberOfVulnerabilitiesForExistingProjects, feature_category: :vulnerability_management do
+  let(:users) { table(:users) }
   let(:scanners) { table(:vulnerability_scanners) }
   let(:namespaces) { table(:namespaces) }
   let(:projects) { table(:projects) }
-  let(:project_settings) { table(:project_settings) }
-  let(:project_statistics) { table(:project_statistics) }
+  let(:security_statistics) { table(:project_security_statistics) }
   let(:vulnerabilities) { table(:vulnerabilities) }
   let(:vulnerability_reads) { table(:vulnerability_reads) }
   let(:vulnerability_occurrences) { table(:vulnerability_occurrences) }
   let(:vulnerability_identifiers) { table(:vulnerability_identifiers) }
 
-  let!(:namespace) { namespaces.create!(name: 'Test Namespace', path: 'namespace-path') }
-  let!(:project_namespace_1) { namespaces.create!(name: 'Project Namespace 1', path: 'project_1') }
-  let!(:project_namespace_2) { namespaces.create!(name: 'Project Namespace 2', path: 'project_2') }
-  let!(:project_namespace_3) { namespaces.create!(name: 'Project Namespace 3', path: 'project_3') }
+  let(:user) { users.create!(email: 'john@doe', username: 'john_doe', projects_limit: 10) }
+  let(:namespace) { namespaces.create!(name: 'Test', path: 'test') }
+  let(:project_namespace_1) { namespaces.create!(name: 'Project Namespace 1', path: 'project_1') }
+  let(:project_namespace_2) { namespaces.create!(name: 'Project Namespace 2', path: 'project_2') }
+  let(:project_namespace_3) { namespaces.create!(name: 'Project Namespace 3', path: 'project_3') }
 
-  let!(:project_1) do
+  let(:project_1) do
     projects.create!(
-      name: "proj1",
-      path: "proj1",
+      name: 'project_1',
+      path: 'project_1',
       namespace_id: namespace.id,
       project_namespace_id: project_namespace_1.id
     )
   end
 
-  let!(:project_2) do
+  let(:project_2) do
     projects.create!(
-      name: "proj2",
-      path: "proj2",
+      name: 'project_2',
+      path: 'project_2',
       namespace_id: namespace.id,
       project_namespace_id: project_namespace_2.id
     )
   end
 
-  let!(:project_3) do
+  let(:project_3) do
     projects.create!(
-      name: "proj3",
-      path: "proj3",
+      name: 'project_3',
+      path: 'project_3',
       namespace_id: namespace.id,
       project_namespace_id: project_namespace_3.id
     )
   end
 
-  let!(:project_statistics_1) do
-    project_statistics.create!(project_id: project_1.id, namespace_id: namespace.id, vulnerability_count: 0)
-  end
-
-  let!(:project_statistics_2) do
-    project_statistics.create!(project_id: project_2.id, namespace_id: namespace.id, vulnerability_count: 0)
-  end
-
-  let!(:project_statistics_3) do
-    project_statistics.create!(project_id: project_3.id, namespace_id: namespace.id, vulnerability_count: 0)
-  end
-
-  let!(:project_settings_1) { project_settings.create!(project_id: project_1.id, has_vulnerabilities: true) }
-  let!(:project_settings_2) { project_settings.create!(project_id: project_2.id, has_vulnerabilities: true) }
-  let!(:project_settings_3) { project_settings.create!(project_id: project_3.id, has_vulnerabilities: false) }
+  let!(:project_statistics_1) { security_statistics.create!(project_id: project_1.id, vulnerability_count: 5) }
+  let!(:project_statistics_2) { security_statistics.create!(project_id: project_2.id, vulnerability_count: 0) }
+  let(:project_statistics_3_relation) { security_statistics.where(project_id: project_3.id) }
 
-  let(:migration_instance) do
+  let(:migration) do
     described_class.new(
-      start_id: project_settings.minimum(:project_id),
-      end_id: project_settings.maximum(:project_id),
-      batch_table: :project_settings,
+      start_id: vulnerability_reads.minimum(:project_id),
+      end_id: vulnerability_reads.maximum(:project_id),
+      batch_table: :vulnerability_reads,
       batch_column: :project_id,
       sub_batch_size: 100,
       pause_ms: 0,
-      connection: project_settings.connection
+      connection: vulnerability_reads.connection
     )
   end
 
-  subject(:perform_migration) { migration_instance.perform }
-
   before do
-    scanner_project_1 = scanners.create!(
-      project_id: project_1.id,
-      external_id: "external-id-1",
-      name: "Scanner 1"
-    )
-    scanner_project_2 = scanners.create!(
-      project_id: project_2.id,
-      external_id: "external-id-2",
-      name: "Scanner 2"
-    )
-    create_vulnerability_read(project_1.id, scanner_project_1.id)
-    create_vulnerability_read(project_1.id, scanner_project_1.id)
-    create_vulnerability_read(project_1.id, scanner_project_1.id)
-    create_vulnerability_read(project_2.id, scanner_project_2.id)
-    create_vulnerability_read(project_2.id, scanner_project_2.id)
+    create_vulnerability_read(project_1.id)
+    create_vulnerability_read(project_2.id)
+    create_vulnerability_read(project_3.id)
   end
 
-  it 'sets the vulnerability_count column' do
-    expect do
-      perform_migration
-      project_statistics_1.reload
-      project_statistics_2.reload
-      project_statistics_3.reload
-    end.to change { project_statistics_1.vulnerability_count }.from(0).to(3)
-       .and change { project_statistics_2.vulnerability_count }.from(0).to(2)
-       .and not_change { project_statistics_3.vulnerability_count }.from(0)
+  it 'updates vulnerability_count for vulnerable projects' do
+    expect { migration.perform }.to change { project_statistics_1.reload.vulnerability_count }.to(1)
+                                .and change { project_statistics_2.reload.vulnerability_count }.to(1)
+                                .and change { project_statistics_3_relation.pick(:vulnerability_count) }.to(1)
   end
 
-  def create_vulnerability_read(project_id, scanner_id)
+  def create_vulnerability_read(project_id)
+    scanner = scanners.create!(
+      project_id: project_id,
+      external_id: 'external-id',
+      name: 'Scanner'
+    )
+
     identifier = vulnerability_identifiers.create!(
       project_id: project_id,
       fingerprint: SecureRandom.bytes(20),
-      external_type: "",
-      external_id: "",
-      name: ""
+      external_type: '',
+      external_id: '',
+      name: ''
     )
 
     finding = vulnerability_occurrences.create!(
       project_id: project_id,
       severity: 1,
       report_type: 1,
-      scanner_id: scanner_id,
+      scanner_id: scanner.id,
       primary_identifier_id: identifier.id,
-      project_fingerprint: "",
-      location_fingerprint: "",
-      name: "name",
-      metadata_version: "sast:1.0",
-      raw_metadata: "",
+      project_fingerprint: '',
+      location_fingerprint: '',
+      name: 'name',
+      metadata_version: '15.0',
+      raw_metadata: '',
       uuid: SecureRandom.uuid
     )
 
@@ -142,7 +117,7 @@ def create_vulnerability_read(project_id, scanner_id)
     vulnerability_reads.create!(
       vulnerability_id: vulnerability.id,
       project_id: project_id,
-      scanner_id: scanner_id,
+      scanner_id: scanner.id,
       report_type: vulnerability.report_type,
       severity: vulnerability.severity,
       state: 1,

lib/gitlab/background_migration/set_project_vulnerability_count.rb → lib/gitlab/background_migration/set_total_number_of_vulnerabilities_for_existing_projects.rb

@@ -2,12 +2,13 @@
 
 module Gitlab
   module BackgroundMigration
-    # Background migration to set the vulnerability count in the project_statistics table.
-    class SetProjectVulnerabilityCount < BatchedMigrationJob
+    # Sets the `vulnerability_count` column of `project_security_statistics` table.
+    class SetTotalNumberOfVulnerabilitiesForExistingProjects < BatchedMigrationJob
       feature_category :vulnerability_management
 
       def perform; end
     end
   end
 end
-Gitlab::BackgroundMigration::SetProjectVulnerabilityCount.prepend_mod
+
+Gitlab::BackgroundMigration::SetTotalNumberOfVulnerabilitiesForExistingProjects.prepend_mod

spec/migrations/20240918122045_reschedule_setting_project_vulnerability_count_spec.rb → spec/migrations/20241001115912_queue_set_total_number_of_vulnerabilities_for_existing_projects_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 require_migration!
 
-RSpec.describe RescheduleSettingProjectVulnerabilityCount, feature_category: :vulnerability_management do
+RSpec.describe QueueSetTotalNumberOfVulnerabilitiesForExistingProjects, feature_category: :vulnerability_management do
   let!(:batched_migration) { described_class::MIGRATION }
 
   it 'schedules a new batched migration' do
@@ -14,8 +14,10 @@
 
       migration.after -> {
         expect(batched_migration).to have_scheduled_batched_migration(
-          table_name: :project_settings,
+          gitlab_schema: :gitlab_sec,
+          table_name: :vulnerability_reads,
           column_name: :project_id,
+          batch_class_name: 'LooseIndexScanBatchingStrategy',
           interval: described_class::DELAY_INTERVAL,
           batch_size: described_class::BATCH_SIZE,
           sub_batch_size: described_class::SUB_BATCH_SIZE