kotovalexarian-likes-gitlab/gitlab-org--gitlab-foss
1
0
Fork 0
Code Releases Activity

MigrationHelper disable_statement_timeout accepts transaction: false

Browse source 
By default statement_timeout will only be enabled during transaction
lifetime, therefore not leaking outside of it.

With `transaction: false` it will set for entire session, but requires
a block to passed. It yields control and cleans up session after block
finishes, also preventing leaking outside of it.
This commit is contained in:
Gabriel Mazetto 2018-07-09 17:34:50 +02:00
parent a3c2b39d10
commit 09e7c75d1b
2 changed files with 167 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
lib/gitlab/database/migration_helpers.rb
View file

@ -58,7 +58,6 @@ module Gitlab
if Database.postgresql?
options = options.merge({ algorithm: :concurrently })
disable_statement_timeout
end
if index_exists?(table_name, column_name, options)
@ -66,7 +65,9 @@ module Gitlab
return
end
add_index(table_name, column_name, options)
disable_statement_timeout(transaction: false) do
add_index(table_name, column_name, options)
end
end
# Removes an existed index, concurrently when supported
@ -87,7 +88,6 @@ module Gitlab
if supports_drop_index_concurrently?
options = options.merge({ algorithm: :concurrently })
disable_statement_timeout
end
unless index_exists?(table_name, column_name, options)
@ -95,7 +95,9 @@ module Gitlab
return
end
remove_index(table_name, options.merge({ column: column_name }))
disable_statement_timeout(transaction: false) do
remove_index(table_name, options.merge({ column: column_name }))
end
end
# Removes an existing index, concurrently when supported
@ -116,7 +118,6 @@ module Gitlab
if supports_drop_index_concurrently?
options = options.merge({ algorithm: :concurrently })
disable_statement_timeout
end
unless index_exists_by_name?(table_name, index_name)
@ -124,7 +125,9 @@ module Gitlab
return
end
remove_index(table_name, options.merge({ name: index_name }))
disable_statement_timeout(transaction: false) do
remove_index(table_name, options.merge({ name: index_name }))
end
end
# Only available on Postgresql >= 9.2
@ -171,8 +174,6 @@ module Gitlab
on_delete = 'SET NULL' if on_delete == :nullify
end
disable_statement_timeout
key_name = concurrent_foreign_key_name(source, column)
unless foreign_key_exists?(source, target, column: column)
@ -199,7 +200,9 @@ module Gitlab
# while running.
#
# Note this is a no-op in case the constraint is VALID already
execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
disable_statement_timeout(transaction: false) do
execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{key_name};")
end
end
def foreign_key_exists?(source, target = nil, column: nil)
@ -224,8 +227,49 @@ module Gitlab
# Long-running migrations may take more than the timeout allowed by
# the database. Disable the session's statement timeout to ensure
# migrations don't get killed prematurely. (PostgreSQL only)
def disable_statement_timeout
execute('SET statement_timeout TO 0') if Database.postgresql?
#
# There are two possible ways to disable the statement timeout:
#
# - Per transaction (this is the preferred and default mode)
# - Per connection (requires a cleanup after the execution)
#
# When using a per connection disable statement, code must be inside a block
# so we can automatically `RESET ALL` after it has executed otherwise the statement
# will still be disabled until connection is dropped or `RESET ALL` is executed
#
# - +transaction:+ true to disable for current transaction only *(default)*
# - +transaction:+ false to disable for current session (requires block)
def disable_statement_timeout(transaction: true)
# bypass disabled_statement logic when not using postgres, but still execute block when one is given
unless Database.postgresql?
if block_given?
yield
end
return
end
if transaction
unless transaction_open?
raise 'disable_statement_timeout() cannot be run without a transaction, ' \
'use it inside a transaction block. Alternatively you can use: ' \
'disable_statement_timeout(transaction: false) { #code here } to make sure ' \
'statement_timeout is reset after the block execution is finished.'
end
execute('SET LOCAL statement_timeout TO 0')
else
unless block_given?
raise ArgumentError, 'disable_statement_timeout(transaction: false) requires a block encapsulating' \
'code that will be executed with the statement_timeout disabled.'
end
execute('SET statement_timeout TO 0')
yield
execute('RESET ALL')
end
end
def true_value
@ -367,30 +411,30 @@ module Gitlab
'in the body of your migration class'
end
disable_statement_timeout
disable_statement_timeout(transaction: false) do
transaction do
if limit
add_column(table, column, type, default: nil, limit: limit)
else
add_column(table, column, type, default: nil)
end
transaction do
if limit
add_column(table, column, type, default: nil, limit: limit)
else
add_column(table, column, type, default: nil)
# Changing the default before the update ensures any newly inserted
# rows already use the proper default value.
change_column_default(table, column, default)
end
# Changing the default before the update ensures any newly inserted
# rows already use the proper default value.
change_column_default(table, column, default)
end
begin
update_column_in_batches(table, column, default, &block)
begin
update_column_in_batches(table, column, default, &block)
change_column_null(table, column, false) unless allow_null
# We want to rescue _all_ exceptions here, even those that don't inherit
# from StandardError.
rescue Exception => error # rubocop: disable all
remove_column(table, column)
change_column_null(table, column, false) unless allow_null
# We want to rescue _all_ exceptions here, even those that don't inherit
# from StandardError.
rescue Exception => error # rubocop: disable all
remove_column(table, column)
raise error
raise error
end
end
end

106
spec/lib/gitlab/database/migration_helpers_spec.rb
View file

@ -51,7 +51,7 @@ describe Gitlab::Database::MigrationHelpers do
context 'using PostgreSQL' do
before do
allow(Gitlab::Database).to receive(:postgresql?).and_return(true)
allow(model).to receive(:disable_statement_timeout)
allow(model).to receive(:disable_statement_timeout).and_call_original
end
it 'creates the index concurrently' do
@ -114,12 +114,12 @@ describe Gitlab::Database::MigrationHelpers do
before do
allow(model).to receive(:transaction_open?).and_return(false)
allow(model).to receive(:index_exists?).and_return(true)
allow(model).to receive(:disable_statement_timeout).and_call_original
end
context 'using PostgreSQL' do
before do
allow(model).to receive(:supports_drop_index_concurrently?).and_return(true)
allow(model).to receive(:disable_statement_timeout)
end
describe 'by column name' do
@ -162,7 +162,7 @@ describe Gitlab::Database::MigrationHelpers do
context 'using MySQL' do
it 'removes an index' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
expect(Gitlab::Database).to receive(:postgresql?).and_return(false).twice
expect(model).to receive(:remove_index)
.with(:users, { column: :foo })
@ -228,17 +228,21 @@ describe Gitlab::Database::MigrationHelpers do
end
it 'creates a concurrent foreign key and validates it' do
expect(model).to receive(:disable_statement_timeout)
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:execute).with(/statement_timeout/)
expect(model).to receive(:execute).ordered.with(/NOT VALID/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
expect(model).to receive(:execute).with(/RESET ALL/)
model.add_concurrent_foreign_key(:projects, :users, column: :user_id)
end
it 'appends a valid ON DELETE statement' do
expect(model).to receive(:disable_statement_timeout)
expect(model).to receive(:disable_statement_timeout).and_call_original
expect(model).to receive(:execute).with(/statement_timeout/)
expect(model).to receive(:execute).with(/ON DELETE SET NULL/)
expect(model).to receive(:execute).ordered.with(/VALIDATE CONSTRAINT/)
expect(model).to receive(:execute).with(/RESET ALL/)
model.add_concurrent_foreign_key(:projects, :users,
column: :user_id,
@ -291,22 +295,98 @@ describe Gitlab::Database::MigrationHelpers do
describe '#disable_statement_timeout' do
context 'using PostgreSQL' do
it 'disables statement timeouts' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
context 'with transaction: true' do
it 'disables statement timeouts to current transaction only' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
expect(model).to receive(:execute).with('SET statement_timeout TO 0')
expect(model).to receive(:execute).with('SET LOCAL statement_timeout TO 0')
model.disable_statement_timeout
model.disable_statement_timeout
end
# this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
context 'with real environment', :postgresql, :delete do
before do
model.execute("SET statement_timeout TO '20000'")
end
after do
model.execute('RESET ALL')
end
it 'defines statement to 0 only for current transaction' do
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
model.connection.transaction do
model.disable_statement_timeout
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
end
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
end
end
end
context 'with transaction: false' do
it 'disables statement timeouts on session level' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
expect(model).to receive(:execute).with('SET statement_timeout TO 0')
expect(model).to receive(:execute).with('RESET ALL')
model.disable_statement_timeout(transaction: false) do
# no-op
end
end
it 'raises error when no block is given' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(true)
expect { model.disable_statement_timeout(transaction: false) }.to raise_error(ArgumentError)
end
# this specs runs without an enclosing transaction (:delete truncation method for db_cleaner)
context 'with real environment', :postgresql, :delete do
before do
model.execute("SET statement_timeout TO '20000'")
end
after do
model.execute('RESET ALL')
end
it 'defines statement to 0 for any code run inside block' do
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s')
model.disable_statement_timeout(transaction: false) do
model.connection.transaction do
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
end
expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('0')
end
end
end
end
end
context 'using MySQL' do
it 'does nothing' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
context 'with transaction: true' do
it 'does nothing' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
expect(model).not_to receive(:execute)
expect(model).not_to receive(:execute)
model.disable_statement_timeout
model.disable_statement_timeout
end
end
context 'with transaction: false' do
it 'executes yielded code' do
expect(Gitlab::Database).to receive(:postgresql?).and_return(false)
expect(model).not_to receive(:execute)
expect { |block| model.disable_statement_timeout(transaction: false, &block) }.to yield_control
end
end
end
end