From 3ab5af14a2e2ea9668b611f7b4bcf56b9f722161 Mon Sep 17 00:00:00 2001
From: Mark Chao <mchao@gitlab.com>
Date: Wed, 19 Dec 2018 16:34:34 +0800
Subject: [PATCH] Check if specific type of background migration are done

Useful for checking progress.
---
 lib/gitlab/background_migration.rb           | 13 ++++++
 spec/lib/gitlab/background_migration_spec.rb | 44 ++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/lib/gitlab/background_migration.rb b/lib/gitlab/background_migration.rb
index d72befce571..6cf40e2d4ca 100644
--- a/lib/gitlab/background_migration.rb
+++ b/lib/gitlab/background_migration.rb
@@ -51,6 +51,19 @@ module Gitlab
       migration_class_for(class_name).new.perform(*arguments)
     end
 
+    def self.exists?(migration_class)
+      enqueued = Sidekiq::Queue.new(self.queue)
+      scheduled = Sidekiq::ScheduledSet.new
+
+      [enqueued, scheduled].each do |queue|
+        queue.each do |job|
+          return true if job.queue == self.queue && job.args.first == migration_class
+        end
+      end
+
+      false
+    end
+
     def self.migration_class_for(class_name)
       const_get(class_name)
     end
diff --git a/spec/lib/gitlab/background_migration_spec.rb b/spec/lib/gitlab/background_migration_spec.rb
index 4ad69aeba43..8a83b76fd94 100644
--- a/spec/lib/gitlab/background_migration_spec.rb
+++ b/spec/lib/gitlab/background_migration_spec.rb
@@ -119,4 +119,48 @@ describe Gitlab::BackgroundMigration do
       described_class.perform('Foo', [10, 20])
     end
   end
+
+  describe '.exists?' do
+    context 'when there are enqueued jobs present' do
+      let(:queue) do
+        [double(args: ['Foo', [10, 20]], queue: described_class.queue)]
+      end
+
+      before do
+        allow(Sidekiq::Queue).to receive(:new)
+          .with(described_class.queue)
+          .and_return(queue)
+      end
+
+      it 'returns true if specific job exists' do
+        expect(described_class.exists?('Foo')).to eq(true)
+      end
+
+      it 'returns false if specific job does not exist' do
+        expect(described_class.exists?('Bar')).to eq(false)
+      end
+    end
+
+    context 'when there are scheduled jobs present', :sidekiq, :redis do
+      before do
+        Sidekiq::Testing.disable! do
+          BackgroundMigrationWorker.perform_in(10.minutes, 'Foo')
+
+          expect(Sidekiq::ScheduledSet.new).to be_one
+        end
+      end
+
+      after do
+        Sidekiq::ScheduledSet.new.clear
+      end
+
+      it 'returns true if specific job exists' do
+        expect(described_class.exists?('Foo')).to eq(true)
+      end
+
+      it 'returns false if specific job does not exist' do
+        expect(described_class.exists?('Bar')).to eq(false)
+      end
+    end
+  end
 end