Add more specs and code improvements

6 files changed, 89 insertions, 56 deletions

diff --git a/db/post_migrate/20190214112022_schedule_sync_issuables_state_id.rb b/db/post_migrate/20190214112022_schedule_sync_issuables_state_id.rb
index d9b77d2f02c..898bad043ac 100644
--- a/db/post_migrate/20190214112022_schedule_sync_issuables_state_id.rb
+++ b/db/post_migrate/20190214112022_schedule_sync_issuables_state_id.rb
@@ -1,22 +1,22 @@
-# frozen_string_literal: true
-
-# See http://doc.gitlab.com/ce/development/migration_style_guide.html
-# for more information on how to write migrations for GitLab.
-
 class ScheduleSyncIssuablesStateId < ActiveRecord::Migration[5.0]
+  # This migration schedules the sync of state_id for issues and merge requests
+  # which are converting the state column from string to integer.
+  # For more information check: https://gitlab.com/gitlab-org/gitlab-ce/issues/51789
+
   include Gitlab::Database::MigrationHelpers
 
   DOWNTIME = false
 
-  # 2019-02-12 Gitlab.com issuable numbers
+  # 2019-02-12 gitlab.com issuable numbers
   # issues count: 13587305
   # merge requests count: 18925274
-  # Using this 25000 as batch size should take around 26 hours
+  #
+  # Using 25000 as batch size should take around 26 hours
   # to migrate both issues and merge requests
   BATCH_SIZE = 25000
   DELAY_INTERVAL = 5.minutes.to_i
-  ISSUE_MIGRATION = 'SyncIssuesStateId'.freeze
-  MERGE_REQUEST_MIGRATION = 'SyncMergeRequestsStateId'.freeze
+  ISSUES_MIGRATION = 'SyncIssuesStateId'.freeze
+  MERGE_REQUESTS_MIGRATION = 'SyncMergeRequestsStateId'.freeze
 
   class Issue < ActiveRecord::Base
     include EachBatch
@@ -32,8 +32,17 @@ class ScheduleSyncIssuablesStateId < ActiveRecord::Migration[5.0]
 
   def up
     Sidekiq::Worker.skipping_transaction_check do
-      queue_background_migration_jobs_by_range_at_intervals(Issue.where(state_id: nil), ISSUE_MIGRATION, DELAY_INTERVAL, batch_size: BATCH_SIZE)
-      queue_background_migration_jobs_by_range_at_intervals(MergeRequest.where(state_id: nil), MERGE_REQUEST_MIGRATION, DELAY_INTERVAL, batch_size: BATCH_SIZE)
+      queue_background_migration_jobs_by_range_at_intervals(Issue.where(state_id: nil),
+                                                            ISSUES_MIGRATION,
+                                                            DELAY_INTERVAL,
+                                                            batch_size: BATCH_SIZE
+                                                           )
+
+      queue_background_migration_jobs_by_range_at_intervals(MergeRequest.where(state_id: nil),
+                                                            MERGE_REQUESTS_MIGRATION,
+                                                            DELAY_INTERVAL,
+                                                            batch_size: BATCH_SIZE
+                                                           )
     end
   end
 
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
diff --git a/spec/migrations/schedule_sync_issuables_state_id_spec.rb b/spec/migrations/schedule_sync_issuables_state_id_spec.rb
index a746fb68c30..6c4c31ae554 100644
--- a/spec/migrations/schedule_sync_issuables_state_id_spec.rb
+++ b/spec/migrations/schedule_sync_issuables_state_id_spec.rb
@@ -15,6 +15,25 @@ describe ScheduleSyncIssuablesStateId, :migration do
     @project = projects.create!(namespace_id: @group.id)
   end
 
+  shared_examples 'scheduling migrations' do
+    before do
+      Sidekiq::Worker.clear_all
+      stub_const("#{described_class.name}::BATCH_SIZE", 2)
+    end
+
+    it 'correctly schedules issuable sync background migration' do
+      Sidekiq::Testing.fake! do
+        Timecop.freeze do
+          migrate!
+
+          expect(migration).to be_scheduled_delayed_migration(5.minutes, resource_1.id, resource_2.id)
+          expect(migration).to be_scheduled_delayed_migration(10.minutes, resource_3.id, resource_4.id)
+          expect(BackgroundMigrationWorker.jobs.size).to eq(2)
+        end
+      end
+    end
+  end
+
   shared_examples 'rescheduling migrations' do
     before do
       Sidekiq::Worker.clear_all
@@ -35,20 +54,6 @@ describe ScheduleSyncIssuablesStateId, :migration do
   end
 
   describe '#up' do
-    # it 'correctly schedules diff file deletion workers' do
-    #   Sidekiq::Testing.fake! do
-    #     Timecop.freeze do
-    #       described_class.new.perform
-
-    #       expect(described_class::MIGRATION).to be_scheduled_delayed_migration(5.minutes, [1, 4, 5])
-
-    #       expect(described_class::MIGRATION).to be_scheduled_delayed_migration(10.minutes, [6])
-
-    #       expect(BackgroundMigrationWorker.jobs.size).to eq(2)
-    #     end
-    #   end
-    # end
-
     context 'issues' do
       it 'migrates state column to integer' do
         opened_issue = issues.create!(description: 'first', state: 'opened')
@@ -64,10 +69,18 @@ describe ScheduleSyncIssuablesStateId, :migration do
         expect(nil_state_issue.reload.state_id).to be_nil
       end
 
+      it_behaves_like 'scheduling migrations' do
+        let(:migration) { described_class::ISSUES_MIGRATION }
+        let!(:resource_1) { issues.create!(description: 'first', state: 'opened') }
+        let!(:resource_2) { issues.create!(description: 'second', state: 'closed') }
+        let!(:resource_3) { issues.create!(description: 'third', state: 'closed') }
+        let!(:resource_4) { issues.create!(description: 'fourth', state: 'closed') }
+      end
+
       it_behaves_like 'rescheduling migrations' do
         let(:worker_class) { Gitlab::BackgroundMigration::SyncIssuesStateId }
-        let(:resource_1) { issues.create!(description: 'first', state: 'opened') }
-        let(:resource_2) { issues.create!(description: 'second', state: 'closed') }
+        let!(:resource_1) { issues.create!(description: 'first', state: 'opened') }
+        let!(:resource_2) { issues.create!(description: 'second', state: 'closed') }
       end
     end
 
@@ -88,6 +101,14 @@ describe ScheduleSyncIssuablesStateId, :migration do
         expect(invalid_state_merge_request.reload.state_id).to be_nil
       end
 
+      it_behaves_like 'scheduling migrations' do
+        let(:migration) { described_class::MERGE_REQUESTS_MIGRATION }
+        let!(:resource_1) { merge_requests.create!(state: 'opened', target_project_id: @project.id, target_branch: 'feature1', source_branch: 'master') }
+        let!(:resource_2) { merge_requests.create!(state: 'closed', target_project_id: @project.id, target_branch: 'feature2', source_branch: 'master') }
+        let!(:resource_3) { merge_requests.create!(state: 'merged', target_project_id: @project.id, target_branch: 'feature3', source_branch: 'master') }
+        let!(:resource_4) { merge_requests.create!(state: 'locked', target_project_id: @project.id, target_branch: 'feature4', source_branch: 'master') }
+      end
+
       it_behaves_like 'rescheduling migrations' do
         let(:worker_class) { Gitlab::BackgroundMigration::SyncMergeRequestsStateId }
         let(:resource_1) { merge_requests.create!(state: 'opened', target_project_id: @project.id, target_branch: 'feature1', source_branch: 'master') }