Calculate checksums

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)

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