diff options
| -rw-r--r-- | lib/gitlab/database/migration_helpers.rb | 75 | ||||
| -rw-r--r-- | spec/lib/gitlab/database/migration_helpers_spec.rb | 53 |
2 files changed, 97 insertions, 31 deletions
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index 0643c56db9b..29e812c3c78 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -221,17 +221,19 @@ module Gitlab # make things _more_ complex). # # rubocop: disable Metrics/AbcSize - def update_column_in_batches(table, column, value) + def update_column_in_batches(table, column, value, &scope) if transaction_open? - raise 'update_column_in_batches can not be run inside a transaction, ' \ - 'you can disable transactions by calling disable_ddl_transaction! ' \ - 'in the body of your migration class' + raise <<-MSG + update_column_in_batches helper can not be run inside a transaction. + You can disable transactions by calling `disable_ddl_transaction!` + method in the body of your migration class. + MSG end - table = Arel::Table.new(table) + table_arel = Arel::Table.new(table) - count_arel = table.project(Arel.star.count.as('count')) - count_arel = yield table, count_arel if block_given? + count_arel = table_arel.project(Arel.star.count.as('count')) + count_arel = yield table_arel, count_arel if block_given? total = exec_query(count_arel.to_sql).to_hash.first['count'].to_i @@ -246,37 +248,56 @@ module Gitlab # rows for GitLab.com. batch_size = max_size if batch_size > max_size + walk_table_in_batches(table, of: batch_size, scope: scope) do + Arel::UpdateManager.new(ActiveRecord::Base) + .table(table_arel) + .set([[table_arel[column], value]]) + end + end + + def walk_table_in_batches(table, of: 1000, scope: nil) + if transaction_open? + raise <<-MSG + walk_table_in_batches helper can not be run inside a transaction. + You can disable transactions by calling `disable_ddl_transaction!` + method in the body of your migration class. + MSG + end + + table = Arel::Table.new(table) + start_arel = table.project(table[:id]).order(table[:id].asc).take(1) - start_arel = yield table, start_arel if block_given? - start_id = exec_query(start_arel.to_sql).to_hash.first['id'].to_i + start_arel = scope.call(table, start_arel) if scope + start_id = exec_query(start_arel.to_sql).to_hash.first.to_h['id'].to_i - loop do + 1.step do |batch| stop_arel = table.project(table[:id]) .where(table[:id].gteq(start_id)) .order(table[:id].asc) .take(1) - .skip(batch_size) - - stop_arel = yield table, stop_arel if block_given? - stop_row = exec_query(stop_arel.to_sql).to_hash.first + .skip(of) - update_arel = Arel::UpdateManager.new(ActiveRecord::Base) - .table(table) - .set([[table[column], value]]) - .where(table[:id].gteq(start_id)) + stop_arel = scope.call(table, stop_arel) if scope + stop_id = exec_query(stop_arel.to_sql) + .to_hash.first.to_h['id'].to_i - if stop_row - stop_id = stop_row['id'].to_i - start_id = stop_id - update_arel = update_arel.where(table[:id].lt(stop_id)) - end + action = yield(batch, start_id, stop_id) - update_arel = yield table, update_arel if block_given? + if action.is_a?(Arel::TreeManager) + exec_arel = action.where(table[:id].gteq(start_id)) + exec_arel = exec_arel.where(table[:id].lt(stop_id)) if stop_id.nonzero? + exec_arel = scope.call(table, exec_arel) if scope - execute(update_arel.to_sql) + execute(exec_arel.to_sql) + end - # There are no more rows left to update. - break unless stop_row + if stop_id.zero? + # there are no more rows left to update + break + else + # next loop + start_id = stop_id + end end end diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index 4259be3f522..65af2e1cb52 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -2,9 +2,7 @@ require 'spec_helper' describe Gitlab::Database::MigrationHelpers, lib: true do let(:model) do - ActiveRecord::Migration.new.extend( - Gitlab::Database::MigrationHelpers - ) + ActiveRecord::Migration.new.extend(described_class) end before do @@ -264,7 +262,7 @@ describe Gitlab::Database::MigrationHelpers, lib: true do describe '#update_column_in_batches' do context 'when running outside of a transaction' do before do - expect(model).to receive(:transaction_open?).and_return(false) + expect(model).to receive(:transaction_open?).twice.and_return(false) create_list(:empty_project, 5) end @@ -313,6 +311,53 @@ describe Gitlab::Database::MigrationHelpers, lib: true do end end + describe '#walk_table_in_batches' do + context 'when running outside of a transaction' do + before do + expect(model).to receive(:transaction_open?).and_return(false) + + create_list(:empty_project, 6) + end + + it 'yields for each batch' do + expect { |b| model.walk_table_in_batches(:projects, of: 2, &b) } + .to yield_control.exactly(3).times + end + + it 'yields successive ranges' do + expect { |b| model.walk_table_in_batches(:projects, of: 2, &b) } + .to yield_successive_args([1, Integer, Integer], + [2, Integer, Integer], + [3, Integer, 0]) + end + + context 'when a scope is provided' do + it 'limits the scope of the statement provided inside the block' do + first_id = Project.first.id + scope = ->(table, query) { query.where(table[:id].eq(first_id)) } + + model.walk_table_in_batches(:projects, scope: scope) do + Arel::UpdateManager.new(ActiveRecord::Base) + .table(Arel::Table.new(:projects)) + .set([[Arel::Table.new(:projects)[:archived], true]]) + end + + expect(Project.where(archived: true).count).to eq(1) + end + end + end + + context 'when running inside the transaction' do + it 'raises RuntimeError' do + expect(model).to receive(:transaction_open?).and_return(true) + + expect do + model.update_column_in_batches(:projects, :star_count, Arel.sql('1+1')) + end.to raise_error(RuntimeError) + end + end + end + describe '#add_column_with_default' do context 'outside of a transaction' do context 'when a column limit is not set' do |