BigW Consortium Gitlab

Skip to content

G
gitlab-ce
Commit af0eeefc by Grzegorz Bizon
Options

Revert recent changes in migration helpers

parent 84265775
Showing with 30 additions and 172 deletions
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
......
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
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment