From ec8450985d455ae76d850285b644a32f52d409eb Mon Sep 17 00:00:00 2001
From: Shubham Kumar <shukumar@gitlab.com>
Date: Tue, 1 Oct 2024 11:48:37 +0000
Subject: [PATCH] Add and backfill project_id for
 approval_project_rules_protected_branche

## What does this MR do and why?

Add and backfill project_id for approval_project_rules_protected_branches.

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/approval_project_rules_protected_branches.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::source code" 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/approval_project_rules_protected_branches.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 @shubhamkrai.

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 create an issue with the
label ~"GitLab Housekeeper" and consider pinging the author of this keep.

Changelog: other
---
 ...roval_project_rules_protected_branches.yml |  1 +
 ...ct_rules_protected_branches_project_id.yml |  9 ++++
 ...proval_project_rules_protected_branches.rb |  9 ++++
 ..._rules_protected_branches_on_project_id.rb | 16 ++++++++
 ..._rules_protected_branches_project_id_fk.rb | 17 ++++++++
 ...s_protected_branches_project_id_trigger.rb | 25 +++++++++++
 ...ect_rules_protected_branches_project_id.rb | 41 +++++++++++++++++++
 db/schema_migrations/20240930160023           |  1 +
 db/schema_migrations/20240930160024           |  1 +
 db/schema_migrations/20240930160025           |  1 +
 db/schema_migrations/20240930160026           |  1 +
 db/schema_migrations/20240930160027           |  1 +
 db/structure.sql                              | 26 +++++++++++-
 ...ect_rules_protected_branches_project_id.rb | 10 +++++
 ...ules_protected_branches_project_id_spec.rb | 16 ++++++++
 ...ules_protected_branches_project_id_spec.rb | 34 +++++++++++++++
 16 files changed, 208 insertions(+), 1 deletion(-)
 create mode 100644 db/docs/batched_background_migrations/backfill_approval_project_rules_protected_branches_project_id.yml
 create mode 100644 db/migrate/20240930160023_add_project_id_to_approval_project_rules_protected_branches.rb
 create mode 100644 db/post_migrate/20240930160024_index_approval_project_rules_protected_branches_on_project_id.rb
 create mode 100644 db/post_migrate/20240930160025_add_approval_project_rules_protected_branches_project_id_fk.rb
 create mode 100644 db/post_migrate/20240930160026_add_approval_project_rules_protected_branches_project_id_trigger.rb
 create mode 100644 db/post_migrate/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id.rb
 create mode 100644 db/schema_migrations/20240930160023
 create mode 100644 db/schema_migrations/20240930160024
 create mode 100644 db/schema_migrations/20240930160025
 create mode 100644 db/schema_migrations/20240930160026
 create mode 100644 db/schema_migrations/20240930160027
 create mode 100644 lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id.rb
 create mode 100644 spec/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id_spec.rb
 create mode 100644 spec/migrations/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id_spec.rb

diff --git a/db/docs/approval_project_rules_protected_branches.yml b/db/docs/approval_project_rules_protected_branches.yml
index b1382866d444a..6ab81dbe594da 100644
--- a/db/docs/approval_project_rules_protected_branches.yml
+++ b/db/docs/approval_project_rules_protected_branches.yml
@@ -19,3 +19,4 @@ desired_sharding_key:
         table: approval_project_rules
         sharding_key: project_id
         belongs_to: approval_project_rule
+desired_sharding_key_migration_job_name: BackfillApprovalProjectRulesProtectedBranchesProjectId
diff --git a/db/docs/batched_background_migrations/backfill_approval_project_rules_protected_branches_project_id.yml b/db/docs/batched_background_migrations/backfill_approval_project_rules_protected_branches_project_id.yml
new file mode 100644
index 0000000000000..d1e4eabfc85b2
--- /dev/null
+++ b/db/docs/batched_background_migrations/backfill_approval_project_rules_protected_branches_project_id.yml
@@ -0,0 +1,9 @@
+---
+migration_job_name: BackfillApprovalProjectRulesProtectedBranchesProjectId
+description: Backfills sharding key `approval_project_rules_protected_branches.project_id` from `approval_project_rules`.
+feature_category: source_code_management
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/167648
+milestone: '17.5'
+queued_migration_version: 20240930160027
+finalize_after: '2024-10-22'
+finalized_by: # version of the migration that finalized this BBM
diff --git a/db/migrate/20240930160023_add_project_id_to_approval_project_rules_protected_branches.rb b/db/migrate/20240930160023_add_project_id_to_approval_project_rules_protected_branches.rb
new file mode 100644
index 0000000000000..f7cc859a0b749
--- /dev/null
+++ b/db/migrate/20240930160023_add_project_id_to_approval_project_rules_protected_branches.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddProjectIdToApprovalProjectRulesProtectedBranches < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+
+  def change
+    add_column :approval_project_rules_protected_branches, :project_id, :bigint
+  end
+end
diff --git a/db/post_migrate/20240930160024_index_approval_project_rules_protected_branches_on_project_id.rb b/db/post_migrate/20240930160024_index_approval_project_rules_protected_branches_on_project_id.rb
new file mode 100644
index 0000000000000..b73e17c365036
--- /dev/null
+++ b/db/post_migrate/20240930160024_index_approval_project_rules_protected_branches_on_project_id.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class IndexApprovalProjectRulesProtectedBranchesOnProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_approval_project_rules_protected_branches_on_project_id'
+
+  def up
+    add_concurrent_index :approval_project_rules_protected_branches, :project_id, name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :approval_project_rules_protected_branches, INDEX_NAME
+  end
+end
diff --git a/db/post_migrate/20240930160025_add_approval_project_rules_protected_branches_project_id_fk.rb b/db/post_migrate/20240930160025_add_approval_project_rules_protected_branches_project_id_fk.rb
new file mode 100644
index 0000000000000..94d516238f2e7
--- /dev/null
+++ b/db/post_migrate/20240930160025_add_approval_project_rules_protected_branches_project_id_fk.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class AddApprovalProjectRulesProtectedBranchesProjectIdFk < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_foreign_key :approval_project_rules_protected_branches, :projects, column: :project_id,
+      on_delete: :cascade
+  end
+
+  def down
+    with_lock_retries do
+      remove_foreign_key :approval_project_rules_protected_branches, column: :project_id
+    end
+  end
+end
diff --git a/db/post_migrate/20240930160026_add_approval_project_rules_protected_branches_project_id_trigger.rb b/db/post_migrate/20240930160026_add_approval_project_rules_protected_branches_project_id_trigger.rb
new file mode 100644
index 0000000000000..091491087b3f6
--- /dev/null
+++ b/db/post_migrate/20240930160026_add_approval_project_rules_protected_branches_project_id_trigger.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+class AddApprovalProjectRulesProtectedBranchesProjectIdTrigger < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+
+  def up
+    install_sharding_key_assignment_trigger(
+      table: :approval_project_rules_protected_branches,
+      sharding_key: :project_id,
+      parent_table: :approval_project_rules,
+      parent_sharding_key: :project_id,
+      foreign_key: :approval_project_rule_id
+    )
+  end
+
+  def down
+    remove_sharding_key_assignment_trigger(
+      table: :approval_project_rules_protected_branches,
+      sharding_key: :project_id,
+      parent_table: :approval_project_rules,
+      parent_sharding_key: :project_id,
+      foreign_key: :approval_project_rule_id
+    )
+  end
+end
diff --git a/db/post_migrate/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id.rb b/db/post_migrate/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id.rb
new file mode 100644
index 0000000000000..f12f7b836e309
--- /dev/null
+++ b/db/post_migrate/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+class QueueBackfillApprovalProjectRulesProtectedBranchesProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.5'
+  restrict_gitlab_migration gitlab_schema: :gitlab_main_cell
+
+  MIGRATION = "BackfillApprovalProjectRulesProtectedBranchesProjectId"
+  DELAY_INTERVAL = 2.minutes
+  BATCH_SIZE = 1000
+  SUB_BATCH_SIZE = 100
+
+  def up
+    queue_batched_background_migration(
+      MIGRATION,
+      :approval_project_rules_protected_branches,
+      :approval_project_rule_id,
+      :project_id,
+      :approval_project_rules,
+      :project_id,
+      :approval_project_rule_id,
+      job_interval: DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      batch_class_name: 'LooseIndexScanBatchingStrategy',
+      sub_batch_size: SUB_BATCH_SIZE
+    )
+  end
+
+  def down
+    delete_batched_background_migration(
+      MIGRATION,
+      :approval_project_rules_protected_branches,
+      :approval_project_rule_id,
+      [
+        :project_id,
+        :approval_project_rules,
+        :project_id,
+        :approval_project_rule_id
+      ]
+    )
+  end
+end
diff --git a/db/schema_migrations/20240930160023 b/db/schema_migrations/20240930160023
new file mode 100644
index 0000000000000..935a1bd81fe67
--- /dev/null
+++ b/db/schema_migrations/20240930160023
@@ -0,0 +1 @@
+e4246085fc454e022e500fdb3b5dbb5d3c78a71fe5107d41d615e945a80f31e5
\ No newline at end of file
diff --git a/db/schema_migrations/20240930160024 b/db/schema_migrations/20240930160024
new file mode 100644
index 0000000000000..4222bfd06e3f0
--- /dev/null
+++ b/db/schema_migrations/20240930160024
@@ -0,0 +1 @@
+d0525a4ec76c7847e59531e5fb6d5137b2121e1dd06ade1680b283bfcf28c53d
\ No newline at end of file
diff --git a/db/schema_migrations/20240930160025 b/db/schema_migrations/20240930160025
new file mode 100644
index 0000000000000..9bbb7ee80afda
--- /dev/null
+++ b/db/schema_migrations/20240930160025
@@ -0,0 +1 @@
+fd42f51e9904d94a6870d24f9ca017b0047b183f92cf1481268667a17e2477f9
\ No newline at end of file
diff --git a/db/schema_migrations/20240930160026 b/db/schema_migrations/20240930160026
new file mode 100644
index 0000000000000..2b2887b90b1cd
--- /dev/null
+++ b/db/schema_migrations/20240930160026
@@ -0,0 +1 @@
+d53da99c66293fa9e53fdef556f3134d490fe66e7bc7bed08093bea9f686f691
\ No newline at end of file
diff --git a/db/schema_migrations/20240930160027 b/db/schema_migrations/20240930160027
new file mode 100644
index 0000000000000..62be59b163e56
--- /dev/null
+++ b/db/schema_migrations/20240930160027
@@ -0,0 +1 @@
+6f49ec85bcf41062ff0c9165444a74a8119ab500392221bc6e8f0a75d3566c45
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index c97e1223e2651..57bd91ce4d06f 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -841,6 +841,22 @@ RETURN NEW;
 END
 $$;
 
+CREATE FUNCTION trigger_0a29d4d42b62() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+IF NEW."project_id" IS NULL THEN
+  SELECT "project_id"
+  INTO NEW."project_id"
+  FROM "approval_project_rules"
+  WHERE "approval_project_rules"."id" = NEW."approval_project_rule_id";
+END IF;
+
+RETURN NEW;
+
+END
+$$;
+
 CREATE FUNCTION trigger_0da002390fdc() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -6620,7 +6636,8 @@ ALTER SEQUENCE approval_project_rules_id_seq OWNED BY approval_project_rules.id;
 
 CREATE TABLE approval_project_rules_protected_branches (
     approval_project_rule_id bigint NOT NULL,
-    protected_branch_id bigint NOT NULL
+    protected_branch_id bigint NOT NULL,
+    project_id bigint
 );
 
 CREATE TABLE approval_project_rules_users (
@@ -27499,6 +27516,8 @@ CREATE INDEX index_approval_project_rules_on_project_id ON approval_project_rule
 
 CREATE INDEX index_approval_project_rules_on_rule_type ON approval_project_rules USING btree (rule_type);
 
+CREATE INDEX index_approval_project_rules_protected_branches_on_project_id ON approval_project_rules_protected_branches USING btree (project_id);
+
 CREATE INDEX index_approval_project_rules_protected_branches_pb_id ON approval_project_rules_protected_branches USING btree (protected_branch_id);
 
 CREATE INDEX index_approval_project_rules_report_type ON approval_project_rules USING btree (report_type);
@@ -33299,6 +33318,8 @@ CREATE TRIGGER trigger_05ce163deddf BEFORE INSERT OR UPDATE ON status_check_resp
 
 CREATE TRIGGER trigger_0a1b0adcf686 BEFORE INSERT OR UPDATE ON packages_debian_project_components FOR EACH ROW EXECUTE FUNCTION trigger_0a1b0adcf686();
 
+CREATE TRIGGER trigger_0a29d4d42b62 BEFORE INSERT OR UPDATE ON approval_project_rules_protected_branches FOR EACH ROW EXECUTE FUNCTION trigger_0a29d4d42b62();
+
 CREATE TRIGGER trigger_0da002390fdc BEFORE INSERT OR UPDATE ON operations_feature_flags_issues FOR EACH ROW EXECUTE FUNCTION trigger_0da002390fdc();
 
 CREATE TRIGGER trigger_0e13f214e504 BEFORE INSERT OR UPDATE ON merge_request_assignment_events FOR EACH ROW EXECUTE FUNCTION trigger_0e13f214e504();
@@ -34910,6 +34931,9 @@ ALTER TABLE ONLY user_preferences
 ALTER TABLE ONLY packages_debian_group_components
     ADD CONSTRAINT fk_e63e8ee3b1 FOREIGN KEY (group_id) REFERENCES namespaces(id) ON DELETE CASCADE;
 
+ALTER TABLE ONLY approval_project_rules_protected_branches
+    ADD CONSTRAINT fk_e6ee913fc2 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY merge_requests
     ADD CONSTRAINT fk_e719a85f8a FOREIGN KEY (author_id) REFERENCES users(id) ON DELETE SET NULL;
 
diff --git a/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id.rb b/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id.rb
new file mode 100644
index 0000000000000..01dddc6c374e5
--- /dev/null
+++ b/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillApprovalProjectRulesProtectedBranchesProjectId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_approval_project_rules_protected_branches_project_id
+      feature_category :source_code_management
+    end
+  end
+end
diff --git a/spec/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id_spec.rb b/spec/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id_spec.rb
new file mode 100644
index 0000000000000..5bd1023f15657
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_approval_project_rules_protected_branches_project_id_spec.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillApprovalProjectRulesProtectedBranchesProjectId,
+  feature_category: :source_code_management,
+  schema: 20240930160023 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :approval_project_rules_protected_branches }
+    let(:batch_column) { :approval_project_rule_id }
+    let(:backfill_column) { :project_id }
+    let(:backfill_via_table) { :approval_project_rules }
+    let(:backfill_via_column) { :project_id }
+    let(:backfill_via_foreign_key) { :approval_project_rule_id }
+  end
+end
diff --git a/spec/migrations/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id_spec.rb b/spec/migrations/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id_spec.rb
new file mode 100644
index 0000000000000..353bacfafa9f2
--- /dev/null
+++ b/spec/migrations/20240930160027_queue_backfill_approval_project_rules_protected_branches_project_id_spec.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillApprovalProjectRulesProtectedBranchesProjectId, feature_category: :source_code_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: :approval_project_rules_protected_branches,
+          column_name: :approval_project_rule_id,
+          interval: described_class::DELAY_INTERVAL,
+          batch_size: described_class::BATCH_SIZE,
+          batch_class_name: 'LooseIndexScanBatchingStrategy',
+          sub_batch_size: described_class::SUB_BATCH_SIZE,
+          gitlab_schema: :gitlab_main_cell,
+          job_arguments: [
+            :project_id,
+            :approval_project_rules,
+            :project_id,
+            :approval_project_rule_id
+          ]
+        )
+      }
+    end
+  end
+end
-- 
GitLab