Add more specs and code improvements

4 files changed, 32 insertions, 29 deletions

diff --git a/lib/gitlab/background_migration/delete_diff_files.rb b/lib/gitlab/background_migration/delete_diff_files.rb
index 4cb068eb71a..0066faa23dd 100644
--- a/lib/gitlab/background_migration/delete_diff_files.rb
+++ b/lib/gitlab/background_migration/delete_diff_files.rb
@@ -20,14 +20,14 @@ module Gitlab
       def perform(ids)
         @ids = ids
 
-        reschedule_if_needed([ids]) do
+        reschedule_if_needed(ids) do
           prune_diff_files
         end
       end
 
       private
 
-      def need_reschedule?
+      def should_reschedule?
         wait_for_deadtuple_vacuum?(MergeRequestDiffFile.table_name)
       end
 
diff --git a/lib/gitlab/background_migration/helpers/reschedulable.rb b/lib/gitlab/background_migration/helpers/reschedulable.rb
index 6087181609e..60ab62cd973 100644
--- a/lib/gitlab/background_migration/helpers/reschedulable.rb
+++ b/lib/gitlab/background_migration/helpers/reschedulable.rb
@@ -5,7 +5,8 @@ module Gitlab
     module Helpers
       # == Reschedulable helper
       #
-      # Allows background migrations to be reschedule itself if a condition is not met.
+      # Allows background migrations to be rescheduled if a condition is met,
+      # the condition should be overridden in classes in #should_reschedule? method.
       #
       # For example, check DeleteDiffFiles migration which is rescheduled if dead tuple count
       # on DB is not acceptable.
@@ -13,18 +14,34 @@ module Gitlab
       module Reschedulable
         extend ActiveSupport::Concern
 
-        def reschedule_if_needed(args, &block)
-          if need_reschedule?
-            BackgroundMigrationWorker.perform_in(vacuum_wait_time, self.class.name.demodulize, args)
+        # Use this method to perform the background migration and it will be rescheduled
+        # if #should_reschedule? returns true.
+        def reschedule_if_needed(*args, &block)
+          if should_reschedule?
+            BackgroundMigrationWorker.perform_in(wait_time, self.class.name.demodulize, args)
           else
             yield
           end
         end
 
         # Override this on base class if you need a different reschedule condition
-        # def need_reschedule?
-        #   raise NotImplementedError, "#{self.class} does not implement #{__method__}"
-        # end
+        def should_reschedule?
+          raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+        end
+
+        # Override in subclass if a different dead tuple threshold
+        def dead_tuples_threshold
+          @dead_tuples_threshold ||= 50_000
+        end
+
+        # Override in subclass if a different wait time
+        def wait_time
+          @wait_time ||= 5.minutes
+        end
+
+        def execute_statement(sql)
+          ActiveRecord::Base.connection.execute(sql)
+        end
 
         def wait_for_deadtuple_vacuum?(table_name)
           return false unless Gitlab::Database.postgresql?
@@ -39,20 +56,6 @@ module Gitlab
 
           dead_tuple&.fetch('n_dead_tup', 0).to_i
         end
-
-        def execute_statement(sql)
-          ActiveRecord::Base.connection.execute(sql)
-        end
-
-        # Override in subclass if you need a different dead tuple threshold
-        def dead_tuples_threshold
-          @dead_tuples_threshold ||= 50_000
-        end
-
-        # Override in subclass if you need a different vacuum wait time
-        def vacuum_wait_time
-          @vacuum_wait_time ||= 5.minutes
-        end
       end
     end
   end
diff --git a/lib/gitlab/background_migration/sync_issues_state_id.rb b/lib/gitlab/background_migration/sync_issues_state_id.rb
index f95c433c64c..50841c3fe4d 100644
--- a/lib/gitlab/background_migration/sync_issues_state_id.rb
+++ b/lib/gitlab/background_migration/sync_issues_state_id.rb
@@ -9,8 +9,8 @@ module Gitlab
       def perform(start_id, end_id)
         Rails.logger.info("Issues - Populating state_id: #{start_id} - #{end_id}")
 
-        reschedule_if_needed([start_id, end_id]) do
-          ActiveRecord::Base.connection.execute <<~SQL
+        reschedule_if_needed(start_id, end_id) do
+          execute_statement <<~SQL
             UPDATE issues
             SET state_id =
               CASE state
@@ -25,7 +25,7 @@ module Gitlab
 
       private
 
-      def need_reschedule?
+      def should_reschedule?
         wait_for_deadtuple_vacuum?('issues')
       end
     end
diff --git a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
index 5d92553f577..4cf8e993c52 100644
--- a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
+++ b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
@@ -9,8 +9,8 @@ module Gitlab
       def perform(start_id, end_id)
         Rails.logger.info("Merge Requests - Populating state_id: #{start_id} - #{end_id}")
 
-        reschedule_if_needed([start_id, end_id]) do
-          ActiveRecord::Base.connection.execute <<~SQL
+        reschedule_if_needed(start_id, end_id) do
+          execute_statement <<~SQL
             UPDATE merge_requests
             SET state_id =
               CASE state
@@ -27,7 +27,7 @@ module Gitlab
 
       private
 
-      def need_reschedule?
+      def should_reschedule?
         wait_for_deadtuple_vacuum?('issues')
       end
     end