From e4a250a50637406ac515a41273791360794e9410 Mon Sep 17 00:00:00 2001
From: Krasimir Angelov <kangelov@gitlab.com>
Date: Thu, 3 Oct 2024 10:26:33 +1300
Subject: [PATCH] Add sharding key columns to uploads

Add new sharding key columns to `uploads` and populate them
on create/update.

For model that is using `uploads` implement
`#uploads_sharding_key` method.

Not all models using `uploads` have sharding key yet,
so the new column may not be populated for them.

https://gitlab.com/gitlab-org/gitlab/-/issues/398199

Changelog: added
---
 app/models/abuse_report.rb                    |  4 +++
 app/models/achievements/achievement.rb        |  4 +++
 app/models/alert_management/metric_image.rb   |  4 +++
 app/models/appearance.rb                      |  4 +++
 app/models/bulk_imports/export_upload.rb      |  7 +++++
 app/models/design_management/action.rb        |  4 +++
 app/models/import_export_upload.rb            |  7 +++++
 app/models/namespace.rb                       |  4 +++
 app/models/note.rb                            |  4 +++
 .../organizations/organization_detail.rb      |  4 +++
 app/models/personal_snippet.rb                |  4 +++
 app/models/project.rb                         |  4 +++
 .../import_export/relation_export_upload.rb   |  4 +++
 app/models/projects/topic.rb                  |  4 +++
 app/models/upload.rb                          | 12 ++++++++
 app/models/user.rb                            |  4 +++
 ...27015108_add_sharding_key_id_to_uploads.rb | 17 +++++++++++
 db/schema_migrations/20241127015108           |  1 +
 db/structure.sql                              |  3 ++
 ee/app/models/ai/vectorizable_file.rb         |  4 +++
 .../dependencies/dependency_list_export.rb    |  8 ++++++
 .../dependency_list_export/part.rb            |  4 +++
 ee/app/models/issuable_metric_image.rb        |  4 +++
 .../models/user_permission_export_upload.rb   |  4 +++
 ee/app/models/vulnerabilities/export.rb       |  4 +++
 ee/app/models/vulnerabilities/export/part.rb  |  4 +++
 ee/app/models/vulnerabilities/remediation.rb  |  4 +++
 ee/spec/models/ai/vectorizable_file_spec.rb   |  9 ++++++
 .../dependency_list_export/part_spec.rb       |  9 ++++++
 .../dependency_list_export_spec.rb            | 21 ++++++++++++++
 ee/spec/models/issuable_metric_image_spec.rb  | 10 +++++++
 .../user_permission_export_upload_spec.rb     |  6 ++++
 .../vulnerabilities/export/part_spec.rb       |  9 ++++++
 ee/spec/models/vulnerabilities/export_spec.rb |  9 ++++++
 .../vulnerabilities/remediation_spec.rb       |  9 ++++++
 spec/db/schema_spec.rb                        |  2 +-
 .../alert_management/metric_images.rb         |  1 +
 spec/lib/gitlab/database/sharding_key_spec.rb |  3 +-
 spec/models/abuse_report_spec.rb              |  6 ++++
 spec/models/achievements/achievement_spec.rb  |  9 ++++++
 .../alert_management/metric_image_spec.rb     |  9 ++++++
 spec/models/appearance_spec.rb                |  8 ++++++
 .../models/bulk_imports/export_upload_spec.rb | 13 ++++++++-
 spec/models/design_management/action_spec.rb  | 10 +++++++
 spec/models/import_export_upload_spec.rb      | 11 +++++++-
 spec/models/namespace_spec.rb                 |  9 ++++++
 spec/models/note_spec.rb                      |  9 ++++++
 .../organizations/organization_detail_spec.rb |  8 ++++++
 spec/models/personal_snippet_spec.rb          |  9 ++++++
 spec/models/project_spec.rb                   |  9 ++++++
 .../relation_export_upload_spec.rb            | 28 ++++++++++++++++---
 spec/models/projects/topic_spec.rb            |  9 ++++++
 spec/models/upload_spec.rb                    | 21 ++++++++++++++
 spec/models/user_spec.rb                      |  8 ++++++
 spec/uploaders/object_storage_spec.rb         |  6 ++++
 55 files changed, 397 insertions(+), 8 deletions(-)
 create mode 100644 db/migrate/20241127015108_add_sharding_key_id_to_uploads.rb
 create mode 100644 db/schema_migrations/20241127015108

diff --git a/app/models/abuse_report.rb b/app/models/abuse_report.rb
index c37e94e41c043b..96c61ed9cc464f 100644
--- a/app/models/abuse_report.rb
+++ b/app/models/abuse_report.rb
@@ -173,6 +173,10 @@ def project
     nil
   end
 
+  def uploads_sharding_key
+    {}
+  end
+
   private
 
   def reported_project
diff --git a/app/models/achievements/achievement.rb b/app/models/achievements/achievement.rb
index 834c12fee5a474..d47868c4b669d9 100644
--- a/app/models/achievements/achievement.rb
+++ b/app/models/achievements/achievement.rb
@@ -17,5 +17,9 @@ class Achievement < ApplicationRecord
       length: { maximum: 255 },
       uniqueness: { case_sensitive: false, scope: [:namespace_id] }
     validates :description, length: { maximum: 1024 }
+
+    def uploads_sharding_key
+      { namespace_id: namespace_id }
+    end
   end
 end
diff --git a/app/models/alert_management/metric_image.rb b/app/models/alert_management/metric_image.rb
index 4ed28c3b1ebdd4..5249e15a3ee15e 100644
--- a/app/models/alert_management/metric_image.rb
+++ b/app/models/alert_management/metric_image.rb
@@ -7,6 +7,10 @@ class MetricImage < ApplicationRecord
 
     belongs_to :alert, class_name: 'AlertManagement::Alert', foreign_key: 'alert_id', inverse_of: :metric_images
 
+    def uploads_sharding_key
+      { project_id: project_id }
+    end
+
     private
 
     def local_path
diff --git a/app/models/appearance.rb b/app/models/appearance.rb
index 2352f6f736c0fa..9a19824d62527d 100644
--- a/app/models/appearance.rb
+++ b/app/models/appearance.rb
@@ -102,6 +102,10 @@ def show_footer?
     footer_message.present?
   end
 
+  def uploads_sharding_key
+    {}
+  end
+
   private
 
   def logo_system_path(logo, mount_type)
diff --git a/app/models/bulk_imports/export_upload.rb b/app/models/bulk_imports/export_upload.rb
index 0560933ed93583..ee837b3c36bd7a 100644
--- a/app/models/bulk_imports/export_upload.rb
+++ b/app/models/bulk_imports/export_upload.rb
@@ -22,5 +22,12 @@ class ExportUpload < ApplicationRecord
     def retrieve_upload(_identifier, paths)
       Upload.find_by(model: self, path: paths)
     end
+
+    def uploads_sharding_key
+      {
+        project_id: export&.project_id,
+        namespace_id: export&.group_id
+      }
+    end
   end
 end
diff --git a/app/models/design_management/action.rb b/app/models/design_management/action.rb
index 48e0c93489b5c8..ea416868e45674 100644
--- a/app/models/design_management/action.rb
+++ b/app/models/design_management/action.rb
@@ -43,5 +43,9 @@ class Action < ApplicationRecord
         raise ArgumentError, "Expected a DesignManagement::Version, got #{version}"
       end
     end
+
+    def uploads_sharding_key
+      { namespace_id: design&.namespace_id }
+    end
   end
 end
diff --git a/app/models/import_export_upload.rb b/app/models/import_export_upload.rb
index fc167e1c3ab931..ce4c1e1b915201 100644
--- a/app/models/import_export_upload.rb
+++ b/app/models/import_export_upload.rb
@@ -44,6 +44,13 @@ def export_archive_exists?
     false
   end
 
+  def uploads_sharding_key
+    {
+      project_id: project_id,
+      namespace_id: group_id
+    }
+  end
+
   private
 
   def carrierwave_export_file
diff --git a/app/models/namespace.rb b/app/models/namespace.rb
index 84264907dee486..f820bef23a668f 100644
--- a/app/models/namespace.rb
+++ b/app/models/namespace.rb
@@ -755,6 +755,10 @@ def deleted?
     !!deleted_at
   end
 
+  def uploads_sharding_key
+    { organization_id: organization_id }
+  end
+
   private
 
   def require_organization?
diff --git a/app/models/note.rb b/app/models/note.rb
index 9b7fe1a4da1855..35c33828de7e08 100644
--- a/app/models/note.rb
+++ b/app/models/note.rb
@@ -723,6 +723,10 @@ def attribute_names_for_serialization
     attributes.keys
   end
 
+  def uploads_sharding_key
+    { namespace_id: namespace_id }
+  end
+
   private
 
   def trigger_note_subscription?
diff --git a/app/models/organizations/organization_detail.rb b/app/models/organizations/organization_detail.rb
index 018e7579c5b8ae..fa077981a0eb30 100644
--- a/app/models/organizations/organization_detail.rb
+++ b/app/models/organizations/organization_detail.rb
@@ -12,5 +12,9 @@ class OrganizationDetail < ApplicationRecord
 
     validates :organization, presence: true
     validates :description, length: { maximum: 1024 }
+
+    def uploads_sharding_key
+      { organization_id: organization_id }
+    end
   end
 end
diff --git a/app/models/personal_snippet.rb b/app/models/personal_snippet.rb
index 0a51bfd55f873b..f12c13fa44f21c 100644
--- a/app/models/personal_snippet.rb
+++ b/app/models/personal_snippet.rb
@@ -14,4 +14,8 @@ def parent_user
   def skip_project_check?
     true
   end
+
+  def uploads_sharding_key
+    { organization_id: organization_id }
+  end
 end
diff --git a/app/models/project.rb b/app/models/project.rb
index 5c2799090d130a..3259cf18a90ffe 100644
--- a/app/models/project.rb
+++ b/app/models/project.rb
@@ -3437,6 +3437,10 @@ def placeholder_reference_store
     )
   end
 
+  def uploads_sharding_key
+    { namespace_id: namespace_id }
+  end
+
   def pages_url
     Gitlab::Pages::UrlBuilder.new(self).pages_url
   end
diff --git a/app/models/projects/import_export/relation_export_upload.rb b/app/models/projects/import_export/relation_export_upload.rb
index 89ff91df563df1..72f2bc1b676812 100644
--- a/app/models/projects/import_export/relation_export_upload.rb
+++ b/app/models/projects/import_export/relation_export_upload.rb
@@ -27,6 +27,10 @@ class RelationExportUpload < ApplicationRecord
         skip_callback :save, :after, :store_export_file!
         set_callback :commit, :after, :store_export_file!
       end
+
+      def uploads_sharding_key
+        { project_id: relation_export&.project_id }
+      end
     end
   end
 end
diff --git a/app/models/projects/topic.rb b/app/models/projects/topic.rb
index e0df6b3b698976..05d5d209b4651d 100644
--- a/app/models/projects/topic.rb
+++ b/app/models/projects/topic.rb
@@ -76,6 +76,10 @@ def update_non_private_projects_counter(ids_before, ids_after, project_visibilit
       end
     end
 
+    def uploads_sharding_key
+      { organization_id: organization_id }
+    end
+
     private
 
     def validate_name_format
diff --git a/app/models/upload.rb b/app/models/upload.rb
index 06fa3b9360fe92..09cd1e59cb25d4 100644
--- a/app/models/upload.rb
+++ b/app/models/upload.rb
@@ -23,6 +23,7 @@ class Upload < ApplicationRecord
   scope :preload_uploaded_by_user, -> { preload(:uploaded_by_user) }
 
   before_save :calculate_checksum!, if: :foreground_checksummable?
+  before_save :ensure_sharding_key
   # as the FileUploader is not mounted, the default CarrierWave ActiveRecord
   # hooks are not executed and the file will not be deleted
   after_destroy :delete_file!, if: -> { uploader_class <= FileUploader }
@@ -182,6 +183,17 @@ def project?
   def update_project_statistics
     ProjectCacheWorker.perform_async(model_id, [], %w[uploads_size])
   end
+
+  def ensure_sharding_key
+    sharding_key = model&.uploads_sharding_key
+    return unless sharding_key.present?
+
+    # This is workaround for some migrations that rely on application code to use
+    # bot users, and creating these fail in tests if the column is not present yet.
+    return unless sharding_key.each_key.all? { |k| respond_to?(k) }
+
+    assign_attributes(sharding_key)
+  end
 end
 
 Upload.prepend_mod_with('Upload')
diff --git a/app/models/user.rb b/app/models/user.rb
index 67cac0bff4f8f4..d50085ea048534 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -2530,6 +2530,10 @@ def has_composite_identity?
     false
   end
 
+  def uploads_sharding_key
+    {}
+  end
+
   protected
 
   # override, from Devise::Validatable
diff --git a/db/migrate/20241127015108_add_sharding_key_id_to_uploads.rb b/db/migrate/20241127015108_add_sharding_key_id_to_uploads.rb
new file mode 100644
index 00000000000000..18c581b5f568e0
--- /dev/null
+++ b/db/migrate/20241127015108_add_sharding_key_id_to_uploads.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+class AddShardingKeyIdToUploads < Gitlab::Database::Migration[2.2]
+  milestone '17.7'
+
+  def up
+    add_column :uploads, :organization_id, :bigint
+    add_column :uploads, :namespace_id, :bigint
+    add_column :uploads, :project_id, :bigint
+  end
+
+  def down
+    remove_column :uploads, :project_id
+    remove_column :uploads, :namespace_id
+    remove_column :uploads, :organization_id
+  end
+end
diff --git a/db/schema_migrations/20241127015108 b/db/schema_migrations/20241127015108
new file mode 100644
index 00000000000000..e62da5f5a6bdd9
--- /dev/null
+++ b/db/schema_migrations/20241127015108
@@ -0,0 +1 @@
+9414c5f6c3d687905fac7ff840361dec14fed7eb165a5c45ec5c554fa6c7d3fc
\ No newline at end of file
diff --git a/db/structure.sql b/db/structure.sql
index f9d894d4389399..2860050ceef4e8 100644
--- a/db/structure.sql
+++ b/db/structure.sql
@@ -20649,6 +20649,9 @@ CREATE TABLE uploads (
     secret character varying,
     version integer DEFAULT 1 NOT NULL,
     uploaded_by_user_id bigint,
+    organization_id bigint,
+    namespace_id bigint,
+    project_id bigint,
     CONSTRAINT check_5e9547379c CHECK ((store IS NOT NULL))
 );
 
diff --git a/ee/app/models/ai/vectorizable_file.rb b/ee/app/models/ai/vectorizable_file.rb
index 4abeb068a7597f..fe03a0fd8ae3cb 100644
--- a/ee/app/models/ai/vectorizable_file.rb
+++ b/ee/app/models/ai/vectorizable_file.rb
@@ -20,5 +20,9 @@ class VectorizableFile < ApplicationRecord
     has_many :attachments, class_name: 'Ai::AgentVersionAttachment', foreign_key: :ai_vectorizable_file_id,
       inverse_of: :file
     has_many :versions, through: :attachments, source: :version
+
+    def uploads_sharding_key
+      { project_id: project_id }
+    end
   end
 end
diff --git a/ee/app/models/dependencies/dependency_list_export.rb b/ee/app/models/dependencies/dependency_list_export.rb
index 211043b67805d3..ceaa0d7d1dd7bd 100644
--- a/ee/app/models/dependencies/dependency_list_export.rb
+++ b/ee/app/models/dependencies/dependency_list_export.rb
@@ -87,6 +87,14 @@ def timed_out?
       created_at < MAX_EXPORT_DURATION.ago
     end
 
+    def uploads_sharding_key
+      {
+        organization_id: organization_id,
+        namespace_id: group_id,
+        project_id: project_id
+      }
+    end
+
     private
 
     def only_one_exportable
diff --git a/ee/app/models/dependencies/dependency_list_export/part.rb b/ee/app/models/dependencies/dependency_list_export/part.rb
index 538ee3fa7825cf..c0143ad57f4ca2 100644
--- a/ee/app/models/dependencies/dependency_list_export/part.rb
+++ b/ee/app/models/dependencies/dependency_list_export/part.rb
@@ -32,6 +32,10 @@ def sbom_occurrences
       end
 
       delegate :exportable, to: :dependency_list_export, private: true
+
+      def uploads_sharding_key
+        { organization_id: organization_id }
+      end
     end
   end
 end
diff --git a/ee/app/models/issuable_metric_image.rb b/ee/app/models/issuable_metric_image.rb
index ac380def4fcb26..4299477460ab44 100644
--- a/ee/app/models/issuable_metric_image.rb
+++ b/ee/app/models/issuable_metric_image.rb
@@ -9,6 +9,10 @@ def self.available_for?(project)
     project&.feature_available?(:incident_metric_upload)
   end
 
+  def uploads_sharding_key
+    { namespace_id: issue&.namespace_id }
+  end
+
   private
 
   def local_path
diff --git a/ee/app/models/user_permission_export_upload.rb b/ee/app/models/user_permission_export_upload.rb
index d03c363355130d..b82a8194e1bbd2 100644
--- a/ee/app/models/user_permission_export_upload.rb
+++ b/ee/app/models/user_permission_export_upload.rb
@@ -30,6 +30,10 @@ class UserPermissionExportUpload < ApplicationRecord
     state :failed, value: 3
   end
 
+  def uploads_sharding_key
+    {}
+  end
+
   private
 
   def file_presence
diff --git a/ee/app/models/vulnerabilities/export.rb b/ee/app/models/vulnerabilities/export.rb
index f4f8453cc10f7b..3f5733eea58474 100644
--- a/ee/app/models/vulnerabilities/export.rb
+++ b/ee/app/models/vulnerabilities/export.rb
@@ -100,6 +100,10 @@ def timed_out?
       created_at < MAX_EXPORT_DURATION.ago
     end
 
+    def uploads_sharding_key
+      { organization_id: organization_id }
+    end
+
     private
 
     def make_project_level_export(project)
diff --git a/ee/app/models/vulnerabilities/export/part.rb b/ee/app/models/vulnerabilities/export/part.rb
index 1d02d1855ea93d..24c2500ebf4258 100644
--- a/ee/app/models/vulnerabilities/export/part.rb
+++ b/ee/app/models/vulnerabilities/export/part.rb
@@ -21,6 +21,10 @@ class Part < Gitlab::Database::SecApplicationRecord
       def retrieve_upload(_identifier, paths)
         Upload.find_by(model: self, path: paths)
       end
+
+      def uploads_sharding_key
+        { organization_id: organization_id }
+      end
     end
   end
 end
diff --git a/ee/app/models/vulnerabilities/remediation.rb b/ee/app/models/vulnerabilities/remediation.rb
index f7e97f7fb9c37b..6710ea55836805 100644
--- a/ee/app/models/vulnerabilities/remediation.rb
+++ b/ee/app/models/vulnerabilities/remediation.rb
@@ -30,5 +30,9 @@ def diff
     def retrieve_upload(_identifier, paths)
       Upload.find_by(model: self, path: paths)
     end
+
+    def uploads_sharding_key
+      { project_id: project_id }
+    end
   end
 end
diff --git a/ee/spec/models/ai/vectorizable_file_spec.rb b/ee/spec/models/ai/vectorizable_file_spec.rb
index 6b0226af6e4eb4..c8083c45f24e8a 100644
--- a/ee/spec/models/ai/vectorizable_file_spec.rb
+++ b/ee/spec/models/ai/vectorizable_file_spec.rb
@@ -31,4 +31,13 @@
       expect(file.errors.full_messages).to include("File is too big (should be at most 100 MiB)")
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id' do
+      project = build_stubbed(:project)
+      file = build_stubbed(:ai_vectorizable_file, project: project)
+
+      expect(file.uploads_sharding_key).to eq(project_id: project.id)
+    end
+  end
 end
diff --git a/ee/spec/models/dependencies/dependency_list_export/part_spec.rb b/ee/spec/models/dependencies/dependency_list_export/part_spec.rb
index 8be556dd048a06..26be54d51f9a37 100644
--- a/ee/spec/models/dependencies/dependency_list_export/part_spec.rb
+++ b/ee/spec/models/dependencies/dependency_list_export/part_spec.rb
@@ -55,4 +55,13 @@
       let_it_be(:model) { create(:dependency_list_export_part, organization: parent) }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      export_part = build_stubbed(:dependency_list_export_part, organization: organization)
+
+      expect(export_part.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/ee/spec/models/dependencies/dependency_list_export_spec.rb b/ee/spec/models/dependencies/dependency_list_export_spec.rb
index d97f92c91483f6..99bbaf5e5162d0 100644
--- a/ee/spec/models/dependencies/dependency_list_export_spec.rb
+++ b/ee/spec/models/dependencies/dependency_list_export_spec.rb
@@ -262,4 +262,25 @@
       let_it_be(:model) { create(:dependency_list_export, group: nil, project: nil, author: nil, organization: parent) }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns one of organization_id, group_id, or porject_id' do
+      parents = { organization: nil, group: nil, project: nil }
+
+      parents.each_key do |parent_type|
+        parent = build_stubbed(parent_type)
+        export = build_stubbed(:dependency_list_export, **parents.merge(parent_type => parent))
+
+        key_name = case parent_type
+                   when :organization then :organization_id
+                   when :group then :namespace_id
+                   when :project then :project_id
+                   end
+
+        expect(export.uploads_sharding_key).to eq(
+          { organization_id: nil, namespace_id: nil, project_id: nil }.merge(key_name => parent.id)
+        )
+      end
+    end
+  end
 end
diff --git a/ee/spec/models/issuable_metric_image_spec.rb b/ee/spec/models/issuable_metric_image_spec.rb
index abdbba41126643..01168efcf64309 100644
--- a/ee/spec/models/issuable_metric_image_spec.rb
+++ b/ee/spec/models/issuable_metric_image_spec.rb
@@ -84,4 +84,14 @@
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns namespace_id' do
+      namespace = build_stubbed(:namespace)
+      issue = build_stubbed(:issue, namespace: namespace)
+      issuable_metric_image = build_stubbed(:issuable_metric_image, issue: issue)
+
+      expect(issuable_metric_image.uploads_sharding_key).to eq(namespace_id: namespace.id)
+    end
+  end
 end
diff --git a/ee/spec/models/user_permission_export_upload_spec.rb b/ee/spec/models/user_permission_export_upload_spec.rb
index 222ab95d2e2f92..ad341afb385218 100644
--- a/ee/spec/models/user_permission_export_upload_spec.rb
+++ b/ee/spec/models/user_permission_export_upload_spec.rb
@@ -46,4 +46,10 @@
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns empty hash' do
+      expect(upload.uploads_sharding_key).to eq({})
+    end
+  end
 end
diff --git a/ee/spec/models/vulnerabilities/export/part_spec.rb b/ee/spec/models/vulnerabilities/export/part_spec.rb
index 3dc8bd6fe08e1a..942e3c4f7325c8 100644
--- a/ee/spec/models/vulnerabilities/export/part_spec.rb
+++ b/ee/spec/models/vulnerabilities/export/part_spec.rb
@@ -35,4 +35,13 @@
       let_it_be(:model) { create(:vulnerability_export_part, organization: parent) }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      export_part = build_stubbed(:vulnerability_export_part, organization: organization)
+
+      expect(export_part.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/ee/spec/models/vulnerabilities/export_spec.rb b/ee/spec/models/vulnerabilities/export_spec.rb
index 28fd3b73e31a74..b018128c21b2a9 100644
--- a/ee/spec/models/vulnerabilities/export_spec.rb
+++ b/ee/spec/models/vulnerabilities/export_spec.rb
@@ -318,4 +318,13 @@
       let_it_be(:model) { create(:vulnerability_export, author_id: parent.id) }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      export = build_stubbed(:vulnerability_export, organization: organization)
+
+      expect(export.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/ee/spec/models/vulnerabilities/remediation_spec.rb b/ee/spec/models/vulnerabilities/remediation_spec.rb
index 0d120e780fc203..eee67f8171ba9a 100644
--- a/ee/spec/models/vulnerabilities/remediation_spec.rb
+++ b/ee/spec/models/vulnerabilities/remediation_spec.rb
@@ -37,4 +37,13 @@
       let_it_be(:model) { create(:vulnerabilities_remediation, project: parent) }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id' do
+      project = build_stubbed(:project)
+      remediation = build_stubbed(:vulnerabilities_remediation, project_id: project.id)
+
+      expect(remediation.uploads_sharding_key).to eq(project_id: project.id)
+    end
+  end
 end
diff --git a/spec/db/schema_spec.rb b/spec/db/schema_spec.rb
index c55080af2bdca1..fa7e6d655972a7 100644
--- a/spec/db/schema_spec.rb
+++ b/spec/db/schema_spec.rb
@@ -216,7 +216,7 @@
       taggings: %w[tag_id taggable_id tagger_id],
       timelogs: %w[user_id],
       todos: %w[target_id commit_id],
-      uploads: %w[model_id],
+      uploads: %w[model_id organization_id namespace_id project_id],
       user_agent_details: %w[subject_id],
       users: %w[color_mode_id color_scheme_id created_by_id theme_id managing_group_id],
       users_star_projects: %w[user_id],
diff --git a/spec/factories/alert_management/metric_images.rb b/spec/factories/alert_management/metric_images.rb
index d7d8182af3ed58..bc1062b63f5e10 100644
--- a/spec/factories/alert_management/metric_images.rb
+++ b/spec/factories/alert_management/metric_images.rb
@@ -4,6 +4,7 @@
   factory :alert_metric_image, class: 'AlertManagement::MetricImage' do
     association :alert, factory: :alert_management_alert
     url { generate(:url) }
+    project_id { alert&.project_id }
 
     trait :local do
       file_store { ObjectStorage::Store::LOCAL }
diff --git a/spec/lib/gitlab/database/sharding_key_spec.rb b/spec/lib/gitlab/database/sharding_key_spec.rb
index f52392c6c2ef1e..70551fd8f7281a 100644
--- a/spec/lib/gitlab/database/sharding_key_spec.rb
+++ b/spec/lib/gitlab/database/sharding_key_spec.rb
@@ -198,7 +198,8 @@
       "oauth_access_tokens" => "https://gitlab.com/gitlab-org/gitlab/-/issues/496717",
       "oauth_access_grants" => "https://gitlab.com/gitlab-org/gitlab/-/issues/496717",
       "oauth_openid_requests" => "https://gitlab.com/gitlab-org/gitlab/-/issues/496717",
-      "oauth_device_grants" => "https://gitlab.com/gitlab-org/gitlab/-/issues/496717"
+      "oauth_device_grants" => "https://gitlab.com/gitlab-org/gitlab/-/issues/496717",
+      "uploads" => "https://gitlab.com/gitlab-org/gitlab/-/issues/398199"
     }
 
     has_lfk = ->(lfks) { lfks.any? { |k| k.options[:column] == 'organization_id' && k.to_table == 'organizations' } }
diff --git a/spec/models/abuse_report_spec.rb b/spec/models/abuse_report_spec.rb
index a9c71328cfac68..9534d06c6cab12 100644
--- a/spec/models/abuse_report_spec.rb
+++ b/spec/models/abuse_report_spec.rb
@@ -438,4 +438,10 @@
 
     it { is_expected.to define_enum_for(:category).with_values(**categories) }
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns empty hash' do
+      expect(report.uploads_sharding_key).to eq({})
+    end
+  end
 end
diff --git a/spec/models/achievements/achievement_spec.rb b/spec/models/achievements/achievement_spec.rb
index ac34efad5bf622..5d8d0e90bb0d58 100644
--- a/spec/models/achievements/achievement_spec.rb
+++ b/spec/models/achievements/achievement_spec.rb
@@ -32,4 +32,13 @@
   it_behaves_like Avatarable do
     let(:model) { create(:achievement, :with_avatar) }
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns namespace_id' do
+      namespace = build_stubbed(:namespace)
+      achievement = build_stubbed(:achievement, namespace: namespace)
+
+      expect(achievement.uploads_sharding_key).to eq(namespace_id: namespace.id)
+    end
+  end
 end
diff --git a/spec/models/alert_management/metric_image_spec.rb b/spec/models/alert_management/metric_image_spec.rb
index ca9104744236da..3b5ed23b85f653 100644
--- a/spec/models/alert_management/metric_image_spec.rb
+++ b/spec/models/alert_management/metric_image_spec.rb
@@ -15,4 +15,13 @@
     it { is_expected.to validate_length_of(:url).is_at_most(255) }
     it { is_expected.to validate_length_of(:url_text).is_at_most(128) }
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id' do
+      project = build_stubbed(:project)
+      metric_image = build_stubbed(:alert_metric_image, project_id: project.id)
+
+      expect(metric_image.uploads_sharding_key).to eq(project_id: project.id)
+    end
+  end
 end
diff --git a/spec/models/appearance_spec.rb b/spec/models/appearance_spec.rb
index 852e70b605b38d..492482765be018 100644
--- a/spec/models/appearance_spec.rb
+++ b/spec/models/appearance_spec.rb
@@ -166,4 +166,12 @@
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns epmty hash' do
+      appearance = build_stubbed(:appearance)
+
+      expect(appearance.uploads_sharding_key).to eq({})
+    end
+  end
 end
diff --git a/spec/models/bulk_imports/export_upload_spec.rb b/spec/models/bulk_imports/export_upload_spec.rb
index ca45fe73b0e760..a2a9ac8c2f9dca 100644
--- a/spec/models/bulk_imports/export_upload_spec.rb
+++ b/spec/models/bulk_imports/export_upload_spec.rb
@@ -3,7 +3,9 @@
 require 'spec_helper'
 
 RSpec.describe BulkImports::ExportUpload, type: :model, feature_category: :importers do
-  subject { described_class.new(export: create(:bulk_import_export)) }
+  let(:export) { create(:bulk_import_export) }
+
+  subject(:bulk_import_export) { described_class.new(export: export) }
 
   describe 'associations' do
     it { is_expected.to belong_to(:export) }
@@ -34,4 +36,13 @@ def find_callback(callbacks, key)
       expect(find_callback(after_save_callbacks, :store_export_file!)).to be_nil
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id or group_id' do
+      expect(bulk_import_export.uploads_sharding_key).to eq(
+        project_id: export.project_id,
+        namespace_id: export.group_id
+      )
+    end
+  end
 end
diff --git a/spec/models/design_management/action_spec.rb b/spec/models/design_management/action_spec.rb
index f2b8fcaa256785..8ce67adf1a9195 100644
--- a/spec/models/design_management/action_spec.rb
+++ b/spec/models/design_management/action_spec.rb
@@ -129,4 +129,14 @@
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns namespace_id' do
+      namespace = build_stubbed(:namespace)
+      design = build_stubbed(:design, namespace_id: namespace.id)
+      design_action = build_stubbed(:design_action, design: design)
+
+      expect(design_action.uploads_sharding_key).to eq(namespace_id: namespace.id)
+    end
+  end
 end
diff --git a/spec/models/import_export_upload_spec.rb b/spec/models/import_export_upload_spec.rb
index 1108b303095595..6ac7ef0611d3e1 100644
--- a/spec/models/import_export_upload_spec.rb
+++ b/spec/models/import_export_upload_spec.rb
@@ -5,7 +5,7 @@
 RSpec.describe ImportExportUpload do
   let(:project) { create(:project) }
 
-  subject { described_class.new(project: project) }
+  subject(:import_export_upload) { described_class.new(project: project) }
 
   shared_examples 'stores the Import/Export file' do |method|
     it 'stores the import file' do
@@ -127,4 +127,13 @@ def find_callback(callbacks, key)
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id / group_id' do
+      expect(import_export_upload.uploads_sharding_key).to eq(
+        project_id: project.id,
+        namespace_id: nil
+      )
+    end
+  end
 end
diff --git a/spec/models/namespace_spec.rb b/spec/models/namespace_spec.rb
index d0bf0b0c00ceb0..f6f321e7e1b760 100644
--- a/spec/models/namespace_spec.rb
+++ b/spec/models/namespace_spec.rb
@@ -2696,4 +2696,13 @@
       it { expect(nested_group.web_url).to include("groups/#{nested_group.full_path}") }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      namespace = build_stubbed(:namespace, organization: organization)
+
+      expect(namespace.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/spec/models/note_spec.rb b/spec/models/note_spec.rb
index 094d76886a5da8..3e17b616d3eafa 100644
--- a/spec/models/note_spec.rb
+++ b/spec/models/note_spec.rb
@@ -2014,4 +2014,13 @@ def retrieve_participants
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns namespace_id' do
+      namespace = build_stubbed(:namespace)
+      note = build_stubbed(:note, namespace: namespace)
+
+      expect(note.uploads_sharding_key).to eq(namespace_id: namespace.id)
+    end
+  end
 end
diff --git a/spec/models/organizations/organization_detail_spec.rb b/spec/models/organizations/organization_detail_spec.rb
index dd49274e7dddfe..ee47e6b4c0f5db 100644
--- a/spec/models/organizations/organization_detail_spec.rb
+++ b/spec/models/organizations/organization_detail_spec.rb
@@ -32,4 +32,12 @@
       let(:uploader_class) { AttachmentUploader }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization_detail = build_stubbed(:organization_detail)
+
+      expect(organization_detail.uploads_sharding_key).to eq(organization_id: organization_detail.organization_id)
+    end
+  end
 end
diff --git a/spec/models/personal_snippet_spec.rb b/spec/models/personal_snippet_spec.rb
index d74ba1f5483e7d..88e1413b710027 100644
--- a/spec/models/personal_snippet_spec.rb
+++ b/spec/models/personal_snippet_spec.rb
@@ -35,4 +35,13 @@
       expect(snippet.parent_user).to eq(snippet.author)
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      snippet = build_stubbed(:personal_snippet, organization: organization)
+
+      expect(snippet.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/spec/models/project_spec.rb b/spec/models/project_spec.rb
index 3f2743515c7350..71e6191e5b7e36 100644
--- a/spec/models/project_spec.rb
+++ b/spec/models/project_spec.rb
@@ -9750,4 +9750,13 @@ def create_build(new_pipeline = pipeline, name = 'test')
       it { expect(project.placeholder_reference_store).to be_nil }
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns namespace_id' do
+      namespace = build_stubbed(:namespace)
+      project = build_stubbed(:project, namespace: namespace)
+
+      expect(project.uploads_sharding_key).to eq(namespace_id: namespace.id)
+    end
+  end
 end
diff --git a/spec/models/projects/import_export/relation_export_upload_spec.rb b/spec/models/projects/import_export/relation_export_upload_spec.rb
index 14c17c4c9a26f4..89b29f97a9cb4e 100644
--- a/spec/models/projects/import_export/relation_export_upload_spec.rb
+++ b/spec/models/projects/import_export/relation_export_upload_spec.rb
@@ -5,7 +5,8 @@
 RSpec.describe Projects::ImportExport::RelationExportUpload, type: :model, feature_category: :importers do
   subject { described_class.new(relation_export: project_relation_export) }
 
-  let_it_be(:project_relation_export) { create(:project_relation_export) }
+  let_it_be(:project) { create(:project) }
+  let_it_be(:project_relation_export) { create(:project_relation_export, project_id: project.id) }
 
   describe 'associations' do
     it { is_expected.to belong_to(:relation_export) }
@@ -15,10 +16,19 @@
     let_it_be(:project_export_job_1) { create(:project_export_job) }
     let_it_be(:project_export_job_2) { create(:project_export_job) }
 
-    let_it_be(:relation_export_1) { create(:project_relation_export, project_export_job: project_export_job_1) }
-    let_it_be(:relation_export_2) { create(:project_relation_export, project_export_job: project_export_job_2) }
+    let_it_be(:relation_export_1) do
+      create(:project_relation_export, project_export_job: project_export_job_1,
+        project_id: project_export_job_1.project_id)
+    end
+
+    let_it_be(:relation_export_2) do
+      create(:project_relation_export, project_export_job: project_export_job_2,
+        project_id: project_export_job_2.project_id)
+    end
+
     let_it_be(:relation_export_3) do
-      create(:project_relation_export, project_export_job: project_export_job_1, relation: 'milestones')
+      create(:project_relation_export, project_export_job: project_export_job_1,
+        project_id: project_export_job_1.project_id, relation: 'milestones')
     end
 
     let_it_be(:relation_export_upload_1) { create(:relation_export_upload, relation_export: relation_export_1) }
@@ -65,4 +75,14 @@ def find_callback(callbacks, key)
       expect(find_callback(after_save_callbacks, :store_export_file!)).to be_nil
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns project_id' do
+      project = build_stubbed(:project)
+      export = build_stubbed(:project_relation_export, project_id: project.id)
+      export_upload = build_stubbed(:relation_export_upload, relation_export: export)
+
+      expect(export_upload.uploads_sharding_key).to eq(project_id: project.id)
+    end
+  end
 end
diff --git a/spec/models/projects/topic_spec.rb b/spec/models/projects/topic_spec.rb
index 8a01dda9fa06e1..3553cec34568bc 100644
--- a/spec/models/projects/topic_spec.rb
+++ b/spec/models/projects/topic_spec.rb
@@ -132,4 +132,13 @@
       expect(topic.title_or_name).to eq('topic')
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns organization_id' do
+      organization = build_stubbed(:organization)
+      topic = build_stubbed(:topic, organization: organization)
+
+      expect(topic.uploads_sharding_key).to eq(organization_id: organization.id)
+    end
+  end
 end
diff --git a/spec/models/upload_spec.rb b/spec/models/upload_spec.rb
index f8a0c31e2f23c3..2e1e3cba49be89 100644
--- a/spec/models/upload_spec.rb
+++ b/spec/models/upload_spec.rb
@@ -48,6 +48,27 @@
       end
     end
 
+    context 'before_save' do
+      it 'sets sharding key on create' do
+        project = build_stubbed(:project)
+        upload = build(:upload, model: project)
+
+        expect { upload.save! }
+          .to change { upload.namespace_id }.from(nil)
+          .to(project.uploads_sharding_key.each_value.first)
+      end
+
+      it 'sets sharding key on update' do
+        project = build_stubbed(:project)
+        upload = create(:upload, model: project)
+        other_project = build_stubbed(:project)
+
+        expect { upload.update!(model: other_project) }
+          .to change { upload.namespace_id }.from(project.uploads_sharding_key.each_value.first)
+          .to(other_project.uploads_sharding_key.each_value.first)
+      end
+    end
+
     describe 'after_destroy' do
       context 'uploader is FileUploader-based' do
         subject { create(:upload, :issuable_upload) }
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index 269f1e72af7777..7fae12955a3a2d 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -9109,4 +9109,12 @@ def owner_class_attribute_default; end
       end
     end
   end
+
+  describe '#uploads_sharding_key' do
+    it 'returns empty hash' do
+      user = build_stubbed(:user)
+
+      expect(user.uploads_sharding_key).to eq({})
+    end
+  end
 end
diff --git a/spec/uploaders/object_storage_spec.rb b/spec/uploaders/object_storage_spec.rb
index 4f9521ac2f0ea6..20fa810b86b8a8 100644
--- a/spec/uploaders/object_storage_spec.rb
+++ b/spec/uploaders/object_storage_spec.rb
@@ -1088,6 +1088,12 @@ def escape_path(path)
           end
 
           context 'when uploaded file remote_id matches a pending direct upload entry' do
+            let(:uploader_class) do
+              Class.new(GitlabUploader) do
+                include ObjectStorage::Concern
+              end
+            end
+
             let(:object) { build_stubbed(:ci_job_artifact) }
             let(:final_path) { '@final/test/123123' }
             let(:fog_config) { Gitlab.config.uploads.object_store }
-- 
GitLab