Control statement timeout

Instead of disabling the statement timeout alltogether, we limit this to
6 hours for now (for index create/drop). This allows us to better reason
about how long the overall reindexing for a single index takes at most.

lib/gitlab/database/reindexing/concurrent_reindex.rb

@@ -5,13 +5,13 @@ module Gitlab
     module Reindexing
       class ConcurrentReindex
         include Gitlab::Utils::StrongMemoize
-        include MigrationHelpers
 
         ReindexError = Class.new(StandardError)
 
         PG_IDENTIFIER_LENGTH = 63
         TEMPORARY_INDEX_PREFIX = 'tmp_reindex_'
         REPLACED_INDEX_PREFIX = 'old_reindex_'
+        STATEMENT_TIMEOUT = 6.hours
 
         attr_reader :index, :logger
 
@@ -47,7 +47,7 @@ module Gitlab
           logger.info("creating replacement index #{replacement_index_name}")
           logger.debug("replacement index definition: #{create_replacement_index_statement}")
 
-          disable_statement_timeout do
+          set_statement_timeout do
             connection.execute(create_replacement_index_statement)
           end
 
@@ -88,7 +88,7 @@ module Gitlab
         def remove_index(schema, name)
           logger.info("Removing index #{schema}.#{name}")
 
-          disable_statement_timeout do
+          set_statement_timeout do
             connection.execute(<<~SQL)
               DROP INDEX CONCURRENTLY
               IF EXISTS #{quote_table_name(schema)}.#{quote_table_name(name)}
@@ -110,6 +110,13 @@ module Gitlab
           Gitlab::Database::WithLockRetries.new(**arguments).run(raise_on_exhaustion: true, &block)
         end
 
+        def set_statement_timeout
+          execute("SET statement_timeout TO #{STATEMENT_TIMEOUT}")
+          yield
+        ensure
+          execute('RESET statement_timeout')
+        end
+
         delegate :execute, :quote_table_name, to: :connection
         def connection
           @connection ||= ActiveRecord::Base.connection

spec/lib/gitlab/database/reindexing/concurrent_reindex_spec.rb

@@ -107,11 +107,11 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
     context 'mocked specs' do
       before do
         allow(subject).to receive(:connection).and_return(connection)
-        allow(subject).to receive(:disable_statement_timeout).and_yield
+        allow(connection).to receive(:execute).and_call_original
       end
 
       it 'replaces the existing index with an identical index' do
-        expect(subject).to receive(:disable_statement_timeout).twice.and_yield
+        expect(connection).to receive(:execute).with('SET statement_timeout TO 21600').twice
 
         expect_to_execute_concurrently_in_order(create_index)
 
@@ -136,7 +136,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
         end
 
         it 'replaces the existing index with an identical index' do
-          expect(subject).to receive(:disable_statement_timeout).exactly(3).times.and_yield
+          expect(connection).to receive(:execute).with('SET statement_timeout TO 21600').exactly(3).times
 
           expect_to_execute_concurrently_in_order(drop_index)
           expect_to_execute_concurrently_in_order(create_index)