From bbcb0efba3bc4c7b6fa577fa14c6c89721255aef Mon Sep 17 00:00:00 2001
From: Shubham Kumar <shukumar@gitlab.com>
Date: Thu, 13 Mar 2025 16:45:37 +0100
Subject: [PATCH] Add & backfill sharding keys for snippet_repositories

Add and backfill multiple sharding keys for snippet_repositories.

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/snippet_repositories.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.

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/snippet_repositories.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::BackfillMultipleDesiredShardingKeySmallTable 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
---
 ...t_repositories_snippet_organization_id.yml |  8 ++++
 ...nippet_repositories_snippet_project_id.yml |  8 ++++
 db/docs/snippet_repositories.yml              |  3 ++
 ...ppet_project_id_to_snippet_repositories.rb |  9 ++++
 ...organization_id_to_snippet_repositories.rb |  9 ++++
 ...ppet_repositories_on_snippet_project_id.rb | 16 +++++++
 ...ppet_repositories_snippet_project_id_fk.rb | 16 +++++++
 ...repositories_snippet_project_id_trigger.rb | 25 ++++++++++
 ...snippet_repositories_snippet_project_id.rb | 40 ++++++++++++++++
 ...repositories_on_snippet_organization_id.rb | 16 +++++++
 ...repositories_snippet_organization_id_fk.rb | 17 +++++++
 ...itories_snippet_organization_id_trigger.rb | 25 ++++++++++
 ...et_repositories_snippet_organization_id.rb | 40 ++++++++++++++++
 db/schema_migrations/20241211134706           |  1 +
 db/schema_migrations/20241211134707           |  1 +
 db/schema_migrations/20241211134708           |  1 +
 db/schema_migrations/20241211134709           |  1 +
 db/schema_migrations/20241211134710           |  1 +
 db/schema_migrations/20241211134711           |  1 +
 db/schema_migrations/20241211134712           |  1 +
 db/schema_migrations/20241211134713           |  1 +
 db/schema_migrations/20241211134714           |  1 +
 db/schema_migrations/20241211134715           |  1 +
 db/structure.sql                              | 48 +++++++++++++++++++
 ...et_repositories_snippet_organization_id.rb | 10 ++++
 ...snippet_repositories_snippet_project_id.rb | 10 ++++
 .../snippets/user_creates_snippet_spec.rb     |  2 +
 ...positories_snippet_organization_id_spec.rb | 16 +++++++
 ...et_repositories_snippet_project_id_spec.rb | 16 +++++++
 ...et_repositories_snippet_project_id_spec.rb | 33 +++++++++++++
 ...positories_snippet_organization_id_spec.rb | 33 +++++++++++++
 .../graphql/mutations/snippets/create_spec.rb |  4 ++
 spec/services/snippets/create_service_spec.rb |  4 ++
 33 files changed, 418 insertions(+)
 create mode 100644 db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_organization_id.yml
 create mode 100644 db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_project_id.yml
 create mode 100644 db/migrate/20241211134706_add_snippet_project_id_to_snippet_repositories.rb
 create mode 100644 db/migrate/20241211134711_add_snippet_organization_id_to_snippet_repositories.rb
 create mode 100644 db/post_migrate/20241211134707_index_snippet_repositories_on_snippet_project_id.rb
 create mode 100644 db/post_migrate/20241211134708_add_snippet_repositories_snippet_project_id_fk.rb
 create mode 100644 db/post_migrate/20241211134709_add_snippet_repositories_snippet_project_id_trigger.rb
 create mode 100644 db/post_migrate/20241211134710_queue_backfill_snippet_repositories_snippet_project_id.rb
 create mode 100644 db/post_migrate/20241211134712_index_snippet_repositories_on_snippet_organization_id.rb
 create mode 100644 db/post_migrate/20241211134713_add_snippet_repositories_snippet_organization_id_fk.rb
 create mode 100644 db/post_migrate/20241211134714_add_snippet_repositories_snippet_organization_id_trigger.rb
 create mode 100644 db/post_migrate/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id.rb
 create mode 100644 db/schema_migrations/20241211134706
 create mode 100644 db/schema_migrations/20241211134707
 create mode 100644 db/schema_migrations/20241211134708
 create mode 100644 db/schema_migrations/20241211134709
 create mode 100644 db/schema_migrations/20241211134710
 create mode 100644 db/schema_migrations/20241211134711
 create mode 100644 db/schema_migrations/20241211134712
 create mode 100644 db/schema_migrations/20241211134713
 create mode 100644 db/schema_migrations/20241211134714
 create mode 100644 db/schema_migrations/20241211134715
 create mode 100644 lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id.rb
 create mode 100644 lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id.rb
 create mode 100644 spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id_spec.rb
 create mode 100644 spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id_spec.rb
 create mode 100644 spec/migrations/20241211134710_queue_backfill_snippet_repositories_snippet_project_id_spec.rb
 create mode 100644 spec/migrations/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id_spec.rb

diff --git a/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_organization_id.yml b/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_organization_id.yml
new file mode 100644
index 000000000000..121472bfeeeb
--- /dev/null
+++ b/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_organization_id.yml
@@ -0,0 +1,8 @@
+---
+migration_job_name: BackfillSnippetRepositoriesSnippetOrganizationId
+description: Backfills sharding key `snippet_repositories.snippet_organization_id` from `snippets`.
+feature_category: source_code_management
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/175410
+milestone: '17.10'
+queued_migration_version: 20241211134715
+finalized_by: # version of the migration that finalized this BBM
diff --git a/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_project_id.yml b/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_project_id.yml
new file mode 100644
index 000000000000..84acd05fa7c3
--- /dev/null
+++ b/db/docs/batched_background_migrations/backfill_snippet_repositories_snippet_project_id.yml
@@ -0,0 +1,8 @@
+---
+migration_job_name: BackfillSnippetRepositoriesSnippetProjectId
+description: Backfills sharding key `snippet_repositories.snippet_project_id` from `snippets`.
+feature_category: source_code_management
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/175410
+milestone: '17.10'
+queued_migration_version: 20241211134710
+finalized_by: # version of the migration that finalized this BBM
diff --git a/db/docs/snippet_repositories.yml b/db/docs/snippet_repositories.yml
index 3d96ecc4ab24..40d10b5236b1 100644
--- a/db/docs/snippet_repositories.yml
+++ b/db/docs/snippet_repositories.yml
@@ -28,3 +28,6 @@ desired_sharding_key:
         table: snippets
         sharding_key: organization_id
         belongs_to: snippet
+desired_sharding_key_migration_job_name:
+- BackfillSnippetRepositoriesSnippetProjectId
+- BackfillSnippetRepositoriesSnippetOrganizationId
diff --git a/db/migrate/20241211134706_add_snippet_project_id_to_snippet_repositories.rb b/db/migrate/20241211134706_add_snippet_project_id_to_snippet_repositories.rb
new file mode 100644
index 000000000000..17a81f99077f
--- /dev/null
+++ b/db/migrate/20241211134706_add_snippet_project_id_to_snippet_repositories.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddSnippetProjectIdToSnippetRepositories < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+
+  def change
+    add_column :snippet_repositories, :snippet_project_id, :bigint
+  end
+end
diff --git a/db/migrate/20241211134711_add_snippet_organization_id_to_snippet_repositories.rb b/db/migrate/20241211134711_add_snippet_organization_id_to_snippet_repositories.rb
new file mode 100644
index 000000000000..bcc8b3eb7647
--- /dev/null
+++ b/db/migrate/20241211134711_add_snippet_organization_id_to_snippet_repositories.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddSnippetOrganizationIdToSnippetRepositories < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+
+  def change
+    add_column :snippet_repositories, :snippet_organization_id, :bigint
+  end
+end
diff --git a/db/post_migrate/20241211134707_index_snippet_repositories_on_snippet_project_id.rb b/db/post_migrate/20241211134707_index_snippet_repositories_on_snippet_project_id.rb
new file mode 100644
index 000000000000..0b9d3814883b
--- /dev/null
+++ b/db/post_migrate/20241211134707_index_snippet_repositories_on_snippet_project_id.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class IndexSnippetRepositoriesOnSnippetProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_snippet_repositories_on_snippet_project_id'
+
+  def up
+    add_concurrent_index :snippet_repositories, :snippet_project_id, name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :snippet_repositories, INDEX_NAME
+  end
+end
diff --git a/db/post_migrate/20241211134708_add_snippet_repositories_snippet_project_id_fk.rb b/db/post_migrate/20241211134708_add_snippet_repositories_snippet_project_id_fk.rb
new file mode 100644
index 000000000000..50428f826383
--- /dev/null
+++ b/db/post_migrate/20241211134708_add_snippet_repositories_snippet_project_id_fk.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class AddSnippetRepositoriesSnippetProjectIdFk < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_foreign_key :snippet_repositories, :projects, column: :snippet_project_id, on_delete: :cascade
+  end
+
+  def down
+    with_lock_retries do
+      remove_foreign_key :snippet_repositories, column: :snippet_project_id
+    end
+  end
+end
diff --git a/db/post_migrate/20241211134709_add_snippet_repositories_snippet_project_id_trigger.rb b/db/post_migrate/20241211134709_add_snippet_repositories_snippet_project_id_trigger.rb
new file mode 100644
index 000000000000..26ae122606fb
--- /dev/null
+++ b/db/post_migrate/20241211134709_add_snippet_repositories_snippet_project_id_trigger.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+class AddSnippetRepositoriesSnippetProjectIdTrigger < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+
+  def up
+    install_sharding_key_assignment_trigger(
+      table: :snippet_repositories,
+      sharding_key: :snippet_project_id,
+      parent_table: :snippets,
+      parent_sharding_key: :project_id,
+      foreign_key: :snippet_id
+    )
+  end
+
+  def down
+    remove_sharding_key_assignment_trigger(
+      table: :snippet_repositories,
+      sharding_key: :snippet_project_id,
+      parent_table: :snippets,
+      parent_sharding_key: :project_id,
+      foreign_key: :snippet_id
+    )
+  end
+end
diff --git a/db/post_migrate/20241211134710_queue_backfill_snippet_repositories_snippet_project_id.rb b/db/post_migrate/20241211134710_queue_backfill_snippet_repositories_snippet_project_id.rb
new file mode 100644
index 000000000000..2b16f78a37f7
--- /dev/null
+++ b/db/post_migrate/20241211134710_queue_backfill_snippet_repositories_snippet_project_id.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+class QueueBackfillSnippetRepositoriesSnippetProjectId < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  restrict_gitlab_migration gitlab_schema: :gitlab_main_cell
+
+  MIGRATION = "BackfillSnippetRepositoriesSnippetProjectId"
+  DELAY_INTERVAL = 2.minutes
+  BATCH_SIZE = 1000
+  SUB_BATCH_SIZE = 100
+
+  def up
+    queue_batched_background_migration(
+      MIGRATION,
+      :snippet_repositories,
+      :snippet_id,
+      :snippet_project_id,
+      :snippets,
+      :project_id,
+      :snippet_id,
+      job_interval: DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      sub_batch_size: SUB_BATCH_SIZE
+    )
+  end
+
+  def down
+    delete_batched_background_migration(
+      MIGRATION,
+      :snippet_repositories,
+      :snippet_id,
+      [
+        :snippet_project_id,
+        :snippets,
+        :project_id,
+        :snippet_id
+      ]
+    )
+  end
+end
diff --git a/db/post_migrate/20241211134712_index_snippet_repositories_on_snippet_organization_id.rb b/db/post_migrate/20241211134712_index_snippet_repositories_on_snippet_organization_id.rb
new file mode 100644
index 000000000000..b6023a7272bb
--- /dev/null
+++ b/db/post_migrate/20241211134712_index_snippet_repositories_on_snippet_organization_id.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class IndexSnippetRepositoriesOnSnippetOrganizationId < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_snippet_repositories_on_snippet_organization_id'
+
+  def up
+    add_concurrent_index :snippet_repositories, :snippet_organization_id, name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :snippet_repositories, INDEX_NAME
+  end
+end
diff --git a/db/post_migrate/20241211134713_add_snippet_repositories_snippet_organization_id_fk.rb b/db/post_migrate/20241211134713_add_snippet_repositories_snippet_organization_id_fk.rb
new file mode 100644
index 000000000000..8c821c32bd33
--- /dev/null
+++ b/db/post_migrate/20241211134713_add_snippet_repositories_snippet_organization_id_fk.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class AddSnippetRepositoriesSnippetOrganizationIdFk < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_foreign_key :snippet_repositories, :organizations, column: :snippet_organization_id,
+      on_delete: :cascade
+  end
+
+  def down
+    with_lock_retries do
+      remove_foreign_key :snippet_repositories, column: :snippet_organization_id
+    end
+  end
+end
diff --git a/db/post_migrate/20241211134714_add_snippet_repositories_snippet_organization_id_trigger.rb b/db/post_migrate/20241211134714_add_snippet_repositories_snippet_organization_id_trigger.rb
new file mode 100644
index 000000000000..7a9cd4d0f561
--- /dev/null
+++ b/db/post_migrate/20241211134714_add_snippet_repositories_snippet_organization_id_trigger.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+class AddSnippetRepositoriesSnippetOrganizationIdTrigger < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+
+  def up
+    install_sharding_key_assignment_trigger(
+      table: :snippet_repositories,
+      sharding_key: :snippet_organization_id,
+      parent_table: :snippets,
+      parent_sharding_key: :organization_id,
+      foreign_key: :snippet_id
+    )
+  end
+
+  def down
+    remove_sharding_key_assignment_trigger(
+      table: :snippet_repositories,
+      sharding_key: :snippet_organization_id,
+      parent_table: :snippets,
+      parent_sharding_key: :organization_id,
+      foreign_key: :snippet_id
+    )
+  end
+end
diff --git a/db/post_migrate/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id.rb b/db/post_migrate/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id.rb
new file mode 100644
index 000000000000..cdba4aa6b35e
--- /dev/null
+++ b/db/post_migrate/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id.rb
@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+class QueueBackfillSnippetRepositoriesSnippetOrganizationId < Gitlab::Database::Migration[2.2]
+  milestone '17.10'
+  restrict_gitlab_migration gitlab_schema: :gitlab_main_cell
+
+  MIGRATION = "BackfillSnippetRepositoriesSnippetOrganizationId"
+  DELAY_INTERVAL = 2.minutes
+  BATCH_SIZE = 1000
+  SUB_BATCH_SIZE = 100
+
+  def up
+    queue_batched_background_migration(
+      MIGRATION,
+      :snippet_repositories,
+      :snippet_id,
+      :snippet_organization_id,
+      :snippets,
+      :organization_id,
+      :snippet_id,
+      job_interval: DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      sub_batch_size: SUB_BATCH_SIZE
+    )
+  end
+
+  def down
+    delete_batched_background_migration(
+      MIGRATION,
+      :snippet_repositories,
+      :snippet_id,
+      [
+        :snippet_organization_id,
+        :snippets,
+        :organization_id,
+        :snippet_id
+      ]
+    )
+  end
+end
diff --git a/db/schema_migrations/20241211134706 b/db/schema_migrations/20241211134706
new file mode 100644
index 000000000000..dbd9892d8112
--- /dev/null
+++ b/db/schema_migrations/20241211134706
@@ -0,0 +1 @@
+48f5f852380a7a3e179c913568219ae8799b32c9b59c0f3219be6dc31a45332b
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134707 b/db/schema_migrations/20241211134707
new file mode 100644
index 000000000000..3666178c5dcc
--- /dev/null
+++ b/db/schema_migrations/20241211134707
@@ -0,0 +1 @@
+4a3a2de5da48f40751d6b7bacb5a4081954eda13a31f535e0e65c17c821c2e3a
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134708 b/db/schema_migrations/20241211134708
new file mode 100644
index 000000000000..3f17f0a8f49b
--- /dev/null
+++ b/db/schema_migrations/20241211134708
@@ -0,0 +1 @@
+d3b3fd1a860be5e1fcf4b1ed8a645d350c9527601db2bf9d0cc5cf8872382026
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134709 b/db/schema_migrations/20241211134709
new file mode 100644
index 000000000000..ed08f8c43a59
--- /dev/null
+++ b/db/schema_migrations/20241211134709
@@ -0,0 +1 @@
+035aea2280dca568860cf636727798b9f2a0d98669d8302d12ed8183f07035c7
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134710 b/db/schema_migrations/20241211134710
new file mode 100644
index 000000000000..c87cf892cf93
--- /dev/null
+++ b/db/schema_migrations/20241211134710
@@ -0,0 +1 @@
+38438dced134d7c7bf46dc0bb955f566c4edaebe6bd47b9a4b7d39027c2f0e98
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134711 b/db/schema_migrations/20241211134711
new file mode 100644
index 000000000000..8c107f5215d5
--- /dev/null
+++ b/db/schema_migrations/20241211134711
@@ -0,0 +1 @@
+7bcb615592724e6a9b499293e5e4eb48daab5a867499a707735fbe336c46b78e
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134712 b/db/schema_migrations/20241211134712
new file mode 100644
index 000000000000..a6ccbf74662c
--- /dev/null
+++ b/db/schema_migrations/20241211134712
@@ -0,0 +1 @@
+e53745adeacf8ec8b3ca1f343927942093bdfbccb76d56ac90b4728f967bf7c8
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134713 b/db/schema_migrations/20241211134713
new file mode 100644
index 000000000000..18e4111c5bc5
--- /dev/null
+++ b/db/schema_migrations/20241211134713
@@ -0,0 +1 @@
+22d6304d4124f9353db4c409c7a30c7b8bc4536233cb7f388706111f0ca48515
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134714 b/db/schema_migrations/20241211134714
new file mode 100644
index 000000000000..98a4d47780ab
--- /dev/null
+++ b/db/schema_migrations/20241211134714
@@ -0,0 +1 @@
+841b185889673d3293c949d91a53baf5501109846d0b2d625d50c2582fe610f2
\ No newline at end of file
diff --git a/db/schema_migrations/20241211134715 b/db/schema_migrations/20241211134715
new file mode 100644
index 000000000000..e92a3224a667
--- /dev/null
+++ b/db/schema_migrations/20241211134715
@@ -0,0 +1 @@
+4af7bbac18c391a5168bc68fbc53e046884f848c25d8c217aae397a37e52fe80
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index e320d2edfdfe..dbe516d873fa 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -1396,6 +1396,22 @@ RETURN NEW;
 END
 $$;
 
+CREATE FUNCTION trigger_1f57c71a69fb() 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_206cbe2dc1a2() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -3539,6 +3555,22 @@ RETURN NEW;
 END
 $$;
 
+CREATE FUNCTION trigger_d9468bfbb0b4() 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_da5fd3d6d75c() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
@@ -21957,6 +21989,8 @@ CREATE TABLE snippet_repositories (
     verification_failure text,
     verification_state smallint DEFAULT 0 NOT NULL,
     verification_started_at timestamp with time zone,
+    snippet_project_id bigint,
+    snippet_organization_id bigint,
     CONSTRAINT snippet_repositories_verification_failure_text_limit CHECK ((char_length(verification_failure) <= 255))
 );
 
@@ -35586,6 +35620,10 @@ CREATE UNIQUE INDEX index_snippet_repositories_on_disk_path ON snippet_repositor
 
 CREATE INDEX index_snippet_repositories_on_shard_id ON snippet_repositories USING btree (shard_id);
 
+CREATE INDEX index_snippet_repositories_on_snippet_organization_id ON snippet_repositories USING btree (snippet_organization_id);
+
+CREATE INDEX index_snippet_repositories_on_snippet_project_id ON snippet_repositories USING btree (snippet_project_id);
+
 CREATE INDEX index_snippet_repositories_pending_verification ON snippet_repositories USING btree (verified_at NULLS FIRST) WHERE (verification_state = 0);
 
 CREATE INDEX index_snippet_repositories_verification_state ON snippet_repositories USING btree (verification_state);
@@ -38882,6 +38920,8 @@ CREATE TRIGGER trigger_1ed40f4d5f4e BEFORE INSERT OR UPDATE ON packages_maven_me
 
 CREATE TRIGGER trigger_1eda1bc6ef53 BEFORE INSERT OR UPDATE ON merge_request_diff_details FOR EACH ROW EXECUTE FUNCTION trigger_1eda1bc6ef53();
 
+CREATE TRIGGER trigger_1f57c71a69fb BEFORE INSERT OR UPDATE ON snippet_repositories FOR EACH ROW EXECUTE FUNCTION trigger_1f57c71a69fb();
+
 CREATE TRIGGER trigger_206cbe2dc1a2 BEFORE INSERT OR UPDATE ON packages_package_files FOR EACH ROW EXECUTE FUNCTION trigger_206cbe2dc1a2();
 
 CREATE TRIGGER trigger_207005e8e995 BEFORE INSERT OR UPDATE ON operations_strategies FOR EACH ROW EXECUTE FUNCTION trigger_207005e8e995();
@@ -39158,6 +39198,8 @@ CREATE TRIGGER trigger_d5c895007948 BEFORE INSERT OR UPDATE ON protected_environ
 
 CREATE TRIGGER trigger_d8c2de748d8c BEFORE INSERT OR UPDATE ON merge_request_predictions FOR EACH ROW EXECUTE FUNCTION trigger_d8c2de748d8c();
 
+CREATE TRIGGER trigger_d9468bfbb0b4 BEFORE INSERT OR UPDATE ON snippet_repositories FOR EACH ROW EXECUTE FUNCTION trigger_d9468bfbb0b4();
+
 CREATE TRIGGER trigger_da5fd3d6d75c BEFORE INSERT OR UPDATE ON packages_composer_metadata FOR EACH ROW EXECUTE FUNCTION trigger_da5fd3d6d75c();
 
 CREATE TRIGGER trigger_dadd660afe2c BEFORE INSERT OR UPDATE ON packages_debian_group_distribution_keys FOR EACH ROW EXECUTE FUNCTION trigger_dadd660afe2c();
@@ -41192,6 +41234,9 @@ ALTER TABLE ONLY merge_request_context_commits
 ALTER TABLE ONLY approval_project_rules
     ADD CONSTRAINT fk_efa5a1e3fb FOREIGN KEY (security_orchestration_policy_configuration_id) REFERENCES security_orchestration_policy_configurations(id) ON DELETE CASCADE;
 
+ALTER TABLE ONLY snippet_repositories
+    ADD CONSTRAINT fk_efaf4ac269 FOREIGN KEY (snippet_organization_id) REFERENCES organizations(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY dora_daily_metrics
     ADD CONSTRAINT fk_efc32a39fa FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE;
 
@@ -41228,6 +41273,9 @@ ALTER TABLE ONLY abuse_reports
 ALTER TABLE ONLY timelogs
     ADD CONSTRAINT fk_f12ef8db70 FOREIGN KEY (timelog_category_id) REFERENCES timelog_categories(id) ON DELETE SET NULL;
 
+ALTER TABLE ONLY snippet_repositories
+    ADD CONSTRAINT fk_f1319bee9d FOREIGN KEY (snippet_project_id) REFERENCES projects(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY boards
     ADD CONSTRAINT fk_f15266b5f9 FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE;
 
diff --git a/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id.rb b/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id.rb
new file mode 100644
index 000000000000..a9fa1321d232
--- /dev/null
+++ b/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillSnippetRepositoriesSnippetOrganizationId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_snippet_repositories_snippet_organization_id
+      feature_category :source_code_management
+    end
+  end
+end
diff --git a/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id.rb b/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id.rb
new file mode 100644
index 000000000000..fba695a972cb
--- /dev/null
+++ b/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    class BackfillSnippetRepositoriesSnippetProjectId < BackfillDesiredShardingKeyJob
+      operation_name :backfill_snippet_repositories_snippet_project_id
+      feature_category :source_code_management
+    end
+  end
+end
diff --git a/spec/features/snippets/user_creates_snippet_spec.rb b/spec/features/snippets/user_creates_snippet_spec.rb
index 23659beb8215..989072ba1fb5 100644
--- a/spec/features/snippets/user_creates_snippet_spec.rb
+++ b/spec/features/snippets/user_creates_snippet_spec.rb
@@ -16,6 +16,8 @@
   let(:snippet_title_field) { 'snippet-title' }
 
   before do
+    create(:organization, :default)
+
     sign_in(user)
 
     visit new_snippet_path
diff --git a/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id_spec.rb b/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id_spec.rb
new file mode 100644
index 000000000000..7c14ed1c0a9d
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_organization_id_spec.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillSnippetRepositoriesSnippetOrganizationId,
+  feature_category: :source_code_management,
+  schema: 20241211134711 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :snippet_repositories }
+    let(:backfill_column) { :snippet_organization_id }
+    let(:batch_column) { :snippet_id }
+    let(:backfill_via_table) { :snippets }
+    let(:backfill_via_column) { :organization_id }
+    let(:backfill_via_foreign_key) { :snippet_id }
+  end
+end
diff --git a/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id_spec.rb b/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id_spec.rb
new file mode 100644
index 000000000000..f592656e0390
--- /dev/null
+++ b/spec/lib/gitlab/background_migration/backfill_snippet_repositories_snippet_project_id_spec.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillSnippetRepositoriesSnippetProjectId,
+  feature_category: :source_code_management,
+  schema: 20241211134706 do
+  include_examples 'desired sharding key backfill job' do
+    let(:batch_table) { :snippet_repositories }
+    let(:backfill_column) { :snippet_project_id }
+    let(:batch_column) { :snippet_id }
+    let(:backfill_via_table) { :snippets }
+    let(:backfill_via_column) { :project_id }
+    let(:backfill_via_foreign_key) { :snippet_id }
+  end
+end
diff --git a/spec/migrations/20241211134710_queue_backfill_snippet_repositories_snippet_project_id_spec.rb b/spec/migrations/20241211134710_queue_backfill_snippet_repositories_snippet_project_id_spec.rb
new file mode 100644
index 000000000000..42dcb11edae1
--- /dev/null
+++ b/spec/migrations/20241211134710_queue_backfill_snippet_repositories_snippet_project_id_spec.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillSnippetRepositoriesSnippetProjectId, 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_repositories,
+          column_name: :snippet_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
diff --git a/spec/migrations/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id_spec.rb b/spec/migrations/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id_spec.rb
new file mode 100644
index 000000000000..0f49a34a42f4
--- /dev/null
+++ b/spec/migrations/20241211134715_queue_backfill_snippet_repositories_snippet_organization_id_spec.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe QueueBackfillSnippetRepositoriesSnippetOrganizationId, 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_repositories,
+          column_name: :snippet_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
diff --git a/spec/requests/api/graphql/mutations/snippets/create_spec.rb b/spec/requests/api/graphql/mutations/snippets/create_spec.rb
index 9f9b312f12e3..a4b65e658a18 100644
--- a/spec/requests/api/graphql/mutations/snippets/create_spec.rb
+++ b/spec/requests/api/graphql/mutations/snippets/create_spec.rb
@@ -37,6 +37,10 @@ def mutation_response
     graphql_mutation_response(:create_snippet)
   end
 
+  before do
+    create(:organization, :default)
+  end
+
   subject { post_graphql_mutation(mutation, current_user: current_user) }
 
   context 'when the user does not have permission' do
diff --git a/spec/services/snippets/create_service_spec.rb b/spec/services/snippets/create_service_spec.rb
index d8fe2ac31c0a..beb07e379468 100644
--- a/spec/services/snippets/create_service_spec.rb
+++ b/spec/services/snippets/create_service_spec.rb
@@ -25,6 +25,10 @@
 
     let(:snippet) { subject.payload[:snippet] }
 
+    before do
+      create(:organization, :default)
+    end
+
     shared_examples 'a service that creates a snippet' do
       it 'creates a snippet with the provided attributes' do
         expect(snippet.title).to eq(opts[:title])
-- 
GitLab