diff options
| -rw-r--r-- | lib/gitlab/database/migration_helpers.rb | 120 | ||||
| -rw-r--r-- | spec/lib/gitlab/database/migration_helpers_spec.rb | 82 |
2 files changed, 30 insertions, 172 deletions
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index ca7e4c8aa7c..0643c56db9b 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -221,19 +221,17 @@ module Gitlab # make things _more_ complex). # # rubocop: disable Metrics/AbcSize - def update_column_in_batches(table, column, value, &scope) + def update_column_in_batches(table, column, value) if transaction_open? - 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 + 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' end - table_arel = Arel::Table.new(table) + table = Arel::Table.new(table) - count_arel = table_arel.project(Arel.star.count.as('count')) - count_arel = yield table_arel, count_arel if block_given? + count_arel = table.project(Arel.star.count.as('count')) + count_arel = yield table, count_arel if block_given? total = exec_query(count_arel.to_sql).to_hash.first['count'].to_i @@ -248,103 +246,37 @@ module Gitlab # rows for GitLab.com. batch_size = max_size if batch_size > max_size - execute_in_batches(table, of: batch_size, scope: scope) do - Arel::UpdateManager.new(ActiveRecord::Base) - .table(table_arel) - .set([[table_arel[column], value]]) - end - end - - ## - # Iterates a table and executes a block for given range. - # - # Yields batch index, start and stop ids. - # - # Optional `scope` keyword argument is a closure that is meant to limit - # the scope the statement is going to be applied onto. - # - # Arel statement this helper will execute must be defined inside the - # block. - # - # Example: - # - # scope = ->(table, query) { query.where(table[:id].gt(100) } - # - # walk_table_in_batches(:table, of: 10, scope: scope) do |index, start, stop| - # # do something here - # 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 = scope.call(table, start_arel) if scope - start_id = exec_query(start_arel.to_sql).to_hash.first.to_h['id'].to_i + start_arel = yield table, start_arel if block_given? + start_id = exec_query(start_arel.to_sql).to_hash.first['id'].to_i - 1.step do |batch| + loop do stop_arel = table.project(table[:id]) .where(table[:id].gteq(start_id)) .order(table[:id].asc) .take(1) - .skip(of) - - 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 + .skip(batch_size) - yield batch, start_id, stop_id + stop_arel = yield table, stop_arel if block_given? + stop_row = exec_query(stop_arel.to_sql).to_hash.first - stop_id.zero? ? break : start_id = stop_id - end - end - - ## - # Executes an SQL statement in batches, created by Arel manager. - # - # Optional `scope` keyword argument is a closure that is meant to limit - # the scope the statement is going to be applied onto. - # - # Arel statement this helper will execute must be defined inside the - # block. - # - # Example: - # - # scope = ->(table, query) { query.where(table[:id].gt(100) } - # - # execute_in_batches(:table, of: 10000, scope: scope) do |table| - # Arel::UpdateManager.new(ActiveRecord::Base) - # .table(table) - # .set([[table[:field], 101]]) - # end - # - def execute_in_batches(table, of: 1000, scope: nil) - if transaction_open? - raise <<-MSG - execute_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 + update_arel = Arel::UpdateManager.new(ActiveRecord::Base) + .table(table) + .set([[table[column], value]]) + .where(table[:id].gteq(start_id)) - raise ArgumentError, 'This method requires a block!' unless block_given? + 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 - table_arel = Arel::Table.new(table) + update_arel = yield table, update_arel if block_given? - walk_table_in_batches(table, of: of, scope: scope) do |_batch, start_id, stop_id| - exec_arel = yield table_arel - exec_arel = exec_arel.where(table_arel[:id].gteq(start_id)) - exec_arel = exec_arel.where(table_arel[:id].lt(stop_id)) if stop_id.nonzero? - exec_arel = scope.call(table_arel, exec_arel) if scope + execute(update_arel.to_sql) - execute(exec_arel.to_sql) + # There are no more rows left to update. + break unless stop_row end end diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index f4a66b7e2a2..4259be3f522 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -2,7 +2,9 @@ require 'spec_helper' describe Gitlab::Database::MigrationHelpers, lib: true do let(:model) do - ActiveRecord::Migration.new.extend(described_class) + ActiveRecord::Migration.new.extend( + Gitlab::Database::MigrationHelpers + ) end before do @@ -262,8 +264,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?) - .at_least(:once).and_return(false) + expect(model).to receive(:transaction_open?).and_return(false) create_list(:empty_project, 5) end @@ -312,81 +313,6 @@ 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)) } - - expect { |b| model.walk_table_in_batches(:projects, of: 1, scope: scope, &b) } - .to yield_control.exactly(:once) - end - end - end - - context 'when running inside the transaction' do - it 'raises RuntimeError' do - expect(model).to receive(:transaction_open?).and_return(true) - - expect { model.walk_table_in_batches(:projects, of: 2) } - .to raise_error(RuntimeError) - end - end - end - - describe '#execute_in_batches' do - context 'when running outside of a transaction' do - before do - expect(model).to receive(:transaction_open?) - .at_least(:once).and_return(false) - - create_list(:empty_project, 6) - 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.execute_in_batches(:projects, scope: scope) do |table| - Arel::UpdateManager.new(ActiveRecord::Base) - .table(table).set([[table[: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 { model.execute_in_batches(:projects)} - .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 |