Class: Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy

Inherits:
BaseStrategy
  • Object
show all
Includes:
Database::DynamicModelHelpers
Defined in:
lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb

Overview

Generic batching class for use with a BatchedBackgroundMigration. Batches over the given table and column combination, returning the MIN() and MAX() values for the next batch as an array.

If no more batches exist in the table, returns nil.

Constant Summary

Constants included from Database::DynamicModelHelpers

Database::DynamicModelHelpers::BATCH_SIZE

Instance Method Summary collapse

Methods included from Database::DynamicModelHelpers

define_batchable_model, #each_batch, #each_batch_range

Methods inherited from BaseStrategy

#initialize

Constructor Details

This class inherits a constructor from Gitlab::BackgroundMigration::BatchingStrategies::BaseStrategy

Instance Method Details

#next_batch(table_name, column_name, batch_min_value:, batch_size:, job_arguments:, job_class: nil) ⇒ Object

Finds and returns the next batch in the table.

table_name - The table to batch over column_name - The column to batch over batch_min_value - The minimum value which the next batch will start at batch_size - The size of the next batch job_arguments - The migration job arguments job_class - The migration job class rubocop:disable Metrics/AbcSize – temporarily contains two branches for cursor and non-cursor batching



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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
 
# File 'lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb', line 23

def next_batch(table_name, column_name, batch_min_value:, batch_size:, job_arguments:, job_class: nil)
  base_class = Gitlab::Database.application_record_for_connection(connection)
  model_class = define_batchable_model(table_name, connection: connection, base_class: base_class)
  next_batch_bounds = nil

  # rubocop:disable Lint/UnreachableLoop -- we need to use each_batch to pull one batch out
  if job_class.cursor?
    cursor_columns = job_class.cursor_columns

    Gitlab::Pagination::Keyset::Iterator.new(
      scope: model_class.order(cursor_columns),
      cursor: cursor_columns.zip(batch_min_value).to_h
    ).each_batch(of: batch_size, load_batch: false) do |batch|
      break unless batch.first && batch.last # skip if the batch is empty for some reason

      next_batch_bounds = [batch.first.values_at(cursor_columns), batch.last.values_at(cursor_columns)]
      break
    end
  else
    arel_column = model_class.arel_table[column_name]
    relation = model_class.where(arel_column.gteq(batch_min_value))
    reset_order = true

    if job_class
      relation = filter_batch(relation,
        table_name: table_name, column_name: column_name,
        job_class: job_class, job_arguments: job_arguments
      )
      reset_order = job_class.reset_order if job_class.respond_to?(:reset_order)
    end

    relation.each_batch(of: batch_size, column: column_name, reset_order: reset_order) do |batch|
      next_batch_bounds = batch.pick(arel_column.minimum, arel_column.maximum)

      break
    end
  end
  # rubocop:enable Lint/UnreachableLoop

  next_batch_bounds
end