lib/gitlab/database/migrations/test_batched_background_runner.rb


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

# frozen_string_literal: true

module Gitlab
  module Database
    module Migrations
      class TestBatchedBackgroundRunner < BaseBackgroundRunner
        attr_reader :connection

        def initialize(result_dir:, connection:)
          super(result_dir: result_dir)
          @connection = connection
        end

        def jobs_by_migration_name
          Gitlab::Database::BackgroundMigration::BatchedMigration
            .executable
            .created_after(2.days.ago) # Simple way to exclude migrations already running before migration testing
            .to_h do |migration|
            batching_strategy = migration.batch_class.new(connection: connection)

            all_migration_jobs = []

            min_value = migration.next_min_value

            while (next_bounds = batching_strategy.next_batch(
              migration.table_name,
              migration.column_name,
              batch_min_value: min_value,
              batch_size: migration.batch_size,
              job_arguments: migration.job_arguments
            ))

              batch_min, batch_max = next_bounds

              all_migration_jobs << migration.create_batched_job!(batch_min, batch_max)
              min_value = batch_max + 1
            end

            [migration.job_class_name, all_migration_jobs]
          end
        end

        def run_job(job)
          Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection).perform(job)
        end
      end
    end
  end
end