Calculate checksums

by copy-pasting in the whole `Upload` class.

Also, fix `Namespace` `model_type` (it should not be `Group`).

3 files changed, 170 insertions, 110 deletions

diff --git a/lib/gitlab/background_migration/populate_untracked_uploads.rb b/lib/gitlab/background_migration/populate_untracked_uploads.rb
index acd424f4558..1773b53bd68 100644
--- a/lib/gitlab/background_migration/populate_untracked_uploads.rb
+++ b/lib/gitlab/background_migration/populate_untracked_uploads.rb
@@ -35,7 +35,7 @@ module Gitlab
           {
             pattern: /\A-\/system\/group\/avatar\/(\d+)/,
             uploader: 'AvatarUploader',
-            model_type: 'Group',
+            model_type: 'Namespace',
           },
           {
             pattern: /\A-\/system\/project\/avatar\/(\d+)/,
@@ -150,8 +150,71 @@ module Gitlab
         end
       end
 
+      # Copy-pasted class for less fragile migration
       class Upload < ActiveRecord::Base
-        self.table_name = 'uploads'
+        self.table_name = 'uploads' # This is the only line different from copy-paste
+
+        # Upper limit for foreground checksum processing
+        CHECKSUM_THRESHOLD = 100.megabytes
+
+        belongs_to :model, polymorphic: true # rubocop:disable Cop/PolymorphicAssociations
+
+        validates :size, presence: true
+        validates :path, presence: true
+        validates :model, presence: true
+        validates :uploader, presence: true
+
+        before_save  :calculate_checksum, if:     :foreground_checksum?
+        after_commit :schedule_checksum,  unless: :foreground_checksum?
+
+        def self.remove_path(path)
+          where(path: path).destroy_all
+        end
+
+        def self.record(uploader)
+          remove_path(uploader.relative_path)
+
+          create(
+            size: uploader.file.size,
+            path: uploader.relative_path,
+            model: uploader.model,
+            uploader: uploader.class.to_s
+          )
+        end
+
+        def absolute_path
+          return path unless relative_path?
+
+          uploader_class.absolute_path(self)
+        end
+
+        def calculate_checksum
+          return unless exist?
+
+          self.checksum = Digest::SHA256.file(absolute_path).hexdigest
+        end
+
+        def exist?
+          File.exist?(absolute_path)
+        end
+
+        private
+
+        def foreground_checksum?
+          size <= CHECKSUM_THRESHOLD
+        end
+
+        def schedule_checksum
+          UploadChecksumWorker.perform_async(id)
+        end
+
+        def relative_path?
+          !path.start_with?('/')
+        end
+
+        def uploader_class
+          Object.const_get(uploader)
+        end
       end
 
       def perform(start_id, end_id)
diff --git a/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb b/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
index ae6a712f2ee..c61a207d012 100644
--- a/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
+++ b/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
@@ -97,61 +97,53 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads::UnhashedUploadFi
     let(:uploaded_file) { fixture_file_upload(fixture) }
 
     context 'for an appearance logo file path' do
-      let(:appearance) { create(:appearance) }
-      let(:unhashed_upload_file) { described_class.create!(path: appearance.logo.file.file) }
+      let(:model) { create(:appearance) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.logo.file.file) }
 
       before do
-        appearance.update!(logo: uploaded_file)
-        appearance.uploads.delete_all
+        model.update!(logo: uploaded_file)
+        model.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/appearance/logo/#{appearance.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => appearance.id,
-          "model_type" => "Appearance",
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "uploads/-/system/appearance/logo/#{model.id}/rails_sample.jpg",
           "uploader"   => "AttachmentUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for an appearance header_logo file path' do
-      let(:appearance) { create(:appearance) }
-      let(:unhashed_upload_file) { described_class.create!(path: appearance.header_logo.file.file) }
+      let(:model) { create(:appearance) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.header_logo.file.file) }
 
       before do
-        appearance.update!(header_logo: uploaded_file)
-        appearance.uploads.delete_all
+        model.update!(header_logo: uploaded_file)
+        model.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/appearance/header_logo/#{appearance.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => appearance.id,
-          "model_type" => "Appearance",
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "uploads/-/system/appearance/header_logo/#{model.id}/rails_sample.jpg",
           "uploader"   => "AttachmentUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for a pre-Markdown Note attachment file path' do
-      let(:note) { create(:note) }
-      let(:unhashed_upload_file) { described_class.create!(path: note.attachment.file.file) }
+      let(:model) { create(:note) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.attachment.file.file) }
 
       before do
-        note.update!(attachment: uploaded_file)
+        model.update!(attachment: uploaded_file)
         upload_class.delete_all
       end
 
@@ -161,115 +153,99 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads::UnhashedUploadFi
         end.to change { upload_class.count }.from(0).to(1)
 
         expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/note/attachment/#{note.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => note.id,
+          "path"       => "uploads/-/system/note/attachment/#{model.id}/rails_sample.jpg",
+          "model_id"   => model.id,
           "model_type" => "Note",
           "uploader"   => "AttachmentUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for a user avatar file path' do
-      let(:user) { create(:user) }
-      let(:unhashed_upload_file) { described_class.create!(path: user.avatar.file.file) }
+      let(:model) { create(:user) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.avatar.file.file) }
 
       before do
-        user.update!(avatar: uploaded_file)
-        upload_class.delete_all
+        model.update!(avatar: uploaded_file)
+        model.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/user/avatar/#{user.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => user.id,
-          "model_type" => "User",
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "uploads/-/system/user/avatar/#{model.id}/rails_sample.jpg",
           "uploader"   => "AvatarUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for a group avatar file path' do
-      let(:group) { create(:group) }
-      let(:unhashed_upload_file) { described_class.create!(path: group.avatar.file.file) }
+      let(:model) { create(:group) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.avatar.file.file) }
 
       before do
-        group.update!(avatar: uploaded_file)
-        upload_class.delete_all
+        model.update!(avatar: uploaded_file)
+        model.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/group/avatar/#{group.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => group.id,
-          "model_type" => "Group",
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "uploads/-/system/group/avatar/#{model.id}/rails_sample.jpg",
+          "model_id"   => model.id,
+          "model_type" => "Namespace", # Explicitly calling this out because it was unexpected to me (I assumed it should be "Group")
           "uploader"   => "AvatarUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for a project avatar file path' do
-      let(:project) { create(:project) }
-      let(:unhashed_upload_file) { described_class.create!(path: project.avatar.file.file) }
+      let(:model) { create(:project) }
+      let(:unhashed_upload_file) { described_class.create!(path: model.avatar.file.file) }
 
       before do
-        project.update!(avatar: uploaded_file)
-        upload_class.delete_all
+        model.update!(avatar: uploaded_file)
+        model.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "uploads/-/system/project/avatar/#{project.id}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => project.id,
-          "model_type" => "Project",
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "uploads/-/system/project/avatar/#{model.id}/rails_sample.jpg",
           "uploader"   => "AvatarUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
     context 'for a project Markdown attachment (notes, issues, MR descriptions) file path' do
-      let(:project) { create(:project) }
-      let(:unhashed_upload_file) { described_class.new(path: "#{Gitlab::BackgroundMigration::PrepareUnhashedUploads::UPLOAD_DIR}/#{project.full_path}/#{project.uploads.first.path}") }
+      let(:model) { create(:project) }
+      let(:unhashed_upload_file) { described_class.new(path: "#{Gitlab::BackgroundMigration::PrepareUnhashedUploads::UPLOAD_DIR}/#{model.full_path}/#{model.uploads.first.path}") }
 
       before do
-        UploadService.new(project, uploaded_file, FileUploader).execute # Markdown upload
+        UploadService.new(model, uploaded_file, FileUploader).execute # Markdown upload
         unhashed_upload_file.save!
-        upload_class.delete_all
+        model.reload.uploads.delete_all
       end
 
       it 'creates an Upload record' do
         expect do
           unhashed_upload_file.add_to_uploads
-        end.to change { upload_class.count }.from(0).to(1)
+        end.to change { model.reload.uploads.count }.from(0).to(1)
 
-        random_hex = unhashed_upload_file.path.match(/\/(\h+)\/rails_sample.jpg/)[1]
-        expect(upload_class.first.attributes).to include({
-          "size"       => 35255,
-          "path"       => "#{random_hex}/rails_sample.jpg",
-          "checksum"   => nil,
-          "model_id"   => project.id,
-          "model_type" => "Project",
+        hex_secret = unhashed_upload_file.path.match(/\/(\h+)\/rails_sample.jpg/)[1]
+        expect(model.uploads.first.attributes).to include({
+          "path"       => "#{hex_secret}/rails_sample.jpg",
           "uploader"   => "FileUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
       end
     end
   end
@@ -441,8 +417,8 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads::UnhashedUploadFi
     context 'for a group avatar file path' do
       let(:unhashed_upload_file) { described_class.create!(path: "#{Gitlab::BackgroundMigration::PrepareUnhashedUploads::UPLOAD_DIR}/-/system/group/avatar/1234/avatar.jpg") }
 
-      it 'returns Group as a string' do
-        expect(unhashed_upload_file.model_type).to eq('Group')
+      it 'returns Namespace as a string' do
+        expect(unhashed_upload_file.model_type).to eq('Namespace')
       end
     end
 
@@ -578,4 +554,11 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads::UnhashedUploadFi
       end
     end
   end
+
+  def rails_sample_jpg_attrs
+    {
+      "size"       => 35255,
+      "checksum"   => 'f2d1fd9d8d8a3368d468fa067888605d74a66f41c16f55979ceaf2af77375844'
+    }
+  end
 end
diff --git a/spec/migrations/track_untracked_uploads_spec.rb b/spec/migrations/track_untracked_uploads_spec.rb
index 8db41786397..9539993be31 100644
--- a/spec/migrations/track_untracked_uploads_spec.rb
+++ b/spec/migrations/track_untracked_uploads_spec.rb
@@ -49,6 +49,7 @@ describe TrackUntrackedUploads, :migration, :sidekiq do
     let(:project1) { create(:project) }
     let(:project2) { create(:project) }
     let(:appearance) { create(:appearance) }
+    let(:uploads) { table(:uploads) }
 
     before do
       fixture = Rails.root.join('spec', 'fixtures', 'rails_sample.jpg')
@@ -73,46 +74,52 @@ describe TrackUntrackedUploads, :migration, :sidekiq do
       appearance.uploads.last.destroy
     end
 
-    it 'tracks untracked migrations' do
+    it 'tracks untracked uploads' do
       Sidekiq::Testing.inline! do
-        migrate!
+        expect do
+          migrate!
+        end.to change { uploads.count }.from(4).to(8)
 
-        # Tracked uploads still exist
-        expect(user1.reload.uploads.first.attributes).to include({
-          "path" => "uploads/-/system/user/avatar/#{user1.id}/rails_sample.jpg",
-          "uploader" => "AvatarUploader"
-        })
-        expect(project1.reload.uploads.first.attributes).to include({
-          "path" => "uploads/-/system/project/avatar/#{project1.id}/rails_sample.jpg",
-          "uploader" => "AvatarUploader"
-        })
-        expect(appearance.reload.uploads.first.attributes).to include({
-          "path" => "uploads/-/system/appearance/logo/#{appearance.id}/rails_sample.jpg",
-          "uploader" => "AttachmentUploader"
-        })
-        expect(project1.uploads.last.attributes).to include({
-          "path" => @project1_markdown_upload_path,
-          "uploader" => "FileUploader"
-        })
-
-        # Untracked uploads are now tracked
         expect(user2.reload.uploads.first.attributes).to include({
           "path" => "uploads/-/system/user/avatar/#{user2.id}/rails_sample.jpg",
           "uploader" => "AvatarUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
         expect(project2.reload.uploads.first.attributes).to include({
           "path" => "uploads/-/system/project/avatar/#{project2.id}/rails_sample.jpg",
           "uploader" => "AvatarUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
         expect(appearance.reload.uploads.count).to eq(2)
         expect(appearance.uploads.last.attributes).to include({
           "path" => "uploads/-/system/appearance/header_logo/#{appearance.id}/rails_sample.jpg",
           "uploader" => "AttachmentUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
         expect(project2.uploads.last.attributes).to include({
           "path" => @project2_markdown_upload_path,
           "uploader" => "FileUploader"
-        })
+        }.merge(rails_sample_jpg_attrs))
+      end
+    end
+
+    it 'ignores already-tracked uploads' do
+      Sidekiq::Testing.inline! do
+        migrate!
+
+        expect(user1.reload.uploads.first.attributes).to include({
+          "path" => "uploads/-/system/user/avatar/#{user1.id}/rails_sample.jpg",
+          "uploader" => "AvatarUploader",
+        }.merge(rails_sample_jpg_attrs))
+        expect(project1.reload.uploads.first.attributes).to include({
+          "path" => "uploads/-/system/project/avatar/#{project1.id}/rails_sample.jpg",
+          "uploader" => "AvatarUploader"
+        }.merge(rails_sample_jpg_attrs))
+        expect(appearance.reload.uploads.first.attributes).to include({
+          "path" => "uploads/-/system/appearance/logo/#{appearance.id}/rails_sample.jpg",
+          "uploader" => "AttachmentUploader"
+        }.merge(rails_sample_jpg_attrs))
+        expect(project1.uploads.last.attributes).to include({
+          "path" => @project1_markdown_upload_path,
+          "uploader" => "FileUploader"
+        }.merge(rails_sample_jpg_attrs))
       end
     end
 
@@ -125,4 +132,11 @@ describe TrackUntrackedUploads, :migration, :sidekiq do
       end
     end
   end
+
+  def rails_sample_jpg_attrs
+    {
+      "size"       => 35255,
+      "checksum"   => 'f2d1fd9d8d8a3368d468fa067888605d74a66f41c16f55979ceaf2af77375844'
+    }
+  end
 end