Migrate pipeline stages in batches instead of single row

lib/gitlab/background_migration/migrate_build_stage.rb

@@ -5,69 +5,45 @@
 module Gitlab
   module BackgroundMigration
     class MigrateBuildStage
-      def perform(id)
-        DatabaseBuild.find_by(id: id).try do |build|
-          MigratableStage.new(build).tap do |stage|
-            break if stage.exists? || stage.legacy?
-
-            stage.ensure!
-            stage.migrate_reference!
-            stage.migrate_status!
-          end
-        end
-      end
-
-      class DatabaseStage < ActiveRecord::Base
-        self.table_name = 'ci_stages'
-      end
-
-      class DatabaseBuild < ActiveRecord::Base
-        self.table_name = 'ci_builds'
-      end
-
-      class MigratableStage
-        def initialize(build)
-          @build = build
-        end
-
-        def exists?
-          @build.reload.stage_id.present?
+      module Migratable
+        class Stage < ActiveRecord::Base
+          self.table_name = 'ci_stages'
         end
 
-        ##
-        # We can have some very old stages that do not have `ci_builds.stage` set.
-        #
-        # In that case we just don't migrate such stage.
-        #
-        def legacy?
-          @build.stage.nil?
-        end
+        class Build < ActiveRecord::Base
+          self.table_name = 'ci_builds'
 
-        def ensure!
-          find || create!
-        end
+          def ensure_stage!
+            find || create!
+          rescue ActiveRecord::RecordNotUnique
+            # TODO
+          end
 
-        def find
-          DatabaseStage.find_by(name: @build.stage,
-                                pipeline_id: @build.commit_id,
-                                project_id: @build.project_id)
-        end
+          def find
+            Stage.find_by(name: self.stage,
+                          pipeline_id: self.commit_id,
+                          project_id: self.project_id)
+          end
 
-        def create!
-          DatabaseStage.create!(name: @build.stage,
-                                pipeline_id: @build.commit_id,
-                                project_id: @build.project_id)
+          def create!
+            Stage.create!(name: self.stage || 'test',
+                          pipeline_id: self.commit_id,
+                          project_id: self.project_id)
+          end
         end
+      end
 
-        def migrate_reference!
-          MigrateBuildStageIdReference.new.perform(@build.id, @build.id)
-        end
+      def perform(start_id, stop_id)
+        # TODO, should we disable_statement_timeout?
+        # TODO, use plain SQL query?
 
-        def migrate_status!
-          raise ArgumentError unless exists?
+        stages = Migratable::Build.where('stage_id IS NULL')
+          .where("id BETWEEN #{start_id.to_i} AND #{stop_id.to_i}")
+          .map { |build| build.ensure_stage! }
+          .compact.map(&:id)
 
-          MigrateStageStatus.new.perform(@build.stage_id, @build.stage_id)
-        end
+        MigrateBuildStageIdReference.new.perform(start_id, stop_id)
+        MigrateStageStatus.new.perform(stages.min, stages.max)
       end
     end
   end

spec/lib/gitlab/background_migration/migrate_build_stage_spec.rb

 require 'spec_helper'
 
-describe Gitlab::BackgroundMigration::MigrateBuildStage, :migration, schema: 20180105101928 do
+describe Gitlab::BackgroundMigration::MigrateBuildStage, :migration, schema: 20180212101928 do
   let(:projects) { table(:projects) }
   let(:pipelines) { table(:ci_pipelines) }
   let(:stages) { table(:ci_stages) }
@@ -10,15 +10,9 @@ describe Gitlab::BackgroundMigration::MigrateBuildStage, :migration, schema: 201
                failed: 4, canceled: 5, skipped: 6, manual: 7 }.freeze
 
   before do
-    ##
-    # Dependencies
-    #
     projects.create!(id: 123, name: 'gitlab', path: 'gitlab-ce')
     pipelines.create!(id: 1, project_id: 123, ref: 'master', sha: 'adf43c3a')
 
-    ##
-    # CI/CD jobs
-    #
     jobs.create!(id: 1, commit_id: 1, project_id: 123,
                  stage_idx: 2, stage: 'build', status: :success)
     jobs.create!(id: 2, commit_id: 1, project_id: 123,
@@ -36,9 +30,7 @@ describe Gitlab::BackgroundMigration::MigrateBuildStage, :migration, schema: 201
   it 'correctly migrates builds stages' do
     expect(stages.count).to be_zero
 
-    jobs.all.find_each do |job|
-      described_class.new.perform(job.id)
-    end
+    described_class.new.perform(1, 6)
 
     expect(stages.count).to eq 3
     expect(stages.all.pluck(:name)).to match_array %w[test build deploy]