another round of fixes

3 files changed, 86 insertions, 0 deletions

diff --git a/lib/tasks/gitlab/uploads/check.rake b/lib/tasks/gitlab/uploads/check.rake
new file mode 100644
index 00000000000..2a5afb82029
--- /dev/null
+++ b/lib/tasks/gitlab/uploads/check.rake
@@ -0,0 +1,30 @@
+require_relative 'helpers.rb'
+
+namespace :gitlab do
+  namespace :uploads do
+    desc 'GitLab | Uploads | Check integrity of uploaded files'
+    task check: :environment do
+      include UploadTaskHelpers
+
+      puts 'Checking integrity of uploaded files'
+
+      uploads_batches do |batch|
+        batch.each do |upload|
+          begin
+            puts "- Checking file (#{upload.id}): #{upload.absolute_path}".color(:green)
+
+            if upload.exist?
+              check_checksum(upload)
+            else
+              puts "  * File does not exist on the file system".color(:red)
+            end
+          rescue ObjectStorage::RemoteStoreError
+            puts "- File (#{upload.id}): File is stored remotely, skipping".color(:yellow)
+          end
+        end
+      end
+
+      puts 'Done!'
+    end
+  end
+end
diff --git a/lib/tasks/gitlab/uploads/helpers.rb b/lib/tasks/gitlab/uploads/helpers.rb
new file mode 100644
index 00000000000..239cba57297
--- /dev/null
+++ b/lib/tasks/gitlab/uploads/helpers.rb
@@ -0,0 +1,23 @@
+module UploadTaskHelpers
+  def batch_size
+    ENV.fetch('BATCH', 200).to_i
+  end
+
+  def calculate_checksum(absolute_path)
+    Digest::SHA256.file(absolute_path).hexdigest
+  end
+
+  def check_checksum(upload)
+    checksum = calculate_checksum(upload.absolute_path)
+
+    if checksum != upload.checksum
+      puts "  * File checksum (#{checksum}) does not match the one in the database (#{upload.checksum})".color(:red)
+    end
+  end
+
+  def uploads_batches(&block)
+    Upload.all.in_batches(of: batch_size, start: ENV['ID_FROM'], finish: ENV['ID_TO']) do |relation| # rubocop: disable Cop/InBatches
+      yield relation
+    end
+  end
+end
diff --git a/lib/tasks/gitlab/uploads/migrate.rake b/lib/tasks/gitlab/uploads/migrate.rake
new file mode 100644
index 00000000000..c26c3ccb3be
--- /dev/null
+++ b/lib/tasks/gitlab/uploads/migrate.rake
@@ -0,0 +1,33 @@
+namespace :gitlab do
+  namespace :uploads do
+    desc 'GitLab | Uploads | Migrate the uploaded files to object storage'
+    task :migrate, [:uploader_class, :model_class, :mounted_as] => :environment do |task, args|
+      batch_size     = ENV.fetch('BATCH', 200).to_i
+      @to_store      = ObjectStorage::Store::REMOTE
+      @mounted_as    = args.mounted_as&.gsub(':', '')&.to_sym
+      @uploader_class = args.uploader_class.constantize
+      @model_class    = args.model_class.constantize
+
+      uploads.each_batch(of: batch_size, &method(:enqueue_batch)) # rubocop: disable Cop/InBatches
+    end
+
+    def enqueue_batch(batch, index)
+      job = ObjectStorage::MigrateUploadsWorker.enqueue!(batch,
+                                                         @mounted_as,
+                                                         @to_store)
+      puts "Enqueued job ##{index}: #{job}"
+    rescue ObjectStorage::MigrateUploadsWorker::SanityCheckError => e
+      # continue for the next batch
+      puts "Could not enqueue batch (#{batch.ids}) #{e.message}".color(:red)
+    end
+
+    def uploads
+      Upload.class_eval { include EachBatch } unless Upload < EachBatch
+
+      Upload
+        .where.not(store: @to_store)
+        .where(uploader: @uploader_class.to_s,
+               model_type: @model_class.base_class.sti_name)
+    end
+  end
+end