Make migrations reschedulable

4 files changed, 47 insertions, 27 deletions

diff --git a/lib/gitlab/background_migration/delete_diff_files.rb b/lib/gitlab/background_migration/delete_diff_files.rb
index 11851c23ee3..4cb068eb71a 100644
--- a/lib/gitlab/background_migration/delete_diff_files.rb
+++ b/lib/gitlab/background_migration/delete_diff_files.rb
@@ -27,7 +27,7 @@ module Gitlab
 
       private
 
-      def should_reschedule?
+      def need_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 3b653059858..6087181609e 100644
--- a/lib/gitlab/background_migration/helpers/reschedulable.rb
+++ b/lib/gitlab/background_migration/helpers/reschedulable.rb
@@ -14,7 +14,7 @@ module Gitlab
         extend ActiveSupport::Concern
 
         def reschedule_if_needed(args, &block)
-          if should_reschedule?
+          if need_reschedule?
             BackgroundMigrationWorker.perform_in(vacuum_wait_time, self.class.name.demodulize, args)
           else
             yield
@@ -22,9 +22,9 @@ module Gitlab
         end
 
         # Override this on base class if you need a different reschedule condition
-        def should_reschedule?
-          raise NotImplementedError, "#{self.class} does not implement #{__method__}"
-        end
+        # def need_reschedule?
+        #   raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+        # end
 
         def wait_for_deadtuple_vacuum?(table_name)
           return false unless Gitlab::Database.postgresql?
diff --git a/lib/gitlab/background_migration/sync_issues_state_id.rb b/lib/gitlab/background_migration/sync_issues_state_id.rb
index 33b997c8533..f95c433c64c 100644
--- a/lib/gitlab/background_migration/sync_issues_state_id.rb
+++ b/lib/gitlab/background_migration/sync_issues_state_id.rb
@@ -4,19 +4,29 @@
 module Gitlab
   module BackgroundMigration
     class SyncIssuesStateId
+      include Helpers::Reschedulable
+
       def perform(start_id, end_id)
         Rails.logger.info("Issues - Populating state_id: #{start_id} - #{end_id}")
 
-        ActiveRecord::Base.connection.execute <<~SQL
-          UPDATE issues
-          SET state_id =
-            CASE state
-            WHEN 'opened' THEN 1
-            WHEN 'closed' THEN 2
-            END
-          WHERE state_id IS NULL
-          AND id BETWEEN #{start_id} AND #{end_id}
-        SQL
+        reschedule_if_needed([start_id, end_id]) do
+          ActiveRecord::Base.connection.execute <<~SQL
+            UPDATE issues
+            SET state_id =
+              CASE state
+              WHEN 'opened' THEN 1
+              WHEN 'closed' THEN 2
+              END
+            WHERE state_id IS NULL
+            AND id BETWEEN #{start_id} AND #{end_id}
+          SQL
+        end
+      end
+
+      private
+
+      def need_reschedule?
+        wait_for_deadtuple_vacuum?('issues')
       end
     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 923ceaeec54..5d92553f577 100644
--- a/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
+++ b/lib/gitlab/background_migration/sync_merge_requests_state_id.rb
@@ -4,21 +4,31 @@
 module Gitlab
   module BackgroundMigration
     class SyncMergeRequestsStateId
+      include Helpers::Reschedulable
+
       def perform(start_id, end_id)
         Rails.logger.info("Merge Requests - Populating state_id: #{start_id} - #{end_id}")
 
-        ActiveRecord::Base.connection.execute <<~SQL
-          UPDATE merge_requests
-          SET state_id =
-            CASE state
-            WHEN 'opened' THEN 1
-            WHEN 'closed' THEN 2
-            WHEN 'merged' THEN 3
-            WHEN 'locked' THEN 4
-            END
-          WHERE state_id IS NULL
-          AND id BETWEEN #{start_id} AND #{end_id}
-        SQL
+        reschedule_if_needed([start_id, end_id]) do
+          ActiveRecord::Base.connection.execute <<~SQL
+            UPDATE merge_requests
+            SET state_id =
+              CASE state
+              WHEN 'opened' THEN 1
+              WHEN 'closed' THEN 2
+              WHEN 'merged' THEN 3
+              WHEN 'locked' THEN 4
+              END
+            WHERE state_id IS NULL
+            AND id BETWEEN #{start_id} AND #{end_id}
+          SQL
+        end
+      end
+
+      private
+
+      def need_reschedule?
+        wait_for_deadtuple_vacuum?('issues')
       end
     end
   end