diff options
Diffstat (limited to 'lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb')
-rw-r--r-- | lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb | 32 |
1 files changed, 27 insertions, 5 deletions
diff --git a/lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb b/lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb index c2f59bf9c76..1ffa4a052e5 100644 --- a/lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb +++ b/lib/gitlab/background_migration/batching_strategies/primary_key_batching_strategy.rb @@ -24,11 +24,19 @@ module Gitlab quoted_column_name = model_class.connection.quote_column_name(column_name) relation = model_class.where("#{quoted_column_name} >= ?", batch_min_value) + + if job_class + relation = filter_batch(relation, + table_name: table_name, column_name: column_name, + job_class: job_class, job_arguments: job_arguments + ) + end + relation = apply_additional_filters(relation, job_arguments: job_arguments, job_class: job_class) next_batch_bounds = nil relation.each_batch(of: batch_size, column: column_name) do |batch| # rubocop:disable Lint/UnreachableLoop - next_batch_bounds = batch.pluck(Arel.sql("MIN(#{quoted_column_name}), MAX(#{quoted_column_name})")).first + next_batch_bounds = batch.pick(Arel.sql("MIN(#{quoted_column_name}), MAX(#{quoted_column_name})")) break end @@ -36,13 +44,27 @@ module Gitlab next_batch_bounds end + # Deprecated + # + # Use `scope_to` to define additional filters on the migration job class. + # + # see https://docs.gitlab.com/ee/development/database/batched_background_migrations.html#adding-additional-filters. def apply_additional_filters(relation, job_arguments: [], job_class: nil) - if job_class.respond_to?(:batching_scope) - return job_class.batching_scope(relation, job_arguments: job_arguments) - end - relation end + + private + + def filter_batch(relation, table_name:, column_name:, job_class:, job_arguments: []) + return relation unless job_class.respond_to?(:generic_instance) + + job = job_class.generic_instance( + batch_table: table_name, batch_column: column_name, + job_arguments: job_arguments, connection: connection + ) + + job.filter_batch(relation) + end end end end |