Merge branch...

Merge branch 'backfill-multiple-desired-sharding-key-small-table-snippet_user_mentions' into 'master' 

Add & backfill sharding keys for snippet_user_mentions

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



Merged-by: Omar Qunsul <oqunsul@gitlab.com>
Approved-by: Omar Qunsul <oqunsul@gitlab.com>
Co-authored-by: Shubham Kumar <shukumar@gitlab.com>

db/schema_migrations/20250209004156

+398980e7624f9591314c8117929c02c65a259a29f00c662e737b94429739c309
\ No newline at end of file

db/schema_migrations/20250209004157

+38c5a2c2fb5e72a01c6b8004076f62b51403fefd0c9e4014ab7bfcad7db7ab98
\ No newline at end of file

db/schema_migrations/20250209004158

+60c6ea7a1a479512f661ae542f430f1ddb6abe30a57bb2a301de4ef4e9bfed53
\ No newline at end of file

db/structure.sql

@@ -2406,6 +2406,22 @@ RETURN NEW;
 END
 $$;
+CREATE FUNCTION trigger_7495f5e0efcb() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+IF NEW."snippet_project_id" IS NULL THEN
+  SELECT "project_id"
+  INTO NEW."snippet_project_id"
+  FROM "snippets"
+  WHERE "snippets"."id" = NEW."snippet_id";
+END IF;
+
+RETURN NEW;
+
+END
+$$;
+
 CREATE FUNCTION trigger_765cae42cd77() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -2883,6 +2899,22 @@ RETURN NEW;
 END
 $$;
+CREATE FUNCTION trigger_93a5b044f4e8() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
+IF NEW."snippet_organization_id" IS NULL THEN
+  SELECT "organization_id"
+  INTO NEW."snippet_organization_id"
+  FROM "snippets"
+  WHERE "snippets"."id" = NEW."snippet_id";
+END IF;
+
+RETURN NEW;
+
+END
+$$;
+
 CREATE FUNCTION trigger_94514aeadc50() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -22076,7 +22108,9 @@ CREATE TABLE snippet_user_mentions (
     mentioned_users_ids bigint[],
     mentioned_projects_ids bigint[],
     mentioned_groups_ids bigint[],
-    note_id bigint
+    note_id bigint,
+    snippet_project_id bigint,
+    snippet_organization_id bigint
 );
 CREATE SEQUENCE snippet_user_mentions_id_seq
@@ -35680,6 +35714,10 @@ CREATE INDEX index_snippet_repository_storage_moves_on_state ON snippet_reposito
 CREATE UNIQUE INDEX index_snippet_user_mentions_on_note_id ON snippet_user_mentions USING btree (note_id) WHERE (note_id IS NOT NULL);
+CREATE INDEX index_snippet_user_mentions_on_snippet_organization_id ON snippet_user_mentions USING btree (snippet_organization_id);
+
+CREATE INDEX index_snippet_user_mentions_on_snippet_project_id ON snippet_user_mentions USING btree (snippet_project_id);
+
 CREATE INDEX index_snippets_on_author_id ON snippets USING btree (author_id);
 CREATE INDEX index_snippets_on_content_trigram ON snippets USING gin (content gin_trgm_ops);
@@ -39076,6 +39114,8 @@ CREATE TRIGGER trigger_70d3f0bba1de BEFORE INSERT OR UPDATE ON compliance_framew
 CREATE TRIGGER trigger_740afa9807b8 BEFORE INSERT OR UPDATE ON subscription_user_add_on_assignments FOR EACH ROW EXECUTE FUNCTION trigger_740afa9807b8();
+CREATE TRIGGER trigger_7495f5e0efcb BEFORE INSERT OR UPDATE ON snippet_user_mentions FOR EACH ROW EXECUTE FUNCTION trigger_7495f5e0efcb();
+
 CREATE TRIGGER trigger_765cae42cd77 BEFORE INSERT OR UPDATE ON bulk_import_trackers FOR EACH ROW EXECUTE FUNCTION trigger_765cae42cd77();
 CREATE TRIGGER trigger_77d9fbad5b12 BEFORE INSERT OR UPDATE ON packages_debian_project_distribution_keys FOR EACH ROW EXECUTE FUNCTION trigger_77d9fbad5b12();
@@ -39142,6 +39182,8 @@ CREATE TRIGGER trigger_90fa5c6951f1 BEFORE INSERT OR UPDATE ON dast_profiles_tag
 CREATE TRIGGER trigger_9259aae92378 BEFORE INSERT OR UPDATE ON packages_build_infos FOR EACH ROW EXECUTE FUNCTION trigger_9259aae92378();
+CREATE TRIGGER trigger_93a5b044f4e8 BEFORE INSERT OR UPDATE ON snippet_user_mentions FOR EACH ROW EXECUTE FUNCTION trigger_93a5b044f4e8();
+
 CREATE TRIGGER trigger_94514aeadc50 BEFORE INSERT OR UPDATE ON deployment_approvals FOR EACH ROW EXECUTE FUNCTION trigger_94514aeadc50();
 CREATE TRIGGER trigger_951ac22c24d7 BEFORE INSERT OR UPDATE ON required_code_owners_sections FOR EACH ROW EXECUTE FUNCTION trigger_951ac22c24d7();
@@ -40255,6 +40297,9 @@ ALTER TABLE ONLY subscription_user_add_on_assignments
 ALTER TABLE ONLY merge_requests_approval_rules_approver_users
     ADD CONSTRAINT fk_725cca295c FOREIGN KEY (approval_rule_id) REFERENCES merge_requests_approval_rules(id) ON DELETE CASCADE;
+ALTER TABLE ONLY snippet_user_mentions
+    ADD CONSTRAINT fk_7280faac49 FOREIGN KEY (snippet_project_id) REFERENCES projects(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY zentao_tracker_data
     ADD CONSTRAINT fk_72a0e59cd8 FOREIGN KEY (instance_integration_id) REFERENCES instance_integrations(id) ON DELETE CASCADE;
@@ -41149,6 +41194,9 @@ ALTER TABLE ONLY project_control_compliance_statuses
 ALTER TABLE ONLY protected_branches
     ADD CONSTRAINT fk_de9216e774 FOREIGN KEY (namespace_id) REFERENCES namespaces(id) ON DELETE CASCADE;
+ALTER TABLE ONLY snippet_user_mentions
+    ADD CONSTRAINT fk_def441dfc3 FOREIGN KEY (snippet_organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY work_item_number_field_values
     ADD CONSTRAINT fk_df27ad8c35 FOREIGN KEY (work_item_id) REFERENCES issues(id) ON DELETE CASCADE;

lib/gitlab/background_migration/backfill_snippet_user_mentions_snippet_organization_id.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillSnippetUserMentionsSnippetOrganizationId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_snippet_user_mentions_snippet_organization_id
+      feature_category :source_code_management
+    end
+  end
+end

lib/gitlab/background_migration/backfill_snippet_user_mentions_snippet_project_id.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillSnippetUserMentionsSnippetProjectId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_snippet_user_mentions_snippet_project_id
+      feature_category :source_code_management
+    end
+  end
+end

spec/lib/gitlab/background_migration/backfill_snippet_user_mentions_snippet_organization_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillSnippetUserMentionsSnippetOrganizationId,
+  feature_category: :source_code_management,
+  schema: 20250209004154 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :snippet_user_mentions }
+    let(:backfill_column) { :snippet_organization_id }
+    let(:backfill_via_table) { :snippets }
+    let(:backfill_via_column) { :organization_id }
+    let(:backfill_via_foreign_key) { :snippet_id }
+  end
+end

spec/lib/gitlab/background_migration/backfill_snippet_user_mentions_snippet_project_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillSnippetUserMentionsSnippetProjectId,
+  feature_category: :source_code_management,
+  schema: 20250209004149 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :snippet_user_mentions }
+    let(:backfill_column) { :snippet_project_id }
+    let(:backfill_via_table) { :snippets }
+    let(:backfill_via_column) { :project_id }
+    let(:backfill_via_foreign_key) { :snippet_id }
+  end
+end

spec/lib/gitlab/database/sharding_key_spec.rb

@@ -211,6 +211,7 @@
     # To add a table to this list, create an issue under https://gitlab.com/groups/gitlab-org/-/epics/11670.
     # Use https://gitlab.com/gitlab-org/gitlab/-/issues/476206 as an example.
     work_in_progress = {
+      "snippet_user_mentions" => "https://gitlab.com/gitlab-org/gitlab/-/issues/517825",
       "bulk_import_failures" => "https://gitlab.com/gitlab-org/gitlab/-/issues/517824",
       "organization_users" => 'https://gitlab.com/gitlab-org/gitlab/-/issues/476210',
       "push_rules" => 'https://gitlab.com/gitlab-org/gitlab/-/issues/476212',

spec/migrations/20250209004153_queue_backfill_snippet_user_mentions_snippet_project_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillSnippetUserMentionsSnippetProjectId, 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: :snippet_user_mentions,
+          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: [
+            :snippet_project_id,
+            :snippets,
+            :project_id,
+            :snippet_id
+          ]
+        )
+      }
+    end
+  end
+end

spec/migrations/20250209004158_queue_backfill_snippet_user_mentions_snippet_organization_id_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillSnippetUserMentionsSnippetOrganizationId, 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: :snippet_user_mentions,
+          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: [
+            :snippet_organization_id,
+            :snippets,
+            :organization_id,
+            :snippet_id
+          ]
+        )
+      }
+    end
+  end
+end