Commit 2917f486 authored by Grzegorz Bizon's avatar Grzegorz Bizon

Extract `execute_in_batches` migrations helper

parent e036a374
...@@ -248,7 +248,7 @@ module Gitlab ...@@ -248,7 +248,7 @@ module Gitlab
# rows for GitLab.com. # rows for GitLab.com.
batch_size = max_size if batch_size > max_size batch_size = max_size if batch_size > max_size
walk_table_in_batches(table, of: batch_size, scope: scope) do execute_in_batches(table, of: batch_size, scope: scope) do
Arel::UpdateManager.new(ActiveRecord::Base) Arel::UpdateManager.new(ActiveRecord::Base)
.table(table_arel) .table(table_arel)
.set([[table_arel[column], value]]) .set([[table_arel[column], value]])
...@@ -281,23 +281,32 @@ module Gitlab ...@@ -281,23 +281,32 @@ module Gitlab
stop_id = exec_query(stop_arel.to_sql) stop_id = exec_query(stop_arel.to_sql)
.to_hash.first.to_h['id'].to_i .to_hash.first.to_h['id'].to_i
action = yield(batch, start_id, stop_id) yield batch, start_id, stop_id
if action.is_a?(Arel::TreeManager) stop_id.zero? ? break : start_id = stop_id
exec_arel = action.where(table[:id].gteq(start_id)) end
exec_arel = exec_arel.where(table[:id].lt(stop_id)) if stop_id.nonzero? end
exec_arel = scope.call(table, exec_arel) if scope
execute(exec_arel.to_sql) def execute_in_batches(table, of: 1000, scope: nil)
end 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
if stop_id.zero? # raise 'This method requires a block!' unless block_given?
# there are no more rows left to update
break table_arel = Arel::Table.new(table)
else
# next loop walk_table_in_batches(table, of: of, scope: scope) do |_batch, start_id, stop_id|
start_id = stop_id exec_arel = yield table_arel
end 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(exec_arel.to_sql)
end end
end end
......
...@@ -262,7 +262,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do ...@@ -262,7 +262,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
describe '#update_column_in_batches' do describe '#update_column_in_batches' do
context 'when running outside of a transaction' do context 'when running outside of a transaction' do
before do before do
expect(model).to receive(:transaction_open?).twice.and_return(false) expect(model).to receive(:transaction_open?)
.at_least(:once).and_return(false)
create_list(:empty_project, 5) create_list(:empty_project, 5)
end end
...@@ -336,10 +337,39 @@ describe Gitlab::Database::MigrationHelpers, lib: true do ...@@ -336,10 +337,39 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
first_id = Project.first.id first_id = Project.first.id
scope = ->(table, query) { query.where(table[:id].eq(first_id)) } scope = ->(table, query) { query.where(table[:id].eq(first_id)) }
model.walk_table_in_batches(:projects, scope: scope) do 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) Arel::UpdateManager.new(ActiveRecord::Base)
.table(Arel::Table.new(:projects)) .table(table).set([[table[:archived], true]])
.set([[Arel::Table.new(:projects)[:archived], true]])
end end
expect(Project.where(archived: true).count).to eq(1) expect(Project.where(archived: true).count).to eq(1)
...@@ -351,9 +381,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do ...@@ -351,9 +381,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
it 'raises RuntimeError' do it 'raises RuntimeError' do
expect(model).to receive(:transaction_open?).and_return(true) expect(model).to receive(:transaction_open?).and_return(true)
expect do expect { model.execute_in_batches(:projects)}
model.update_column_in_batches(:projects, :star_count, Arel.sql('1+1')) .to raise_error(RuntimeError)
end.to raise_error(RuntimeError)
end end
end end
end end
......
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