diff options
| -rw-r--r-- | lib/gitlab/database/migration_helpers.rb | 106 | ||||
| -rw-r--r-- | spec/lib/gitlab/database/migration_helpers_spec.rb | 106 |
2 files changed, 168 insertions, 44 deletions
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index f39b3b6eb5b..1e6a620dd66 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -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 - - transaction do - if limit - add_column(table, column, type, default: nil, limit: limit) - else - add_column(table, column, type, default: nil) + 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 + + # 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 diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index eb7148ff108..c993fa12803 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -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 LOCAL statement_timeout TO 0') + + 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 - expect(model).to receive(:execute).with('SET statement_timeout TO 0') + it 'defines statement to 0 only for current transaction' do + expect(model.execute('SHOW statement_timeout').first['statement_timeout']).to eq('20s') - model.disable_statement_timeout + 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 + end + end - model.disable_statement_timeout + 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 |