Add latest changes from gitlab-org/gitlab@15-9-stable-eev15.9.0-rc42

40 files changed, 779 insertions, 174 deletions

diff --git a/lib/gitlab/database/indexing_exclusive_lease_guard.rb b/lib/gitlab/database/async_ddl_exclusive_lease_guard.rb
index fb45de347e6..5742e96c9b3 100644
--- a/lib/gitlab/database/indexing_exclusive_lease_guard.rb
+++ b/lib/gitlab/database/async_ddl_exclusive_lease_guard.rb
@@ -2,16 +2,16 @@
 
 module Gitlab
   module Database
-    module IndexingExclusiveLeaseGuard
+    module AsyncDdlExclusiveLeaseGuard
       extend ActiveSupport::Concern
       include ExclusiveLeaseGuard
 
       def lease_key
-        @lease_key ||= "gitlab/database/indexing/actions/#{database_config_name}"
+        @lease_key ||= "gitlab/database/asyncddl/actions/#{database_config_name}"
       end
 
       def database_config_name
-        Gitlab::Database.db_config_name(connection)
+        connection_db_config.name
       end
     end
   end
diff --git a/lib/gitlab/database/async_foreign_keys.rb b/lib/gitlab/database/async_foreign_keys.rb
new file mode 100644
index 00000000000..115ae9ba2e8
--- /dev/null
+++ b/lib/gitlab/database/async_foreign_keys.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncForeignKeys
+      DEFAULT_ENTRIES_PER_INVOCATION = 2
+
+      def self.validate_pending_entries!(how_many: DEFAULT_ENTRIES_PER_INVOCATION)
+        PostgresAsyncForeignKeyValidation.ordered.limit(how_many).each do |record|
+          ForeignKeyValidator.new(record).perform
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/async_foreign_keys/foreign_key_validator.rb b/lib/gitlab/database/async_foreign_keys/foreign_key_validator.rb
new file mode 100644
index 00000000000..5958c56a45a
--- /dev/null
+++ b/lib/gitlab/database/async_foreign_keys/foreign_key_validator.rb
@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncForeignKeys
+      class ForeignKeyValidator
+        include AsyncDdlExclusiveLeaseGuard
+
+        TIMEOUT_PER_ACTION = 1.day
+        STATEMENT_TIMEOUT = 12.hours
+
+        def initialize(async_validation)
+          @async_validation = async_validation
+        end
+
+        def perform
+          try_obtain_lease do
+            if foreign_key_exists?
+              log_index_info("Starting to validate foreign key")
+              validate_foreign_with_error_handling
+              log_index_info("Finished validating foreign key")
+            else
+              log_index_info(skip_log_message)
+              async_validation.destroy!
+            end
+          end
+        end
+
+        private
+
+        attr_reader :async_validation
+
+        delegate :connection, :name, :table_name, :connection_db_config, to: :async_validation
+
+        def foreign_key_exists?
+          relation = if table_name =~ Gitlab::Database::FULLY_QUALIFIED_IDENTIFIER
+                       Gitlab::Database::PostgresForeignKey.by_constrained_table_identifier(table_name)
+                     else
+                       Gitlab::Database::PostgresForeignKey.by_constrained_table_name(table_name)
+                     end
+
+          relation.by_name(name).exists?
+        end
+
+        def validate_foreign_with_error_handling
+          validate_foreign_key
+          async_validation.destroy!
+        rescue StandardError => error
+          async_validation.handle_exception!(error)
+
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error)
+          Gitlab::AppLogger.error(message: error.message, **logging_options)
+        end
+
+        def validate_foreign_key
+          set_statement_timeout do
+            connection.execute(<<~SQL.squish)
+              ALTER TABLE #{connection.quote_table_name(table_name)}
+              VALIDATE CONSTRAINT #{connection.quote_column_name(name)};
+            SQL
+          end
+        end
+
+        def set_statement_timeout
+          connection.execute(format("SET statement_timeout TO '%ds'", STATEMENT_TIMEOUT))
+          yield
+        ensure
+          connection.execute('RESET statement_timeout')
+        end
+
+        def lease_timeout
+          TIMEOUT_PER_ACTION
+        end
+
+        def log_index_info(message)
+          Gitlab::AppLogger.info(message: message, **logging_options)
+        end
+
+        def skip_log_message
+          "Skipping #{name} validation since it does not exist. " \
+            "The queuing entry will be deleted"
+        end
+
+        def logging_options
+          {
+            fk_name: name,
+            table_name: table_name,
+            class: self.class.name.to_s
+          }
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/async_foreign_keys/migration_helpers.rb b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
new file mode 100644
index 00000000000..b8b9fc6d156
--- /dev/null
+++ b/lib/gitlab/database/async_foreign_keys/migration_helpers.rb
@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncForeignKeys
+      module MigrationHelpers
+        # Prepares a foreign key for asynchronous validation.
+        #
+        # Stores the FK information in the postgres_async_foreign_key_validations
+        # table to be executed later.
+        #
+        def prepare_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          fk_name = name || concurrent_foreign_key_name(table_name, column_name)
+
+          unless foreign_key_exists?(table_name, name: fk_name)
+            raise missing_schema_object_message(table_name, "foreign key", fk_name)
+          end
+
+          async_validation = PostgresAsyncForeignKeyValidation
+            .find_or_create_by!(name: fk_name, table_name: table_name)
+
+          Gitlab::AppLogger.info(
+            message: 'Prepared FK for async validation',
+            table_name: async_validation.table_name,
+            fk_name: async_validation.name)
+
+          async_validation
+        end
+
+        def unprepare_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          fk_name = name || concurrent_foreign_key_name(table_name, column_name)
+
+          PostgresAsyncForeignKeyValidation.find_by(name: fk_name).try(&:destroy)
+        end
+
+        def prepare_partitioned_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          Gitlab::Database::PostgresPartitionedTable.each_partition(table_name) do |partition|
+            prepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
+          end
+        end
+
+        def unprepare_partitioned_async_foreign_key_validation(table_name, column_name = nil, name: nil)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
+          return unless async_fk_validation_available?
+
+          Gitlab::Database::PostgresPartitionedTable.each_partition(table_name) do |partition|
+            unprepare_async_foreign_key_validation(partition.identifier, column_name, name: name)
+          end
+        end
+
+        private
+
+        def async_fk_validation_available?
+          connection.table_exists?(:postgres_async_foreign_key_validations)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/async_foreign_keys/postgres_async_foreign_key_validation.rb b/lib/gitlab/database/async_foreign_keys/postgres_async_foreign_key_validation.rb
new file mode 100644
index 00000000000..de69a3d496f
--- /dev/null
+++ b/lib/gitlab/database/async_foreign_keys/postgres_async_foreign_key_validation.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncForeignKeys
+      class PostgresAsyncForeignKeyValidation < SharedModel
+        include QueueErrorHandlingConcern
+
+        self.table_name = 'postgres_async_foreign_key_validations'
+
+        MAX_IDENTIFIER_LENGTH = Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH
+        MAX_LAST_ERROR_LENGTH = 10_000
+
+        validates :name, presence: true, uniqueness: true, length: { maximum: MAX_IDENTIFIER_LENGTH }
+        validates :table_name, presence: true, length: { maximum: MAX_IDENTIFIER_LENGTH }
+
+        scope :ordered, -> { order(attempts: :asc, id: :asc) }
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/async_indexes.rb b/lib/gitlab/database/async_indexes.rb
index 6f301a66803..581c7e7ff94 100644
--- a/lib/gitlab/database/async_indexes.rb
+++ b/lib/gitlab/database/async_indexes.rb
@@ -16,6 +16,15 @@ module Gitlab
           IndexDestructor.new(async_index).perform
         end
       end
+
+      def self.execute_pending_actions!(how_many: DEFAULT_INDEXES_PER_INVOCATION)
+        queue_ids = PostgresAsyncIndex.ordered.limit(how_many).pluck(:id)
+        removal_actions = PostgresAsyncIndex.where(id: queue_ids).to_drop.ordered
+        creation_actions = PostgresAsyncIndex.where(id: queue_ids).to_create.ordered
+
+        removal_actions.each { |async_index| IndexDestructor.new(async_index).perform }
+        creation_actions.each { |async_index| IndexCreator.new(async_index).perform }
+      end
     end
   end
 end
diff --git a/lib/gitlab/database/async_indexes/index_base.rb b/lib/gitlab/database/async_indexes/index_base.rb
new file mode 100644
index 00000000000..bde75e12295
--- /dev/null
+++ b/lib/gitlab/database/async_indexes/index_base.rb
@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module AsyncIndexes
+      class IndexBase
+        include AsyncDdlExclusiveLeaseGuard
+        extend ::Gitlab::Utils::Override
+
+        TIMEOUT_PER_ACTION = 1.day
+
+        def initialize(async_index)
+          @async_index = async_index
+        end
+
+        def perform
+          try_obtain_lease do
+            if preconditions_met?
+              log_index_info("Starting async index #{action_type}")
+              execute_action_with_error_handling
+              log_index_info("Finished async index #{action_type}")
+            else
+              log_index_info(skip_log_message)
+              async_index.destroy!
+            end
+          end
+        end
+
+        private
+
+        attr_reader :async_index
+
+        delegate :connection, :connection_db_config, to: :async_index
+
+        def preconditions_met?
+          raise NotImplementedError, 'must implement preconditions_met?'
+        end
+
+        def action_type
+          raise NotImplementedError, 'must implement action_type'
+        end
+
+        def execute_action_with_error_handling
+          around_execution { execute_action }
+        rescue StandardError => error
+          async_index.handle_exception!(error)
+
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error)
+          Gitlab::AppLogger.error(message: error.message, **logging_options)
+        end
+
+        def around_execution
+          yield
+        end
+
+        def execute_action
+          connection.execute(async_index.definition)
+          async_index.destroy!
+        end
+
+        def index_exists?
+          connection.indexes(async_index.table_name).any? do |index|
+            index.name == async_index.name
+          end
+        end
+
+        def lease_timeout
+          TIMEOUT_PER_ACTION
+        end
+
+        def log_index_info(message)
+          Gitlab::AppLogger.info(message: message, **logging_options)
+        end
+
+        def skip_log_message
+          "Skipping index #{action_type} since preconditions are not met. " \
+            "The queuing entry will be deleted"
+        end
+
+        def logging_options
+          {
+            table_name: async_index.table_name,
+            index_name: async_index.name,
+            class: self.class.name.to_s
+          }
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/async_indexes/index_creator.rb b/lib/gitlab/database/async_indexes/index_creator.rb
index 3ae2bb7b3e5..c5f4c5f30ad 100644
--- a/lib/gitlab/database/async_indexes/index_creator.rb
+++ b/lib/gitlab/database/async_indexes/index_creator.rb
@@ -3,48 +3,24 @@
 module Gitlab
   module Database
     module AsyncIndexes
-      class IndexCreator
-        include IndexingExclusiveLeaseGuard
-
-        TIMEOUT_PER_ACTION = 1.day
+      class IndexCreator < AsyncIndexes::IndexBase
         STATEMENT_TIMEOUT = 20.hours
 
-        def initialize(async_index)
-          @async_index = async_index
-        end
-
-        def perform
-          try_obtain_lease do
-            if index_exists?
-              log_index_info('Skipping index creation as the index exists')
-            else
-              log_index_info('Creating async index')
-
-              set_statement_timeout do
-                connection.execute(async_index.definition)
-              end
-
-              log_index_info('Finished creating async index')
-            end
-
-            async_index.destroy
-          end
-        end
-
         private
 
-        attr_reader :async_index
-
-        def index_exists?
-          connection.indexes(async_index.table_name).any? { |index| index.name == async_index.name }
+        override :preconditions_met?
+        def preconditions_met?
+          !index_exists?
         end
 
-        def connection
-          @connection ||= async_index.connection
+        override :action_type
+        def action_type
+          'creation'
         end
 
-        def lease_timeout
-          TIMEOUT_PER_ACTION
+        override :around_execution
+        def around_execution(&block)
+          set_statement_timeout(&block)
         end
 
         def set_statement_timeout
@@ -53,10 +29,6 @@ module Gitlab
         ensure
           connection.execute('RESET statement_timeout')
         end
-
-        def log_index_info(message)
-          Gitlab::AppLogger.info(message: message, table_name: async_index.table_name, index_name: async_index.name)
-        end
       end
     end
   end
diff --git a/lib/gitlab/database/async_indexes/index_destructor.rb b/lib/gitlab/database/async_indexes/index_destructor.rb
index 66955df9d04..5596e099cb6 100644
--- a/lib/gitlab/database/async_indexes/index_destructor.rb
+++ b/lib/gitlab/database/async_indexes/index_destructor.rb
@@ -3,58 +3,29 @@
 module Gitlab
   module Database
     module AsyncIndexes
-      class IndexDestructor
-        include IndexingExclusiveLeaseGuard
-
-        TIMEOUT_PER_ACTION = 1.day
-
-        def initialize(async_index)
-          @async_index = async_index
-        end
-
-        def perform
-          try_obtain_lease do
-            if !index_exists?
-              log_index_info('Skipping dropping as the index does not exist')
-            else
-              log_index_info('Dropping async index')
-
-              retries = Gitlab::Database::WithLockRetriesOutsideTransaction.new(
-                connection: connection,
-                timing_configuration: Gitlab::Database::Reindexing::REMOVE_INDEX_RETRY_CONFIG,
-                klass: self.class,
-                logger: Gitlab::AppLogger
-              )
-
-              retries.run(raise_on_exhaustion: false) do
-                connection.execute(async_index.definition)
-              end
-
-              log_index_info('Finished dropping async index')
-            end
-
-            async_index.destroy
-          end
-        end
-
+      class IndexDestructor < AsyncIndexes::IndexBase
         private
 
-        attr_reader :async_index
-
-        def index_exists?
-          connection.indexes(async_index.table_name).any? { |index| index.name == async_index.name }
+        override :preconditions_met?
+        def preconditions_met?
+          index_exists?
         end
 
-        def connection
-          @connection ||= async_index.connection
+        override :action_type
+        def action_type
+          'removal'
         end
 
-        def lease_timeout
-          TIMEOUT_PER_ACTION
-        end
+        override :around_execution
+        def around_execution(&block)
+          retries = Gitlab::Database::WithLockRetriesOutsideTransaction.new(
+            connection: connection,
+            timing_configuration: Gitlab::Database::Reindexing::REMOVE_INDEX_RETRY_CONFIG,
+            klass: self.class,
+            logger: Gitlab::AppLogger
+          )
 
-        def log_index_info(message)
-          Gitlab::AppLogger.info(message: message, table_name: async_index.table_name, index_name: async_index.name)
+          retries.run(raise_on_exhaustion: false, &block)
         end
       end
     end
diff --git a/lib/gitlab/database/async_indexes/migration_helpers.rb b/lib/gitlab/database/async_indexes/migration_helpers.rb
index c8f6761534c..f459c43e0ee 100644
--- a/lib/gitlab/database/async_indexes/migration_helpers.rb
+++ b/lib/gitlab/database/async_indexes/migration_helpers.rb
@@ -22,7 +22,7 @@ module Gitlab
           return unless async_index_creation_available?
 
           PostgresAsyncIndex.find_by(name: index_name).try do |async_index|
-            async_index.destroy
+            async_index.destroy!
           end
         end
 
diff --git a/lib/gitlab/database/async_indexes/postgres_async_index.rb b/lib/gitlab/database/async_indexes/postgres_async_index.rb
index dc932482d40..9f5f39613ed 100644
--- a/lib/gitlab/database/async_indexes/postgres_async_index.rb
+++ b/lib/gitlab/database/async_indexes/postgres_async_index.rb
@@ -4,6 +4,8 @@ module Gitlab
   module Database
     module AsyncIndexes
       class PostgresAsyncIndex < SharedModel
+        include QueueErrorHandlingConcern
+
         self.table_name = 'postgres_async_indexes'
 
         MAX_IDENTIFIER_LENGTH = Gitlab::Database::MigrationHelpers::MAX_IDENTIFIER_NAME_LENGTH
@@ -15,6 +17,7 @@ module Gitlab
 
         scope :to_create, -> { where("definition ILIKE 'CREATE%'") }
         scope :to_drop, -> { where("definition ILIKE 'DROP%'") }
+        scope :ordered, -> { order(attempts: :asc, id: :asc) }
 
         def to_s
           definition
diff --git a/lib/gitlab/database/bulk_update.rb b/lib/gitlab/database/bulk_update.rb
index 4b4a9b38fd8..36dbb157b0d 100644
--- a/lib/gitlab/database/bulk_update.rb
+++ b/lib/gitlab/database/bulk_update.rb
@@ -43,15 +43,7 @@ module Gitlab
         end
 
         def update!
-          if without_prepared_statement?
-            # A workaround for https://github.com/rails/rails/issues/24893
-            # When prepared statements are prevented (such as when using the
-            # query counter or in omnibus by default), we cannot call
-            # `exec_update`, since that will discard the bindings.
-            connection.send(:exec_no_cache, sql, log_name, params) # rubocop: disable GitlabSecurity/PublicSend
-          else
-            connection.exec_update(sql, log_name, params)
-          end
+          connection.exec_update(sql, log_name, params)
         end
 
         def self.column_definitions(model, columns)
@@ -93,14 +85,6 @@ module Gitlab
           end
         end
 
-        # A workaround for https://github.com/rails/rails/issues/24893
-        # We need to detect if prepared statements have been disabled.
-        def without_prepared_statement?
-          strong_memoize(:without_prepared_statement) do
-            connection.send(:without_prepared_statement?, [1]) # rubocop: disable GitlabSecurity/PublicSend
-          end
-        end
-
         def query_attribute(column, key, values)
           value = values[column.name]
           key[column.name] = value if key.try(:id) # optimistic update
diff --git a/lib/gitlab/database/connection_timer.rb b/lib/gitlab/database/connection_timer.rb
index f9b893ffd0f..4eb214e74f4 100644
--- a/lib/gitlab/database/connection_timer.rb
+++ b/lib/gitlab/database/connection_timer.rb
@@ -27,7 +27,7 @@ module Gitlab
         end
 
         def current_clock_value
-          Concurrent.monotonic_time
+          Process.clock_gettime(Process::CLOCK_MONOTONIC)
         end
       end
 
diff --git a/lib/gitlab/database/load_balancing/connection_proxy.rb b/lib/gitlab/database/load_balancing/connection_proxy.rb
index f0343f9d8b5..622e310ead3 100644
--- a/lib/gitlab/database/load_balancing/connection_proxy.rb
+++ b/lib/gitlab/database/load_balancing/connection_proxy.rb
@@ -97,11 +97,11 @@ module Gitlab
           if current_session.use_primary? &&
               !current_session.use_replicas_for_read_queries?
             @load_balancer.read_write do |connection|
-              connection.send(...)
+              connection.public_send(...)
             end
           else
             @load_balancer.read do |connection|
-              connection.send(...)
+              connection.public_send(...)
             end
           end
         end
@@ -117,7 +117,7 @@ module Gitlab
           end
 
           @load_balancer.read_write do |connection|
-            connection.send(...)
+            connection.public_send(...)
           end
         end
 
diff --git a/lib/gitlab/database/load_balancing/load_balancer.rb b/lib/gitlab/database/load_balancing/load_balancer.rb
index cb3a378ad64..23476e1f5e9 100644
--- a/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -105,11 +105,9 @@ module Gitlab
         def read_write
           connection = nil
           transaction_open = nil
-          attempts = 3
 
-          if prevent_load_balancer_retries_in_transaction?
-            attempts = 1 if pool.connection.transaction_open?
-          end
+          # Retry only once when in a transaction (see https://gitlab.com/gitlab-org/gitlab/-/issues/220242)
+          attempts = pool.connection.transaction_open? ? 1 : 3
 
           # In the event of a failover the primary may be briefly unavailable.
           # Instead of immediately grinding to a halt we'll retry the operation
@@ -348,10 +346,6 @@ module Gitlab
           row = ar_connection.select_all(sql).first
           row['location'] if row
         end
-
-        def prevent_load_balancer_retries_in_transaction?
-          Gitlab::Utils.to_boolean(ENV['PREVENT_LOAD_BALANCER_RETRIES_IN_TRANSACTION'], default: false)
-        end
       end
     end
   end
diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb
index 8e5dc98e96e..f5cb83e398a 100644
--- a/lib/gitlab/database/load_balancing/sticking.rb
+++ b/lib/gitlab/database/load_balancing/sticking.rb
@@ -121,19 +121,19 @@ module Gitlab
         end
 
         def unstick(namespace, id)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.del(redis_key_for(namespace, id))
           end
         end
 
         def set_write_location_for(namespace, id, location)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
           end
         end
 
         def last_write_location_for(namespace, id)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.get(redis_key_for(namespace, id))
           end
         end
@@ -143,6 +143,12 @@ module Gitlab
 
           "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
         end
+
+        private
+
+        def with_redis(&block)
+          Gitlab::Redis::DbLoadBalancing.with(&block)
+        end
       end
     end
   end
diff --git a/lib/gitlab/database/lock_writes_manager.rb b/lib/gitlab/database/lock_writes_manager.rb
index 2e08e1ffb42..83884e89d6e 100644
--- a/lib/gitlab/database/lock_writes_manager.rb
+++ b/lib/gitlab/database/lock_writes_manager.rb
@@ -10,18 +10,6 @@ module Gitlab
       # See https://www.postgresql.org/message-id/16934.1568989957%40sss.pgh.pa.us
       EXPECTED_TRIGGER_RECORD_COUNT = 3
 
-      def self.tables_to_lock(connection)
-        Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
-          yield table_name, schema_name
-        end
-
-        Gitlab::Database::SharedModel.using_connection(connection) do
-          Postgresql::DetachedPartition.find_each do |detached_partition|
-            yield detached_partition.fully_qualified_table_name, detached_partition.table_schema
-          end
-        end
-      end
-
       def initialize(table_name:, connection:, database_name:, with_retries: true, logger: nil, dry_run: false)
         @table_name = table_name
         @connection = connection
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index e41107370ec..9c1cb8e352c 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -14,6 +14,7 @@ module Gitlab
       include DynamicModelHelpers
       include RenameTableHelpers
       include AsyncIndexes::MigrationHelpers
+      include AsyncForeignKeys::MigrationHelpers
 
       def define_batchable_model(table_name, connection: self.connection)
         super(table_name, connection: connection)
diff --git a/lib/gitlab/database/migrations/background_migration_helpers.rb b/lib/gitlab/database/migrations/background_migration_helpers.rb
index 25e75a10bb3..60df3370046 100644
--- a/lib/gitlab/database/migrations/background_migration_helpers.rb
+++ b/lib/gitlab/database/migrations/background_migration_helpers.rb
@@ -200,11 +200,14 @@ module Gitlab
           end
         end
 
+        # rubocop: disable Style/ArgumentsForwarding
+        # Reason: the default argument will not apply if we just forward via `...`
         def migrate_in(*args, coordinator: coordinator_for_tracking_database)
           with_migration_context do
             coordinator.perform_in(*args)
           end
         end
+        # rubocop: enable Style/ArgumentsForwarding
 
         def delete_queued_jobs(class_name)
           coordinator_for_tracking_database.steal(class_name) do |job|
diff --git a/lib/gitlab/database/migrations/base_background_runner.rb b/lib/gitlab/database/migrations/base_background_runner.rb
index 8975c04e33a..840add8783d 100644
--- a/lib/gitlab/database/migrations/base_background_runner.rb
+++ b/lib/gitlab/database/migrations/base_background_runner.rb
@@ -38,13 +38,15 @@ module Gitlab
         def run_jobs_for_migration(migration_name:, jobs:, run_until:)
           per_background_migration_result_dir = File.join(@result_dir, migration_name)
 
-          instrumentation = Instrumentation.new(result_dir: per_background_migration_result_dir)
+          instrumentation = Instrumentation.new(result_dir: per_background_migration_result_dir,
+                                                observer_classes: observers)
+
           batch_names = (1..).each.lazy.map { |i| "batch_#{i}" }
 
           jobs.each do |j|
             break if run_until <= Time.current
 
-            meta = migration_meta(j)
+            meta = { job_meta: job_meta(j) }
 
             instrumentation.observe(version: nil,
                                     name: batch_names.next,
@@ -55,9 +57,13 @@ module Gitlab
           end
         end
 
-        def migration_meta(_job)
+        def job_meta(_job)
           {}
         end
+
+        def observers
+          ::Gitlab::Database::Migrations::Observers.all_observers
+        end
       end
     end
   end
diff --git a/lib/gitlab/database/migrations/observation.rb b/lib/gitlab/database/migrations/observation.rb
index 80388c4dbbb..cd048beac96 100644
--- a/lib/gitlab/database/migrations/observation.rb
+++ b/lib/gitlab/database/migrations/observation.rb
@@ -4,16 +4,12 @@
 module Gitlab
   module Database
     module Migrations
-      Observation = Struct.new(
-        :version,
-        :name,
-        :walltime,
-        :success,
-        :total_database_size_change,
-        :meta,
-        :query_statistics,
-        keyword_init: true
-      )
+      Observation = Struct.new(:version, :name, :walltime, :success, :total_database_size_change,
+        :meta, :query_statistics, keyword_init: true) do
+        def to_json(...)
+          as_json.except('meta').to_json(...)
+        end
+      end
     end
   end
 end
diff --git a/lib/gitlab/database/migrations/observers/batch_details.rb b/lib/gitlab/database/migrations/observers/batch_details.rb
new file mode 100644
index 00000000000..0f8cdcf3cd6
--- /dev/null
+++ b/lib/gitlab/database/migrations/observers/batch_details.rb
@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        class BatchDetails < MigrationObserver
+          FILE_NAME = 'batch-details.json'
+
+          def before
+            @started_at = get_time
+          end
+
+          def after
+            @finished_at = get_time
+          end
+
+          def record
+            File.open(path, 'wb') { |file| file.write(file_contents.to_json) }
+          end
+
+          private
+
+          attr_reader :started_at, :finished_at
+
+          def file_contents
+            { time_spent: time_spent }.merge(job_meta)
+          end
+
+          def get_time
+            Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          end
+
+          def time_spent
+            @time_spent ||= (finished_at - started_at).round(2)
+          end
+
+          def path
+            File.join(output_dir, FILE_NAME)
+          end
+
+          def job_meta
+            meta = observation.meta
+
+            return {} unless meta
+
+            meta[:job_meta].to_h
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/migrations/test_batched_background_runner.rb b/lib/gitlab/database/migrations/test_batched_background_runner.rb
index c123d01f327..01fdba22c19 100644
--- a/lib/gitlab/database/migrations/test_batched_background_runner.rb
+++ b/lib/gitlab/database/migrations/test_batched_background_runner.rb
@@ -6,6 +6,8 @@ module Gitlab
       class TestBatchedBackgroundRunner < BaseBackgroundRunner
         include Gitlab::Database::DynamicModelHelpers
 
+        MIGRATION_DETAILS_FILE_NAME = 'details.json'
+
         def initialize(result_dir:, connection:, from_id:)
           super(result_dir: result_dir, connection: connection)
           @connection = connection
@@ -51,6 +53,7 @@ module Gitlab
                     migration.column_name,
                     batch_min_value: batch_start,
                     batch_size: migration.batch_size,
+                    job_class: migration.job_class,
                     job_arguments: migration.job_arguments
                   )
 
@@ -64,7 +67,11 @@ module Gitlab
                 end
               end
 
-              [migration.job_class_name, jobs_to_sample]
+              job_class_name = migration.job_class_name
+
+              export_migration_details(job_class_name, migration.slice(:interval, :total_tuple_count, :max_batch_size))
+
+              [job_class_name, jobs_to_sample]
             end
           end
         end
@@ -112,11 +119,25 @@ module Gitlab
           Gitlab::Database::SharedModel.using_connection(connection, &block)
         end
 
-        def migration_meta(job)
+        def job_meta(job)
           set_shared_model_connection do
-            job.batched_migration.slice(:max_batch_size, :total_tuple_count, :interval)
+            job.slice(:min_value, :max_value, :batch_size, :sub_batch_size, :pause_ms)
           end
         end
+
+        def export_migration_details(migration_name, attributes)
+          directory = result_dir.join(migration_name)
+
+          FileUtils.mkdir_p(directory) unless Dir.exist?(directory)
+
+          File.write(directory.join(MIGRATION_DETAILS_FILE_NAME), attributes.to_json)
+        end
+
+        def observers
+          ::Gitlab::Database::Migrations::Observers.all_observers + [
+            ::Gitlab::Database::Migrations::Observers::BatchDetails
+          ]
+        end
       end
     end
   end
diff --git a/lib/gitlab/database/postgres_constraint.rb b/lib/gitlab/database/postgres_constraint.rb
index fa590914332..fa3870cb9c7 100644
--- a/lib/gitlab/database/postgres_constraint.rb
+++ b/lib/gitlab/database/postgres_constraint.rb
@@ -4,7 +4,6 @@ module Gitlab
   module Database
     # Backed by the postgres_constraints view
     class PostgresConstraint < SharedModel
-      IDENTIFIER_REGEX = /^\w+\.\w+$/.freeze
       self.primary_key = :oid
 
       scope :check_constraints, -> { where(constraint_type: 'c') }
@@ -18,7 +17,7 @@ module Gitlab
       scope :valid, -> { where(constraint_valid: true) }
 
       scope :by_table_identifier, ->(identifier) do
-        unless identifier =~ IDENTIFIER_REGEX
+        unless identifier =~ Gitlab::Database::FULLY_QUALIFIED_IDENTIFIER
           raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}"
         end
 
diff --git a/lib/gitlab/database/postgres_foreign_key.rb b/lib/gitlab/database/postgres_foreign_key.rb
index d3ede45fe86..04ef574a451 100644
--- a/lib/gitlab/database/postgres_foreign_key.rb
+++ b/lib/gitlab/database/postgres_foreign_key.rb
@@ -5,17 +5,23 @@ module Gitlab
     class PostgresForeignKey < SharedModel
       self.primary_key = :oid
 
-      # These values come from the possible confdeltype values in pg_constraint
-      enum on_delete_action: {
+      # These values come from the possible confdeltype / confupdtype values in pg_constraint
+      ACTION_TYPES = {
         restrict: 'r',
         cascade: 'c',
         nullify: 'n',
         set_default: 'd',
         no_action: 'a'
-      }
+      }.freeze
+
+      enum on_delete_action: ACTION_TYPES, _prefix: :on_delete
+
+      enum on_update_action: ACTION_TYPES, _prefix: :on_update
 
       scope :by_referenced_table_identifier, ->(identifier) do
-        raise ArgumentError, "Referenced table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+        unless identifier =~ Database::FULLY_QUALIFIED_IDENTIFIER
+          raise ArgumentError, "Referenced table name is not fully qualified with a schema: #{identifier}"
+        end
 
         where(referenced_table_identifier: identifier)
       end
@@ -23,7 +29,9 @@ module Gitlab
       scope :by_referenced_table_name, ->(name) { where(referenced_table_name: name) }
 
       scope :by_constrained_table_identifier, ->(identifier) do
-        raise ArgumentError, "Constrained table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+        unless identifier =~ Database::FULLY_QUALIFIED_IDENTIFIER
+          raise ArgumentError, "Constrained table name is not fully qualified with a schema: #{identifier}"
+        end
 
         where(constrained_table_identifier: identifier)
       end
@@ -43,6 +51,12 @@ module Gitlab
 
         where(on_delete_action: on_delete)
       end
+
+      scope :by_on_update_action, ->(on_update) do
+        raise ArgumentError, "Invalid on_update action #{on_update}" unless on_update_actions.key?(on_update)
+
+        where(on_update_action: on_update)
+      end
     end
   end
 end
diff --git a/lib/gitlab/database/postgres_index.rb b/lib/gitlab/database/postgres_index.rb
index 4a9d8728c83..50009cadf5d 100644
--- a/lib/gitlab/database/postgres_index.rb
+++ b/lib/gitlab/database/postgres_index.rb
@@ -14,7 +14,9 @@ module Gitlab
       has_many :queued_reindexing_actions, class_name: 'Gitlab::Database::Reindexing::QueuedAction', foreign_key: :index_identifier
 
       scope :by_identifier, ->(identifier) do
-        raise ArgumentError, "Index name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+        unless identifier =~ Gitlab::Database::FULLY_QUALIFIED_IDENTIFIER
+          raise ArgumentError, "Index name is not fully qualified with a schema: #{identifier}"
+        end
 
         find(identifier)
       end
diff --git a/lib/gitlab/database/postgres_partition.rb b/lib/gitlab/database/postgres_partition.rb
index e4f70ee1745..36dc6818157 100644
--- a/lib/gitlab/database/postgres_partition.rb
+++ b/lib/gitlab/database/postgres_partition.rb
@@ -8,7 +8,9 @@ module Gitlab
       belongs_to :postgres_partitioned_table, foreign_key: 'parent_identifier', primary_key: 'identifier'
 
       scope :for_identifier, ->(identifier) do
-        raise ArgumentError, "Partition name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+        unless identifier =~ Gitlab::Database::FULLY_QUALIFIED_IDENTIFIER
+          raise ArgumentError, "Partition name is not fully qualified with a schema: #{identifier}"
+        end
 
         where(primary_key => identifier)
       end
diff --git a/lib/gitlab/database/postgres_partitioned_table.rb b/lib/gitlab/database/postgres_partitioned_table.rb
index 3bd342f940f..fead7379e43 100644
--- a/lib/gitlab/database/postgres_partitioned_table.rb
+++ b/lib/gitlab/database/postgres_partitioned_table.rb
@@ -10,7 +10,9 @@ module Gitlab
       has_many :postgres_partitions, foreign_key: 'parent_identifier', primary_key: 'identifier'
 
       scope :by_identifier, ->(identifier) do
-        raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+        unless identifier =~ Gitlab::Database::FULLY_QUALIFIED_IDENTIFIER
+          raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}"
+        end
 
         find(identifier)
       end
@@ -19,6 +21,13 @@ module Gitlab
         find_by("identifier = concat(current_schema(), '.', ?)", name)
       end
 
+      def self.each_partition(table_name, &block)
+        find_by_name_in_current_schema(table_name)
+        .postgres_partitions
+        .order(:name)
+        .each(&block)
+      end
+
       def dynamic?
         DYNAMIC_PARTITION_STRATEGIES.include?(strategy)
       end
diff --git a/lib/gitlab/database/query_analyzers/restrict_allowed_schemas.rb b/lib/gitlab/database/query_analyzers/restrict_allowed_schemas.rb
index c51282c9a55..4ae3622479f 100644
--- a/lib/gitlab/database/query_analyzers/restrict_allowed_schemas.rb
+++ b/lib/gitlab/database/query_analyzers/restrict_allowed_schemas.rb
@@ -9,7 +9,18 @@ module Gitlab
         DMLNotAllowedError = Class.new(UnsupportedSchemaError)
         DMLAccessDeniedError = Class.new(UnsupportedSchemaError)
 
-        IGNORED_SCHEMAS = %i[gitlab_shared gitlab_internal].freeze
+        # Re-map schemas observed schemas to a single cluster mode
+        # - symbol:
+        #     The mapped schema indicates that it contains all data in a single-cluster mode
+        # - nil:
+        #     Inidicates that changes made to this schema are ignored and always allowed
+        SCHEMA_MAPPING = {
+          gitlab_shared: nil,
+          gitlab_internal: nil,
+
+          # Pods specific changes
+          gitlab_main_clusterwide: :gitlab_main
+        }.freeze
 
         class << self
           def enabled?
@@ -90,7 +101,13 @@ module Gitlab
 
           def dml_schemas(tables)
             extra_schemas = ::Gitlab::Database::GitlabSchema.table_schemas(tables)
-            extra_schemas.subtract(IGNORED_SCHEMAS)
+
+            SCHEMA_MAPPING.each do |schema, mapped_schema|
+              next unless extra_schemas.delete?(schema)
+
+              extra_schemas.add(mapped_schema) if mapped_schema
+            end
+
             extra_schemas
           end
 
diff --git a/lib/gitlab/database/queue_error_handling_concern.rb b/lib/gitlab/database/queue_error_handling_concern.rb
new file mode 100644
index 00000000000..7d35426d029
--- /dev/null
+++ b/lib/gitlab/database/queue_error_handling_concern.rb
@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module QueueErrorHandlingConcern
+      extend ActiveSupport::Concern
+
+      MAX_LAST_ERROR_LENGTH = 10_000
+
+      included do
+        validates :last_error, length: { maximum: MAX_LAST_ERROR_LENGTH },
+          if: ->(record) { record.respond_to?(:last_error) }
+      end
+
+      def handle_exception!(error)
+        transaction do
+          increment!(:attempts)
+          update!(last_error: format_last_error(error))
+        end
+      end
+
+      private
+
+      def format_last_error(error)
+        [error.message]
+          .concat(error.backtrace)
+          .join("\n")
+          .truncate(MAX_LAST_ERROR_LENGTH)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/reindexing.rb b/lib/gitlab/database/reindexing.rb
index aba45fcc57b..78de7161a0f 100644
--- a/lib/gitlab/database/reindexing.rb
+++ b/lib/gitlab/database/reindexing.rb
@@ -28,6 +28,7 @@ module Gitlab
           # Hack: Before we do actual reindexing work, create async indexes
           Gitlab::Database::AsyncIndexes.create_pending_indexes! if Feature.enabled?(:database_async_index_creation, type: :ops)
           Gitlab::Database::AsyncIndexes.drop_pending_indexes!
+          Gitlab::Database::AsyncForeignKeys.validate_pending_entries! if Feature.enabled?(:database_async_foreign_key_validation, type: :ops)
 
           automatic_reindexing
         end
diff --git a/lib/gitlab/database/reindexing/coordinator.rb b/lib/gitlab/database/reindexing/coordinator.rb
index eca118a4ff2..57e2e0c1beb 100644
--- a/lib/gitlab/database/reindexing/coordinator.rb
+++ b/lib/gitlab/database/reindexing/coordinator.rb
@@ -4,7 +4,7 @@ module Gitlab
   module Database
     module Reindexing
       class Coordinator
-        include IndexingExclusiveLeaseGuard
+        include AsyncDdlExclusiveLeaseGuard
 
         # Maximum lease time for the global Redis lease
         # This should be higher than the maximum time for any
@@ -54,7 +54,7 @@ module Gitlab
 
         private
 
-        delegate :connection, to: :index
+        delegate :connection, :connection_db_config, to: :index
 
         def with_notifications(action)
           notifier.notify_start(action)
diff --git a/lib/gitlab/database/schema_validation/database.rb b/lib/gitlab/database/schema_validation/database.rb
new file mode 100644
index 00000000000..dfc845f0b44
--- /dev/null
+++ b/lib/gitlab/database/schema_validation/database.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaValidation
+      class Database
+        def initialize(connection)
+          @connection = connection
+        end
+
+        def fetch_index_by_name(index_name)
+          index_map[index_name]
+        end
+
+        def indexes
+          index_map.values
+        end
+
+        private
+
+        def index_map
+          @index_map ||=
+            fetch_indexes.transform_values! do |index_stmt|
+              Index.new(PgQuery.parse(index_stmt).tree.stmts.first.stmt.index_stmt)
+            end
+        end
+
+        attr_reader :connection
+
+        def fetch_indexes
+          sql = <<~SQL
+            SELECT indexname, indexdef
+            FROM pg_indexes
+            WHERE indexname NOT LIKE '%_pkey' AND schemaname IN ('public', 'gitlab_partitions_static');
+          SQL
+
+          @fetch_indexes ||= connection.exec_query(sql).rows.to_h
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/schema_validation/index.rb b/lib/gitlab/database/schema_validation/index.rb
new file mode 100644
index 00000000000..af0d5f31f4e
--- /dev/null
+++ b/lib/gitlab/database/schema_validation/index.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaValidation
+      class Index
+        def initialize(parsed_stmt)
+          @parsed_stmt = parsed_stmt
+        end
+
+        def name
+          parsed_stmt.idxname
+        end
+
+        def statement
+          @statement ||= PgQuery.deparse_stmt(parsed_stmt)
+        end
+
+        private
+
+        attr_reader :parsed_stmt
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/schema_validation/indexes.rb b/lib/gitlab/database/schema_validation/indexes.rb
new file mode 100644
index 00000000000..b7c3705bde9
--- /dev/null
+++ b/lib/gitlab/database/schema_validation/indexes.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaValidation
+      class Indexes
+        def initialize(structure_sql, database)
+          @structure_sql = structure_sql
+          @database = database
+        end
+
+        def missing_indexes
+          structure_sql.indexes.map(&:name) - database.indexes.map(&:name)
+        end
+
+        def extra_indexes
+          database.indexes.map(&:name) - structure_sql.indexes.map(&:name)
+        end
+
+        def wrong_indexes
+          structure_sql.indexes.filter_map do |structure_sql_index|
+            database_index = database.fetch_index_by_name(structure_sql_index.name)
+
+            next if database_index.nil?
+            next if database_index.statement == structure_sql_index.statement
+
+            structure_sql_index.name
+          end
+        end
+
+        private
+
+        attr_reader :structure_sql, :database
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/schema_validation/structure_sql.rb b/lib/gitlab/database/schema_validation/structure_sql.rb
new file mode 100644
index 00000000000..32c69a0e5e7
--- /dev/null
+++ b/lib/gitlab/database/schema_validation/structure_sql.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaValidation
+      class StructureSql
+        def initialize(structure_file_path)
+          @structure_file_path = structure_file_path
+        end
+
+        def indexes
+          @indexes ||= index_statements.map do |index_statement|
+            index_statement.relation.schemaname = "public" if index_statement.relation.schemaname == ''
+
+            Index.new(index_statement)
+          end
+        end
+
+        private
+
+        attr_reader :structure_file_path
+
+        def index_statements
+          parsed_structure_file.tree.stmts.filter_map { |s| s.stmt.index_stmt }
+        end
+
+        def parsed_structure_file
+          PgQuery.parse(File.read(structure_file_path))
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/shared_model.rb b/lib/gitlab/database/shared_model.rb
index 877866b9b23..41c3a27bc5b 100644
--- a/lib/gitlab/database/shared_model.rb
+++ b/lib/gitlab/database/shared_model.rb
@@ -44,6 +44,11 @@ module Gitlab
           end
         end
 
+        # in case the connection has been switched with using_connection
+        def connection_pool
+          connection.pool
+        end
+
         private
 
         def overriding_connection
diff --git a/lib/gitlab/database/tables_locker.rb b/lib/gitlab/database/tables_locker.rb
new file mode 100644
index 00000000000..c417ce716e8
--- /dev/null
+++ b/lib/gitlab/database/tables_locker.rb
@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class TablesLocker
+      GITLAB_SCHEMAS_TO_IGNORE = %i[gitlab_geo].freeze
+
+      def initialize(logger: nil, dry_run: false)
+        @logger = logger
+        @dry_run = dry_run
+      end
+
+      def unlock_writes
+        Gitlab::Database::EachDatabase.each_database_connection do |connection, database_name|
+          tables_to_lock(connection) do |table_name, schema_name|
+            # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
+            next if schema_name.in? GITLAB_SCHEMAS_TO_IGNORE
+
+            lock_writes_manager(table_name, connection, database_name).unlock_writes
+          end
+        end
+      end
+
+      def lock_writes
+        Gitlab::Database::EachDatabase.each_database_connection(include_shared: false) do |connection, database_name|
+          schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
+
+          tables_to_lock(connection) do |table_name, schema_name|
+            # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
+            next if schema_name.in? GITLAB_SCHEMAS_TO_IGNORE
+
+            if schemas_for_connection.include?(schema_name)
+              lock_writes_manager(table_name, connection, database_name).unlock_writes
+            else
+              lock_writes_manager(table_name, connection, database_name).lock_writes
+            end
+          end
+        end
+      end
+
+      private
+
+      def tables_to_lock(connection, &block)
+        Gitlab::Database::GitlabSchema.tables_to_schema.each(&block)
+
+        Gitlab::Database::SharedModel.using_connection(connection) do
+          Postgresql::DetachedPartition.find_each do |detached_partition|
+            yield detached_partition.fully_qualified_table_name, detached_partition.table_schema
+          end
+        end
+      end
+
+      def lock_writes_manager(table_name, connection, database_name)
+        Gitlab::Database::LockWritesManager.new(
+          table_name: table_name,
+          connection: connection,
+          database_name: database_name,
+          with_retries: true,
+          logger: @logger,
+          dry_run: @dry_run
+        )
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/tables_truncate.rb b/lib/gitlab/database/tables_truncate.rb
index daef0402742..a6430d1758b 100644
--- a/lib/gitlab/database/tables_truncate.rb
+++ b/lib/gitlab/database/tables_truncate.rb
@@ -71,17 +71,25 @@ module Gitlab
         @connection ||= Gitlab::Database.database_base_models[database_name].connection
       end
 
+      def remove_schema_name(table_with_schema)
+        ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
+          .extract_schema_qualified_name(table_with_schema)
+          .identifier
+      end
+
+      def disable_locks_on_table(table)
+        sql_statement = "SELECT set_config('lock_writes.#{table}', 'false', false)"
+        logger&.info(sql_statement)
+        connection.execute(sql_statement) unless dry_run
+      end
+
       def truncate_tables_in_batches(tables_sorted)
         truncated_tables = []
 
         tables_sorted.flatten.each do |table|
-          table_name_without_schema = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
-            .extract_schema_qualified_name(table)
-            .identifier
+          table_name_without_schema = remove_schema_name(table)
 
-          sql_statement = "SELECT set_config('lock_writes.#{table_name_without_schema}', 'false', false)"
-          logger&.info(sql_statement)
-          connection.execute(sql_statement) unless dry_run
+          disable_locks_on_table(table_name_without_schema)
 
           # Temporarily unlocking writes on the attached partitions of the table.
           # Because in some cases they might have been locked for writes as well, when they used to be
@@ -89,13 +97,7 @@ module Gitlab
           Gitlab::Database::SharedModel.using_connection(connection) do
             table_partitions = Gitlab::Database::PostgresPartition.for_parent_table(table_name_without_schema)
             table_partitions.each do |table_partition|
-              partition_name_without_schema = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
-                .extract_schema_qualified_name(table_partition.identifier)
-                .identifier
-
-              sql_statement = "SELECT set_config('lock_writes.#{partition_name_without_schema}', 'false', false)"
-              logger&.info(sql_statement)
-              connection.execute(sql_statement) unless dry_run
+              disable_locks_on_table(remove_schema_name(table_partition.identifier))
             end
           end
         end
diff --git a/lib/gitlab/database/transaction_timeout_settings.rb b/lib/gitlab/database/transaction_timeout_settings.rb
new file mode 100644
index 00000000000..9485b8d4cbc
--- /dev/null
+++ b/lib/gitlab/database/transaction_timeout_settings.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class TransactionTimeoutSettings
+      SETTING = 'idle_in_transaction_session_timeout'
+
+      def initialize(connection)
+        @connection = connection
+      end
+
+      def disable_timeouts
+        @connection.execute("SET #{SETTING} = 0")
+      end
+
+      def restore_timeouts
+        @connection.execute("RESET #{SETTING}")
+      end
+    end
+  end
+end