From 8709db853647b5d2afa117c34d1ffda55dde47d1 Mon Sep 17 00:00:00 2001
From: bmarjanovic <bmarjanovic@gitlab.com>
Date: Fri, 21 Jun 2024 13:59:49 +0200
Subject: [PATCH] Add and backfill project_id for
 vulnerability_finding_evidences

## What does this MR do and why?

Add and backfill project_id for vulnerability_finding_evidences.

This table has a
[desired sharding key](https://docs.gitlab.com/ee/development/database/multiple_databases.html#define-a-desired_sharding_key-to-automatically-backfill-a-sharding_key)
configured ([view configuration](https://gitlab.com/gitlab-org/gitlab/-/blob/master/db/docs/vulnerability_finding_evidences.yml)).

This merge request is the first step towards transforming the desired sharding key into a
[sharding key](https://docs.gitlab.com/ee/development/database/multiple_databases.html#defining-a-sharding-key-for-all-cell-local-tables).

This involves three changes:

- Adding a new column that will serve as the sharding key (along with the relevant index and foreign key).
- Populating the sharding key when new records are created by adding a database function and trigger.
- Scheduling a [batched background migration](https://docs.gitlab.com/ee/development/database/batched_background_migrations.html)
  to set the sharding key for existing records.

Once the background migration has completed, a second merge request will be created to finalize the background
migration and validate the not null constraint.

## How to verify

We have assigned a random backend engineer from ~"group::threat insights" to review these changes. Please review this merge
request from a ~backend perspective. The main thing we are looking to verify is that the added column and association
match the values specified by the [desired sharding key](https://gitlab.com/gitlab-org/gitlab/-/blob/master/db/docs/vulnerability_finding_evidences.yml)
configuration and that backfilling the column from this other table makes sense in the context of this feature.

When you are finished, please:

1. Trigger the [database testing pipeline](https://docs.gitlab.com/ee/development/database/database_migration_pipeline.html)
   as instructed by Danger.
1. Request a review from the ~backend maintainer and ~database reviewer suggested by Danger.

If you have any questions or concerns, reach out to `@tigerwnz` or @manojmj.

This merge request was generated by a once off keep implemented in
https://gitlab.com/gitlab-org/gitlab/-/merge_requests/143774

This change was generated by
[gitlab-housekeeper](https://gitlab.com/gitlab-org/gitlab/-/tree/master/gems/gitlab-housekeeper)
using the Keeps::BackfillDesiredShardingKeySmallTable keep.

To provide feedback on your experience with `gitlab-housekeeper` please comment in
<https://gitlab.com/gitlab-org/gitlab/-/issues/442003>.

Changelog: other
---
 ...erability_finding_evidences_project_id.yml |  9 +++++
 db/docs/vulnerability_finding_evidences.yml   |  1 +
 ...t_id_to_vulnerability_finding_evidences.rb |  9 +++++
 ...ability_finding_evidences_on_project_id.rb | 16 ++++++++
 ...ability_finding_evidences_project_id_fk.rb | 16 ++++++++
 ...ty_finding_evidences_project_id_trigger.rb | 25 ++++++++++++
 ...nerability_finding_evidences_project_id.rb | 40 +++++++++++++++++++
 db/schema_migrations/20240621115727           |  1 +
 db/schema_migrations/20240621115728           |  1 +
 db/schema_migrations/20240621115729           |  1 +
 db/schema_migrations/20240621115730           |  1 +
 db/schema_migrations/20240621115731           |  1 +
 db/structure.sql                              | 26 +++++++++++-
 ...nerability_finding_evidences_project_id.rb | 10 +++++
 ...ility_finding_evidences_project_id_spec.rb | 15 +++++++
 ...ility_finding_evidences_project_id_spec.rb | 33 +++++++++++++++
 16 files changed, 204 insertions(+), 1 deletion(-)
 create mode 100644 db/docs/batched_background_migrations/backfill_vulnerability_finding_evidences_project_id.yml
 create mode 100644 db/migrate/20240621115727_add_project_id_to_vulnerability_finding_evidences.rb
 create mode 100644 db/post_migrate/20240621115728_index_vulnerability_finding_evidences_on_project_id.rb
 create mode 100644 db/post_migrate/20240621115729_add_vulnerability_finding_evidences_project_id_fk.rb
 create mode 100644 db/post_migrate/20240621115730_add_vulnerability_finding_evidences_project_id_trigger.rb
 create mode 100644 db/post_migrate/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id.rb
 create mode 100644 db/schema_migrations/20240621115727
 create mode 100644 db/schema_migrations/20240621115728
 create mode 100644 db/schema_migrations/20240621115729
 create mode 100644 db/schema_migrations/20240621115730
 create mode 100644 db/schema_migrations/20240621115731
 create mode 100644 lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id.rb
 create mode 100644 spec/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id_spec.rb
 create mode 100644 spec/migrations/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id_spec.rb

diff --git a/db/docs/batched_background_migrations/backfill_vulnerability_finding_evidences_project_id.yml b/db/docs/batched_background_migrations/backfill_vulnerability_finding_evidences_project_id.yml
new file mode 100644
index 0000000000000..26cb65f9c0c1e
--- /dev/null
+++ b/db/docs/batched_background_migrations/backfill_vulnerability_finding_evidences_project_id.yml
@@ -0,0 +1,9 @@
+---
+migration_job_name: BackfillVulnerabilityFindingEvidencesProjectId
+description: Backfills sharding key `vulnerability_finding_evidences.project_id` from `vulnerability_occurrences`.
+feature_category: vulnerability_management
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/157032
+milestone: '17.2'
+queued_migration_version: 20240621115731
+finalize_after: '2024-07-22'
+finalized_by: # version of the migration that finalized this BBM
diff --git a/db/docs/vulnerability_finding_evidences.yml b/db/docs/vulnerability_finding_evidences.yml
index 35dc895aeaaa8..654c89dcae616 100644
--- a/db/docs/vulnerability_finding_evidences.yml
+++ b/db/docs/vulnerability_finding_evidences.yml
@@ -19,3 +19,4 @@ desired_sharding_key:
         table: vulnerability_occurrences
         sharding_key: project_id
         belongs_to: finding
+desired_sharding_key_migration_job_name: BackfillVulnerabilityFindingEvidencesProjectId
diff --git a/db/migrate/20240621115727_add_project_id_to_vulnerability_finding_evidences.rb b/db/migrate/20240621115727_add_project_id_to_vulnerability_finding_evidences.rb
new file mode 100644
index 0000000000000..eb4363f651d4f
--- /dev/null
+++ b/db/migrate/20240621115727_add_project_id_to_vulnerability_finding_evidences.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddProjectIdToVulnerabilityFindingEvidences < Gitlab::Database::Migration[2.2]
+  milestone '17.2'
+
+  def change
+    add_column :vulnerability_finding_evidences, :project_id, :bigint
+  end
+end
diff --git a/db/post_migrate/20240621115728_index_vulnerability_finding_evidences_on_project_id.rb b/db/post_migrate/20240621115728_index_vulnerability_finding_evidences_on_project_id.rb
new file mode 100644
index 0000000000000..089a47ba5eb63
--- /dev/null
+++ b/db/post_migrate/20240621115728_index_vulnerability_finding_evidences_on_project_id.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class IndexVulnerabilityFindingEvidencesOnProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.2'
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_vulnerability_finding_evidences_on_project_id'
+
+  def up
+    add_concurrent_index :vulnerability_finding_evidences, :project_id, name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :vulnerability_finding_evidences, INDEX_NAME
+  end
+end
diff --git a/db/post_migrate/20240621115729_add_vulnerability_finding_evidences_project_id_fk.rb b/db/post_migrate/20240621115729_add_vulnerability_finding_evidences_project_id_fk.rb
new file mode 100644
index 0000000000000..fa0e70d5d1035
--- /dev/null
+++ b/db/post_migrate/20240621115729_add_vulnerability_finding_evidences_project_id_fk.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class AddVulnerabilityFindingEvidencesProjectIdFk < Gitlab::Database::Migration[2.2]
+  milestone '17.2'
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_foreign_key :vulnerability_finding_evidences, :projects, column: :project_id, on_delete: :cascade
+  end
+
+  def down
+    with_lock_retries do
+      remove_foreign_key :vulnerability_finding_evidences, column: :project_id
+    end
+  end
+end
diff --git a/db/post_migrate/20240621115730_add_vulnerability_finding_evidences_project_id_trigger.rb b/db/post_migrate/20240621115730_add_vulnerability_finding_evidences_project_id_trigger.rb
new file mode 100644
index 0000000000000..a0d2ac9621c8c
--- /dev/null
+++ b/db/post_migrate/20240621115730_add_vulnerability_finding_evidences_project_id_trigger.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+class AddVulnerabilityFindingEvidencesProjectIdTrigger < Gitlab::Database::Migration[2.2]
+  milestone '17.2'
+
+  def up
+    install_sharding_key_assignment_trigger(
+      table: :vulnerability_finding_evidences,
+      sharding_key: :project_id,
+      parent_table: :vulnerability_occurrences,
+      parent_sharding_key: :project_id,
+      foreign_key: :vulnerability_occurrence_id
+    )
+  end
+
+  def down
+    remove_sharding_key_assignment_trigger(
+      table: :vulnerability_finding_evidences,
+      sharding_key: :project_id,
+      parent_table: :vulnerability_occurrences,
+      parent_sharding_key: :project_id,
+      foreign_key: :vulnerability_occurrence_id
+    )
+  end
+end
diff --git a/db/post_migrate/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id.rb b/db/post_migrate/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id.rb
new file mode 100644
index 0000000000000..f8e08f3b5aab1
--- /dev/null
+++ b/db/post_migrate/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+class QueueBackfillVulnerabilityFindingEvidencesProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.2'
+  restrict_gitlab_migration gitlab_schema: :gitlab_main_cell
+
+  MIGRATION = "BackfillVulnerabilityFindingEvidencesProjectId"
+  DELAY_INTERVAL = 2.minutes
+  BATCH_SIZE = 1000
+  SUB_BATCH_SIZE = 100
+
+  def up
+    queue_batched_background_migration(
+      MIGRATION,
+      :vulnerability_finding_evidences,
+      :id,
+      :project_id,
+      :vulnerability_occurrences,
+      :project_id,
+      :vulnerability_occurrence_id,
+      job_interval: DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      sub_batch_size: SUB_BATCH_SIZE
+    )
+  end
+
+  def down
+    delete_batched_background_migration(
+      MIGRATION,
+      :vulnerability_finding_evidences,
+      :id,
+      [
+        :project_id,
+        :vulnerability_occurrences,
+        :project_id,
+        :vulnerability_occurrence_id
+      ]
+    )
+  end
+end
diff --git a/db/schema_migrations/20240621115727 b/db/schema_migrations/20240621115727
new file mode 100644
index 0000000000000..da693d775f804
--- /dev/null
+++ b/db/schema_migrations/20240621115727
@@ -0,0 +1 @@
+87e10a60c27d45dc9ab5e31d0290aad07070e39c25aa80bb29d1e52f5612d96c
\ No newline at end of file
diff --git a/db/schema_migrations/20240621115728 b/db/schema_migrations/20240621115728
new file mode 100644
index 0000000000000..b1400497211ff
--- /dev/null
+++ b/db/schema_migrations/20240621115728
@@ -0,0 +1 @@
+f53a3e8e6812bdd14ff9890e6e95178c660f4cc7ab88eb5be4ae20bc1225086f
\ No newline at end of file
diff --git a/db/schema_migrations/20240621115729 b/db/schema_migrations/20240621115729
new file mode 100644
index 0000000000000..16d9bc97e6527
--- /dev/null
+++ b/db/schema_migrations/20240621115729
@@ -0,0 +1 @@
+1cc562767546e5cfb017dc51a1bf0946893ac4b6540cd33b5984e2b274fae7a5
\ No newline at end of file
diff --git a/db/schema_migrations/20240621115730 b/db/schema_migrations/20240621115730
new file mode 100644
index 0000000000000..c248c122ed71e
--- /dev/null
+++ b/db/schema_migrations/20240621115730
@@ -0,0 +1 @@
+3740b89f47ed72101f15f1b936d99f9a43ebf5d9625c574d86a75c51a47ac72c
\ No newline at end of file
diff --git a/db/schema_migrations/20240621115731 b/db/schema_migrations/20240621115731
new file mode 100644
index 0000000000000..7bbca3c208cfc
--- /dev/null
+++ b/db/schema_migrations/20240621115731
@@ -0,0 +1 @@
+11873f1781c7bc9ae4ffa733c91c7cdcf10d6f660eb1bd653da762c6d67bd55d
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index 7397914c3a891..7338f6171c90d 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -1409,6 +1409,22 @@ RETURN NEW;
 END
 $$;
 
+CREATE FUNCTION trigger_a7e0fb195210() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+IF NEW."project_id" IS NULL THEN
+  SELECT "project_id"
+  INTO NEW."project_id"
+  FROM "vulnerability_occurrences"
+  WHERE "vulnerability_occurrences"."id" = NEW."vulnerability_occurrence_id";
+END IF;
+
+RETURN NEW;
+
+END
+$$;
+
 CREATE FUNCTION trigger_af3f17817e4d() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -18679,7 +18695,8 @@ CREATE TABLE vulnerability_finding_evidences (
     created_at timestamp with time zone NOT NULL,
     updated_at timestamp with time zone NOT NULL,
     vulnerability_occurrence_id bigint NOT NULL,
-    data jsonb DEFAULT '{}'::jsonb NOT NULL
+    data jsonb DEFAULT '{}'::jsonb NOT NULL,
+    project_id bigint
 );
 
 CREATE SEQUENCE vulnerability_finding_evidences_id_seq
@@ -29139,6 +29156,8 @@ CREATE INDEX index_vulnerability_feedback_on_merge_request_id ON vulnerability_f
 
 CREATE INDEX index_vulnerability_feedback_on_pipeline_id ON vulnerability_feedback USING btree (pipeline_id);
 
+CREATE INDEX index_vulnerability_finding_evidences_on_project_id ON vulnerability_finding_evidences USING btree (project_id);
+
 CREATE INDEX index_vulnerability_finding_signatures_on_finding_id ON vulnerability_finding_signatures USING btree (finding_id);
 
 CREATE INDEX index_vulnerability_finding_signatures_on_project_id ON vulnerability_finding_signatures USING btree (project_id);
@@ -31291,6 +31310,8 @@ CREATE TRIGGER trigger_a253cb3cacdf BEFORE INSERT OR UPDATE ON dora_daily_metric
 
 CREATE TRIGGER trigger_a4e4fb2451d9 BEFORE INSERT OR UPDATE ON epic_user_mentions FOR EACH ROW EXECUTE FUNCTION trigger_a4e4fb2451d9();
 
+CREATE TRIGGER trigger_a7e0fb195210 BEFORE INSERT OR UPDATE ON vulnerability_finding_evidences FOR EACH ROW EXECUTE FUNCTION trigger_a7e0fb195210();
+
 CREATE TRIGGER trigger_af3f17817e4d BEFORE INSERT OR UPDATE ON protected_tag_create_access_levels FOR EACH ROW EXECUTE FUNCTION trigger_af3f17817e4d();
 
 CREATE TRIGGER trigger_b2612138515d BEFORE INSERT OR UPDATE ON project_relation_exports FOR EACH ROW EXECUTE FUNCTION trigger_b2612138515d();
@@ -32760,6 +32781,9 @@ ALTER TABLE ONLY protected_tag_create_access_levels
 ALTER TABLE ONLY application_settings
     ADD CONSTRAINT fk_f9867b3540 FOREIGN KEY (web_ide_oauth_application_id) REFERENCES oauth_applications(id) ON DELETE SET NULL;
 
+ALTER TABLE ONLY vulnerability_finding_evidences
+    ADD CONSTRAINT fk_fa3efd4e94 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE;
+
 ALTER TABLE p_ci_stages
     ADD CONSTRAINT fk_fb57e6cc56 FOREIGN KEY (pipeline_id) REFERENCES ci_pipelines(id) ON DELETE CASCADE;
 
diff --git a/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id.rb b/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id.rb
new file mode 100644
index 0000000000000..a517cf40c1415
--- /dev/null
+++ b/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillVulnerabilityFindingEvidencesProjectId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_vulnerability_finding_evidences_project_id
+      feature_category :vulnerability_management
+    end
+  end
+end
diff --git a/spec/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id_spec.rb b/spec/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id_spec.rb
new file mode 100644
index 0000000000000..bde817089b718
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_vulnerability_finding_evidences_project_id_spec.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillVulnerabilityFindingEvidencesProjectId,
+  feature_category: :vulnerability_management,
+  schema: 20240621115727 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :vulnerability_finding_evidences }
+    let(:backfill_column) { :project_id }
+    let(:backfill_via_table) { :vulnerability_occurrences }
+    let(:backfill_via_column) { :project_id }
+    let(:backfill_via_foreign_key) { :vulnerability_occurrence_id }
+  end
+end
diff --git a/spec/migrations/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id_spec.rb b/spec/migrations/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id_spec.rb
new file mode 100644
index 0000000000000..9371ed33f2155
--- /dev/null
+++ b/spec/migrations/20240621115731_queue_backfill_vulnerability_finding_evidences_project_id_spec.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillVulnerabilityFindingEvidencesProjectId, feature_category: :vulnerability_management do
+  let!(:batched_migration) { described_class::MIGRATION }
+
+  it 'schedules a new batched migration' do
+    reversible_migration do |migration|
+      migration.before -> {
+        expect(batched_migration).not_to have_scheduled_batched_migration
+      }
+
+      migration.after -> {
+        expect(batched_migration).to have_scheduled_batched_migration(
+          table_name: :vulnerability_finding_evidences,
+          column_name: :id,
+          interval: described_class::DELAY_INTERVAL,
+          batch_size: described_class::BATCH_SIZE,
+          sub_batch_size: described_class::SUB_BATCH_SIZE,
+          gitlab_schema: :gitlab_main_cell,
+          job_arguments: [
+            :project_id,
+            :vulnerability_occurrences,
+            :project_id,
+            :vulnerability_occurrence_id
+          ]
+        )
+      }
+    end
+  end
+end
-- 
GitLab