Add latest changes from gitlab-org/gitlab@15-0-stable-eev15.0.0-rc42

24 files changed, 399 insertions, 121 deletions

diff --git a/lib/gitlab/database/background_migration/batch_optimizer.rb b/lib/gitlab/database/background_migration/batch_optimizer.rb
index 58c4a214077..c8fdf8281cd 100644
--- a/lib/gitlab/database/background_migration/batch_optimizer.rb
+++ b/lib/gitlab/database/background_migration/batch_optimizer.rb
@@ -41,7 +41,7 @@ module Gitlab
         end
 
         def optimize!
-          return unless Feature.enabled?(:optimize_batched_migrations, type: :ops, default_enabled: :yaml)
+          return unless Feature.enabled?(:optimize_batched_migrations, type: :ops)
 
           if multiplier = batch_size_multiplier
             max_batch = migration.max_batch_size || MAX_BATCH_SIZE
diff --git a/lib/gitlab/database/background_migration/batched_migration.rb b/lib/gitlab/database/background_migration/batched_migration.rb
index d94bf060d05..a90cae7aea2 100644
--- a/lib/gitlab/database/background_migration/batched_migration.rb
+++ b/lib/gitlab/database/background_migration/batched_migration.rb
@@ -28,6 +28,8 @@ module Gitlab
         # on_hold_until is a temporary runtime status which puts execution "on hold"
         scope :executable, -> { with_status(:active).where('on_hold_until IS NULL OR on_hold_until < NOW()') }
 
+        scope :created_after, ->(time) { where('created_at > ?', time) }
+
         scope :for_configuration, ->(job_class_name, table_name, column_name, job_arguments) do
           where(job_class_name: job_class_name, table_name: table_name, column_name: column_name)
             .where("job_arguments = ?", job_arguments.to_json) # rubocop:disable Rails/WhereEquals
diff --git a/lib/gitlab/database/background_migration/batched_migration_wrapper.rb b/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
index ec68f401ca2..5f4b2be3da8 100644
--- a/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
+++ b/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
@@ -39,7 +39,40 @@ module Gitlab
         end
 
         def execute_batch(tracking_record)
-          job_instance = migration_instance_for(tracking_record.migration_job_class)
+          job_instance = execute_job(tracking_record)
+
+          if job_instance.respond_to?(:batch_metrics)
+            tracking_record.metrics = job_instance.batch_metrics
+          end
+        end
+
+        def execute_job(tracking_record)
+          job_class = tracking_record.migration_job_class
+
+          if job_class < Gitlab::BackgroundMigration::BatchedMigrationJob
+            execute_batched_migration_job(job_class, tracking_record)
+          else
+            execute_legacy_job(job_class, tracking_record)
+          end
+        end
+
+        def execute_batched_migration_job(job_class, tracking_record)
+          job_instance = job_class.new(
+            start_id: tracking_record.min_value,
+            end_id: tracking_record.max_value,
+            batch_table: tracking_record.migration_table_name,
+            batch_column: tracking_record.migration_column_name,
+            sub_batch_size: tracking_record.sub_batch_size,
+            pause_ms: tracking_record.pause_ms,
+            connection: connection)
+
+          job_instance.perform(*tracking_record.migration_job_arguments)
+
+          job_instance
+        end
+
+        def execute_legacy_job(job_class, tracking_record)
+          job_instance = job_class.new
 
           job_instance.perform(
             tracking_record.min_value,
@@ -50,17 +83,7 @@ module Gitlab
             tracking_record.pause_ms,
             *tracking_record.migration_job_arguments)
 
-          if job_instance.respond_to?(:batch_metrics)
-            tracking_record.metrics = job_instance.batch_metrics
-          end
-        end
-
-        def migration_instance_for(job_class)
-          if job_class < Gitlab::BackgroundMigration::BaseJob
-            job_class.new(connection: connection)
-          else
-            job_class.new
-          end
+          job_instance
         end
       end
     end
diff --git a/lib/gitlab/database/gitlab_schemas.yml b/lib/gitlab/database/gitlab_schemas.yml
index ae0ea919b62..036ce7d7631 100644
--- a/lib/gitlab/database/gitlab_schemas.yml
+++ b/lib/gitlab/database/gitlab_schemas.yml
@@ -217,7 +217,6 @@ geo_event_log: :gitlab_main
 geo_events: :gitlab_main
 geo_hashed_storage_attachments_events: :gitlab_main
 geo_hashed_storage_migrated_events: :gitlab_main
-geo_job_artifact_deleted_events: :gitlab_main
 geo_lfs_object_deleted_events: :gitlab_main
 geo_node_namespace_links: :gitlab_main
 geo_nodes: :gitlab_main
@@ -327,6 +326,7 @@ namespace_aggregation_schedules: :gitlab_main
 namespace_limits: :gitlab_main
 namespace_package_settings: :gitlab_main
 namespace_root_storage_statistics: :gitlab_main
+namespace_ci_cd_settings: :gitlab_main
 namespace_settings: :gitlab_main
 namespaces: :gitlab_main
 namespaces_sync_events: :gitlab_main
@@ -348,6 +348,7 @@ operations_strategies: :gitlab_main
 operations_strategies_user_lists: :gitlab_main
 operations_user_lists: :gitlab_main
 packages_build_infos: :gitlab_main
+packages_cleanup_policies: :gitlab_main
 packages_composer_cache_files: :gitlab_main
 packages_composer_metadata: :gitlab_main
 packages_conan_file_metadata: :gitlab_main
@@ -388,6 +389,7 @@ plan_limits: :gitlab_main
 plans: :gitlab_main
 pool_repositories: :gitlab_main
 postgres_async_indexes: :gitlab_shared
+postgres_autovacuum_activity: :gitlab_shared
 postgres_foreign_keys: :gitlab_shared
 postgres_index_bloat_estimates: :gitlab_shared
 postgres_indexes: :gitlab_shared
diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/lib/gitlab/database/load_balancing/configuration.rb
index 3f03d9e2c12..0ddc745ebae 100644
--- a/lib/gitlab/database/load_balancing/configuration.rb
+++ b/lib/gitlab/database/load_balancing/configuration.rb
@@ -90,7 +90,7 @@ module Gitlab
           return false unless ::Gitlab::SafeRequestStore.active?
 
           ::Gitlab::SafeRequestStore.fetch(:force_no_sharing_primary_model) do
-            ::Feature::FlipperFeature.table_exists? && ::Feature.enabled?(:force_no_sharing_primary_model, default_enabled: :yaml)
+            ::Feature::FlipperFeature.table_exists? && ::Feature.enabled?(:force_no_sharing_primary_model)
           end
         end
 
diff --git a/lib/gitlab/database/load_balancing/load_balancer.rb b/lib/gitlab/database/load_balancing/load_balancer.rb
index 1e27bcfc55d..191ebe18b8a 100644
--- a/lib/gitlab/database/load_balancing/load_balancer.rb
+++ b/lib/gitlab/database/load_balancing/load_balancer.rb
@@ -255,6 +255,7 @@ module Gitlab
         # ActiveRecord::ConnectionAdapters::ConnectionHandler handles fetching,
         # and caching for connections pools for each "connection", so we
         # leverage that.
+        # rubocop:disable Database/MultipleDatabases
         def pool
           ActiveRecord::Base.connection_handler.retrieve_connection_pool(
             @configuration.primary_connection_specification_name,
@@ -262,6 +263,7 @@ module Gitlab
             shard: ActiveRecord::Base.default_shard
           ) || raise(::ActiveRecord::ConnectionNotEstablished)
         end
+        # rubocop:enable Database/MultipleDatabases
 
         def wal_diff(location1, location2)
           read_write do |connection|
diff --git a/lib/gitlab/database/migration.rb b/lib/gitlab/database/migration.rb
index dc695a74a4b..038af570dbc 100644
--- a/lib/gitlab/database/migration.rb
+++ b/lib/gitlab/database/migration.rb
@@ -37,18 +37,19 @@ module Gitlab
       class V1_0 < ActiveRecord::Migration[6.1] # rubocop:disable Naming/ClassAndModuleCamelCase
         include LockRetriesConcern
         include Gitlab::Database::MigrationHelpers::V2
-      end
-
-      class V2_0 < V1_0 # rubocop:disable Naming/ClassAndModuleCamelCase
-        include Gitlab::Database::MigrationHelpers::RestrictGitlabSchema
 
         # When running migrations, the `db:migrate` switches connection of
         # ActiveRecord::Base depending where the migration runs.
         # This helper class is provided to avoid confusion using `ActiveRecord::Base`
         class MigrationRecord < ActiveRecord::Base
+          self.abstract_class = true # Prevent STI behavior
         end
       end
 
+      class V2_0 < V1_0 # rubocop:disable Naming/ClassAndModuleCamelCase
+        include Gitlab::Database::MigrationHelpers::RestrictGitlabSchema
+      end
+
       def self.[](version)
         version = version.to_s
         name = "V#{version.tr('.', '_')}"
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index d016dea224b..0453b81d67d 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -3,6 +3,7 @@
 module Gitlab
   module Database
     module MigrationHelpers
+      include Migrations::ReestablishedConnectionStack
       include Migrations::BackgroundMigrationHelpers
       include Migrations::BatchedBackgroundMigrationHelpers
       include DynamicModelHelpers
@@ -943,7 +944,7 @@ module Gitlab
         execute("DELETE FROM batched_background_migrations WHERE #{conditions}")
       end
 
-      def ensure_batched_background_migration_is_finished(job_class_name:, table_name:, column_name:, job_arguments:)
+      def ensure_batched_background_migration_is_finished(job_class_name:, table_name:, column_name:, job_arguments:, finalize: true)
         migration = Gitlab::Database::BackgroundMigration::BatchedMigration
           .for_configuration(job_class_name, table_name, column_name, job_arguments).first
 
@@ -954,14 +955,18 @@ module Gitlab
           job_arguments: job_arguments
         }
 
-        if migration.nil?
-          Gitlab::AppLogger.warn "Could not find batched background migration for the given configuration: #{configuration}"
-        elsif !migration.finished?
+        return Gitlab::AppLogger.warn "Could not find batched background migration for the given configuration: #{configuration}" if migration.nil?
+
+        return if migration.finished?
+
+        finalize_batched_background_migration(job_class_name: job_class_name, table_name: table_name, column_name: column_name, job_arguments: job_arguments) if finalize
+
+        unless migration.reload.finished? # rubocop:disable Cop/ActiveRecordAssociationReload
           raise "Expected batched background migration for the given configuration to be marked as 'finished', " \
             "but it is '#{migration.status_name}':" \
             "\t#{configuration}" \
             "\n\n" \
-            "Finalize it manualy by running" \
+            "Finalize it manually by running" \
             "\n\n" \
             "\tsudo gitlab-rake gitlab:background_migrations:finalize[#{job_class_name},#{table_name},#{column_name},'#{job_arguments.to_json.gsub(',', '\,')}']" \
             "\n\n" \
diff --git a/lib/gitlab/database/migration_helpers/restrict_gitlab_schema.rb b/lib/gitlab/database/migration_helpers/restrict_gitlab_schema.rb
index 5a25128f3a9..d8d07fcaf2d 100644
--- a/lib/gitlab/database/migration_helpers/restrict_gitlab_schema.rb
+++ b/lib/gitlab/database/migration_helpers/restrict_gitlab_schema.rb
@@ -27,7 +27,7 @@ module Gitlab
             return
           end
 
-          Gitlab::Database::QueryAnalyzer.instance.within([validator_class]) do
+          Gitlab::Database::QueryAnalyzer.instance.within([validator_class, connection_validator_class]) do
             validator_class.allowed_gitlab_schemas = self.allowed_gitlab_schemas
 
             super
@@ -45,6 +45,10 @@ module Gitlab
           Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas
         end
 
+        def connection_validator_class
+          Gitlab::Database::QueryAnalyzers::GitlabSchemasValidateConnection
+        end
+
         def unmatched_schemas
           (self.allowed_gitlab_schemas || []) - allowed_schemas_for_connection
         end
diff --git a/lib/gitlab/database/migrations/background_migration_helpers.rb b/lib/gitlab/database/migrations/background_migration_helpers.rb
index 7e5c002d072..9bffed43077 100644
--- a/lib/gitlab/database/migrations/background_migration_helpers.rb
+++ b/lib/gitlab/database/migrations/background_migration_helpers.rb
@@ -41,6 +41,25 @@ module Gitlab
         #       end
         #     end
         def queue_background_migration_jobs_by_range_at_intervals(model_class, job_class_name, delay_interval, batch_size: BATCH_SIZE, other_job_arguments: [], initial_delay: 0, track_jobs: false, primary_column_name: :id)
+          if transaction_open?
+            raise 'The `#queue_background_migration_jobs_by_range_at_intervals` can not be run inside a transaction, ' \
+              'you can disable transactions by calling disable_ddl_transaction! ' \
+              'in the body of your migration class'
+          end
+
+          # Background Migrations do not work well for in cases requiring to update `gitlab_shared`
+          # Once the decomposition is done, enqueued jobs for `gitlab_shared` tables (on CI database)
+          # will not be executed since the queue (which is stored in Redis) is tied to main database, not to schema.
+          # The batched background migrations do not have those limitations since the tracking tables
+          # are properly database-only.
+          if background_migration_restrict_gitlab_migration_schemas&.include?(:gitlab_shared)
+            raise 'The `#queue_background_migration_jobs_by_range_at_intervals` cannot " \
+              "use `restrict_gitlab_migration:` " with `:gitlab_shared`. ' \
+              'Background migrations do encode migration worker which is tied to a given database. ' \
+              'After split this worker will not be properly duplicated into decomposed database. ' \
+              'Use batched background migrations instead that do support well working across all databases.'
+          end
+
           raise "#{model_class} does not have an ID column of #{primary_column_name} to use for batch ranges" unless model_class.column_names.include?(primary_column_name.to_s)
           raise "#{primary_column_name} is not an integer or string column" unless [:integer, :string].include?(model_class.columns_hash[primary_column_name.to_s].type)
 
@@ -90,6 +109,18 @@ module Gitlab
         # delay_interval - The duration between each job's scheduled time
         # batch_size - The maximum number of jobs to fetch to memory from the database.
         def requeue_background_migration_jobs_by_range_at_intervals(job_class_name, delay_interval, batch_size: BATCH_SIZE, initial_delay: 0)
+          if transaction_open?
+            raise 'The `#requeue_background_migration_jobs_by_range_at_intervals` can not be run inside a transaction, ' \
+              'you can disable transactions by calling disable_ddl_transaction! ' \
+              'in the body of your migration class'
+          end
+
+          if background_migration_restrict_gitlab_migration_schemas&.any?
+            raise 'The `#requeue_background_migration_jobs_by_range_at_intervals` cannot use `restrict_gitlab_migration:`. ' \
+              'The `#requeue_background_migration_jobs_by_range_at_intervals` needs to be executed on all databases since ' \
+              'each database has its own queue of background migrations.'
+          end
+
           job_coordinator = coordinator_for_tracking_database
 
           # To not overload the worker too much we enforce a minimum interval both
@@ -133,23 +164,40 @@ module Gitlab
         # This method does not garauntee that all jobs completed successfully.
         # It can only be used if the previous background migration used the queue_background_migration_jobs_by_range_at_intervals helper.
         def finalize_background_migration(class_name, delete_tracking_jobs: ['succeeded'])
+          if transaction_open?
+            raise 'The `#finalize_background_migration` can not be run inside a transaction, ' \
+              'you can disable transactions by calling disable_ddl_transaction! ' \
+              'in the body of your migration class'
+          end
+
+          if background_migration_restrict_gitlab_migration_schemas&.any?
+            raise 'The `#finalize_background_migration` cannot use `restrict_gitlab_migration:`. ' \
+              'The `#finalize_background_migration` needs to be executed on all databases since ' \
+              'each database has its own queue of background migrations.'
+          end
+
           job_coordinator = coordinator_for_tracking_database
 
-          # Empty the sidekiq queue.
-          job_coordinator.steal(class_name)
+          with_restored_connection_stack do
+            # Since we are running trusted code (background migration class) allow to execute any type of finalize
+            Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.with_suppressed do
+              # Empty the sidekiq queue.
+              job_coordinator.steal(class_name)
 
-          # Process pending tracked jobs.
-          jobs = Gitlab::Database::BackgroundMigrationJob.pending.for_migration_class(class_name)
+              # Process pending tracked jobs.
+              jobs = Gitlab::Database::BackgroundMigrationJob.pending.for_migration_class(class_name)
 
-          jobs.find_each do |job|
-            job_coordinator.perform(job.class_name, job.arguments)
-          end
+              jobs.find_each do |job|
+                job_coordinator.perform(job.class_name, job.arguments)
+              end
 
-          # Empty the sidekiq queue.
-          job_coordinator.steal(class_name)
+              # Empty the sidekiq queue.
+              job_coordinator.steal(class_name)
 
-          # Delete job tracking rows.
-          delete_job_tracking(class_name, status: delete_tracking_jobs) if delete_tracking_jobs
+              # Delete job tracking rows.
+              delete_job_tracking(class_name, status: delete_tracking_jobs) if delete_tracking_jobs
+            end
+          end
         end
 
         def migrate_in(*args, coordinator: coordinator_for_tracking_database)
@@ -174,6 +222,10 @@ module Gitlab
 
         private
 
+        def background_migration_restrict_gitlab_migration_schemas
+          self.allowed_gitlab_schemas if self.respond_to?(:allowed_gitlab_schemas)
+        end
+
         def with_migration_context(&block)
           Gitlab::ApplicationContext.with_context(caller_id: self.class.to_s, &block)
         end
@@ -183,11 +235,9 @@ module Gitlab
         end
 
         def coordinator_for_tracking_database
-          Gitlab::BackgroundMigration.coordinator_for_database(tracking_database)
-        end
+          tracking_database = Gitlab::Database.db_config_name(connection)
 
-        def tracking_database
-          Gitlab::BackgroundMigration::DEFAULT_TRACKING_DATABASE
+          Gitlab::BackgroundMigration.coordinator_for_database(tracking_database)
         end
       end
     end
diff --git a/lib/gitlab/database/migrations/base_background_runner.rb b/lib/gitlab/database/migrations/base_background_runner.rb
new file mode 100644
index 00000000000..2772502140e
--- /dev/null
+++ b/lib/gitlab/database/migrations/base_background_runner.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      class BaseBackgroundRunner
+        attr_reader :result_dir
+
+        def initialize(result_dir:)
+          @result_dir = result_dir
+        end
+
+        def jobs_by_migration_name
+          raise NotImplementedError, 'subclass must implement'
+        end
+
+        def run_job(job)
+          raise NotImplementedError, 'subclass must implement'
+        end
+
+        def run_jobs(for_duration:)
+          jobs_to_run = jobs_by_migration_name
+          return if jobs_to_run.empty?
+
+          # without .to_f, we do integer division
+          # For example, 3.minutes / 2 == 1.minute whereas 3.minutes / 2.to_f == (1.minute + 30.seconds)
+          duration_per_migration_type = for_duration / jobs_to_run.count.to_f
+          jobs_to_run.each do |migration_name, jobs|
+            run_until = duration_per_migration_type.from_now
+
+            run_jobs_for_migration(migration_name: migration_name, jobs: jobs, run_until: run_until)
+          end
+        end
+
+        private
+
+        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)
+          batch_names = (1..).each.lazy.map { |i| "batch_#{i}"}
+
+          jobs.shuffle.each do |j|
+            break if run_until <= Time.current
+
+            instrumentation.observe(version: nil,
+                                    name: batch_names.next,
+                                    connection: ActiveRecord::Migration.connection) do
+              run_job(j)
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
index 0261ade0fe7..7113c3686f1 100644
--- a/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
+++ b/lib/gitlab/database/migrations/batched_background_migration_helpers.rb
@@ -122,6 +122,22 @@ module Gitlab
           migration.save!
           migration
         end
+
+        def finalize_batched_background_migration(job_class_name:, table_name:, column_name:, job_arguments:)
+          database_name = Gitlab::Database.db_config_name(connection)
+
+          unless ActiveRecord::Base.configurations.primary?(database_name)
+            raise 'The `#finalize_background_migration` is currently not supported when running in decomposed database, ' \
+              'and this database is not `main:`. For more information visit: ' \
+              'https://docs.gitlab.com/ee/development/database/migrations_for_multiple_databases.html'
+          end
+
+          migration = Gitlab::Database::BackgroundMigration::BatchedMigration.find_for_configuration(job_class_name, table_name, column_name, job_arguments)
+
+          raise 'Could not find batched background migration' if migration.nil?
+
+          Gitlab::Database::BackgroundMigration::BatchedMigrationRunner.finalize(job_class_name, table_name, column_name, job_arguments, connection: connection)
+        end
       end
     end
   end
diff --git a/lib/gitlab/database/migrations/observers/query_log.rb b/lib/gitlab/database/migrations/observers/query_log.rb
index 8ca57bb7df9..543e6b8e302 100644
--- a/lib/gitlab/database/migrations/observers/query_log.rb
+++ b/lib/gitlab/database/migrations/observers/query_log.rb
@@ -6,7 +6,7 @@ module Gitlab
       module Observers
         class QueryLog < MigrationObserver
           def before
-            @logger_was = ActiveRecord::Base.logger
+            @logger_was = ActiveRecord::Base.logger # rubocop:disable Database/MultipleDatabases
             file_path = File.join(output_dir, "migration.log")
             @logger = Logger.new(file_path)
             ActiveRecord::Base.logger = @logger
diff --git a/lib/gitlab/database/migrations/observers/query_statistics.rb b/lib/gitlab/database/migrations/observers/query_statistics.rb
index 54504646a79..2d026f0c8d2 100644
--- a/lib/gitlab/database/migrations/observers/query_statistics.rb
+++ b/lib/gitlab/database/migrations/observers/query_statistics.rb
@@ -22,6 +22,7 @@ module Gitlab
             observation.query_statistics = connection.execute(<<~SQL)
               SELECT query, calls, total_time, max_time, mean_time, rows
               FROM pg_stat_statements
+              WHERE pg_get_userbyid(userid) = current_user
               ORDER BY total_time DESC
             SQL
           end
diff --git a/lib/gitlab/database/migrations/reestablished_connection_stack.rb b/lib/gitlab/database/migrations/reestablished_connection_stack.rb
new file mode 100644
index 00000000000..d7cf482c32a
--- /dev/null
+++ b/lib/gitlab/database/migrations/reestablished_connection_stack.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module ReestablishedConnectionStack
+        # This is workaround for `db:migrate` that switches `ActiveRecord::Base.connection`
+        # depending on execution. This is subject to be removed once proper fix is implemented:
+        # https://gitlab.com/gitlab-org/gitlab/-/issues/362341
+        #
+        # In some cases when we run application code we need to restore application connection stack:
+        # - ApplicationRecord (in fact ActiveRecord::Base): points to main
+        # - Ci::ApplicationRecord: points to ci
+        #
+        # rubocop:disable Database/MultipleDatabases
+        def with_restored_connection_stack(&block)
+          original_handler = ActiveRecord::Base.connection_handler
+
+          original_db_config = ActiveRecord::Base.connection_db_config
+          return yield if ActiveRecord::Base.configurations.primary?(original_db_config.name)
+
+          # If the `ActiveRecord::Base` connection is different than `:main`
+          # re-establish and configure `SharedModel` context accordingly
+          # to previously established `ActiveRecord::Base` to allow the application
+          # code to use `ApplicationRecord` and `Ci::ApplicationRecord` usual way.
+          # We swap a connection handler as migration context does hold an actual
+          # connection which we cannot close.
+          base_model = Gitlab::Database.database_base_models.fetch(original_db_config.name.to_sym)
+
+          # copy connections over to new connection handler
+          db_configs = original_handler.connection_pool_names.map do |connection_pool_name|
+            [connection_pool_name.constantize, connection_pool_name.constantize.connection_db_config]
+          end
+
+          new_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+          ActiveRecord::Base.connection_handler = new_handler
+
+          db_configs.each do |klass, db_config|
+            new_handler.establish_connection(db_config, owner_name: klass)
+          end
+
+          # re-establish ActiveRecord::Base to main
+          ActiveRecord::Base.establish_connection :main # rubocop:disable Database/EstablishConnection
+
+          Gitlab::Database::SharedModel.using_connection(base_model.connection) do
+            yield
+          end
+        ensure
+          ActiveRecord::Base.connection_handler = original_handler
+          new_handler&.clear_all_connections!
+        end
+        # rubocop:enable Database/MultipleDatabases
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/migrations/runner.rb b/lib/gitlab/database/migrations/runner.rb
index 3b6f52b43a8..4404b5bf961 100644
--- a/lib/gitlab/database/migrations/runner.rb
+++ b/lib/gitlab/database/migrations/runner.rb
@@ -21,6 +21,18 @@ module Gitlab
             TestBackgroundRunner.new(result_dir: BASE_RESULT_DIR.join('background_migrations'))
           end
 
+          def batched_background_migrations(for_database:)
+            runner = nil
+
+            # Only one loop iteration since we pass `only:` here
+            Gitlab::Database::EachDatabase.each_database_connection(only: for_database) do |connection|
+              runner = Gitlab::Database::Migrations::TestBatchedBackgroundRunner
+                         .new(result_dir: BASE_RESULT_DIR.join('background_migrations'), connection: connection)
+            end
+
+            runner
+          end
+
           def migration_context
             @migration_context ||= ApplicationRecord.connection.migration_context
           end
diff --git a/lib/gitlab/database/migrations/test_background_runner.rb b/lib/gitlab/database/migrations/test_background_runner.rb
index 74e54d62e05..f7713237b38 100644
--- a/lib/gitlab/database/migrations/test_background_runner.rb
+++ b/lib/gitlab/database/migrations/test_background_runner.rb
@@ -3,11 +3,9 @@
 module Gitlab
   module Database
     module Migrations
-      class TestBackgroundRunner
-        attr_reader :result_dir
-
+      class TestBackgroundRunner < BaseBackgroundRunner
         def initialize(result_dir:)
-          @result_dir = result_dir
+          super(result_dir: result_dir)
           @job_coordinator = Gitlab::BackgroundMigration.coordinator_for_database(Gitlab::Database::MAIN_DATABASE_NAME)
         end
 
@@ -15,37 +13,12 @@ module Gitlab
           @job_coordinator.pending_jobs
         end
 
-        def run_jobs(for_duration:)
-          jobs_to_run = traditional_background_migrations.group_by { |j| class_name_for_job(j) }
-          return if jobs_to_run.empty?
-
-          # without .to_f, we do integer division
-          # For example, 3.minutes / 2 == 1.minute whereas 3.minutes / 2.to_f == (1.minute + 30.seconds)
-          duration_per_migration_type = for_duration / jobs_to_run.count.to_f
-          jobs_to_run.each do |migration_name, jobs|
-            run_until = duration_per_migration_type.from_now
-
-            run_jobs_for_migration(migration_name: migration_name, jobs: jobs, run_until: run_until)
-          end
+        def jobs_by_migration_name
+          traditional_background_migrations.group_by { |j| class_name_for_job(j) }
         end
 
         private
 
-        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)
-          batch_names = (1..).each.lazy.map { |i| "batch_#{i}"}
-
-          jobs.shuffle.each do |j|
-            break if run_until <= Time.current
-
-            instrumentation.observe(version: nil, name: batch_names.next, connection: ActiveRecord::Migration.connection) do
-              run_job(j)
-            end
-          end
-        end
-
         def run_job(job)
           Gitlab::BackgroundMigration.perform(job.args[0], job.args[1])
         end
diff --git a/lib/gitlab/database/migrations/test_batched_background_runner.rb b/lib/gitlab/database/migrations/test_batched_background_runner.rb
new file mode 100644
index 00000000000..0c6a8d3d856
--- /dev/null
+++ b/lib/gitlab/database/migrations/test_batched_background_runner.rb
@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      class TestBatchedBackgroundRunner < BaseBackgroundRunner
+        attr_reader :connection
+
+        def initialize(result_dir:, connection:)
+          super(result_dir: result_dir)
+          @connection = connection
+        end
+
+        def jobs_by_migration_name
+          Gitlab::Database::BackgroundMigration::BatchedMigration
+            .executable
+            .created_after(2.days.ago) # Simple way to exclude migrations already running before migration testing
+            .to_h do |migration|
+            batching_strategy = migration.batch_class.new(connection: connection)
+
+            all_migration_jobs = []
+
+            min_value = migration.next_min_value
+
+            while (next_bounds = batching_strategy.next_batch(
+              migration.table_name,
+              migration.column_name,
+              batch_min_value: min_value,
+              batch_size: migration.batch_size,
+              job_arguments: migration.job_arguments
+            ))
+
+              batch_min, batch_max = next_bounds
+
+              all_migration_jobs << migration.create_batched_job!(batch_min, batch_max)
+              min_value = batch_max + 1
+            end
+
+            [migration.job_class_name, all_migration_jobs]
+          end
+        end
+
+        def run_job(job)
+          Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new(connection: connection).perform(job)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
index 034e18ec9f4..a541ecf5316 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
@@ -6,7 +6,6 @@ module Gitlab
       module TableManagementHelpers
         include ::Gitlab::Database::SchemaHelpers
         include ::Gitlab::Database::MigrationHelpers
-        include ::Gitlab::Database::Migrations::BackgroundMigrationHelpers
 
         ALLOWED_TABLES = %w[audit_events web_hook_logs].freeze
         ERROR_SCOPE = 'table partitioning'
diff --git a/lib/gitlab/database/query_analyzer.rb b/lib/gitlab/database/query_analyzer.rb
index 0c78dda734c..6f64d04270f 100644
--- a/lib/gitlab/database/query_analyzer.rb
+++ b/lib/gitlab/database/query_analyzer.rb
@@ -30,52 +30,25 @@ module Gitlab
         end
       end
 
-      def within(user_analyzers = nil)
-        # Due to singleton nature of analyzers
-        # only an outer invocation of the `.within`
-        # is allowed to initialize them
-        if already_within?
-          raise 'Query analyzers are already defined, cannot re-define them.' if user_analyzers
-
-          return yield
-        end
-
-        begin!(user_analyzers || all_analyzers)
+      def within(analyzers = all_analyzers)
+        newly_enabled_analyzers = begin!(analyzers)
 
         begin
           yield
         ensure
-          end!
+          end!(newly_enabled_analyzers)
         end
       end
 
-      def already_within?
-        # If analyzers are set they are already configured
-        !enabled_analyzers.nil?
-      end
+      # Enable query analyzers (only the ones that were not yet enabled)
+      # Returns a list of newly enabled analyzers
+      def begin!(analyzers)
+        analyzers.select do |analyzer|
+          next if enabled_analyzers.include?(analyzer)
 
-      def process_sql(sql, connection)
-        analyzers = enabled_analyzers
-        return unless analyzers&.any?
-
-        parsed = parse(sql, connection)
-        return unless parsed
-
-        analyzers.each do |analyzer|
-          next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
-
-          analyzer.analyze(parsed)
-        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
-          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
-          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
-        end
-      end
-
-      # Enable query analyzers
-      def begin!(analyzers = all_analyzers)
-        analyzers = analyzers.select do |analyzer|
           if analyzer.enabled?
             analyzer.begin!
+            enabled_analyzers.append(analyzer)
 
             true
           end
@@ -84,25 +57,40 @@ module Gitlab
 
           false
         end
-
-        Thread.current[:query_analyzer_enabled_analyzers] = analyzers
       end
 
-      # Disable enabled query analyzers
-      def end!
-        enabled_analyzers.select do |analyzer|
+      # Disable enabled query analyzers (only the ones that were enabled previously)
+      def end!(analyzers)
+        analyzers.each do |analyzer|
+          next unless enabled_analyzers.delete(analyzer)
+
           analyzer.end!
         rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
-
-        Thread.current[:query_analyzer_enabled_analyzers] = nil
       end
 
       private
 
       def enabled_analyzers
-        Thread.current[:query_analyzer_enabled_analyzers]
+        Thread.current[:query_analyzer_enabled_analyzers] ||= []
+      end
+
+      def process_sql(sql, connection)
+        analyzers = enabled_analyzers
+        return unless analyzers&.any?
+
+        parsed = parse(sql, connection)
+        return unless parsed
+
+        analyzers.each do |analyzer|
+          next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
+
+          analyzer.analyze(parsed)
+        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
+          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        end
       end
 
       def parse(sql, connection)
diff --git a/lib/gitlab/database/query_analyzers/gitlab_schemas_validate_connection.rb b/lib/gitlab/database/query_analyzers/gitlab_schemas_validate_connection.rb
new file mode 100644
index 00000000000..3de9e8011fb
--- /dev/null
+++ b/lib/gitlab/database/query_analyzers/gitlab_schemas_validate_connection.rb
@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module QueryAnalyzers
+      # The purpose of this analyzer is to validate if tables observed
+      # are properly used according to schema used by current connection
+      class GitlabSchemasValidateConnection < Base
+        CrossSchemaAccessError = Class.new(QueryAnalyzerError)
+
+        class << self
+          def enabled?
+            true
+          end
+
+          def analyze(parsed)
+            tables = parsed.pg.select_tables + parsed.pg.dml_tables
+            table_schemas = ::Gitlab::Database::GitlabSchema.table_schemas(tables)
+            return if table_schemas.empty?
+
+            allowed_schemas = ::Gitlab::Database.gitlab_schemas_for_connection(parsed.connection)
+            return unless allowed_schemas
+
+            invalid_schemas = table_schemas - allowed_schemas
+            if invalid_schemas.any?
+              message = "The query tried to access #{tables} (of #{table_schemas.to_a}) "
+              message += "which is outside of allowed schemas (#{allowed_schemas}) "
+              message += "for the current connection '#{Gitlab::Database.db_config_name(parsed.connection)}'"
+
+              raise CrossSchemaAccessError, message
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb b/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
index a53da514df2..e0cb803b872 100644
--- a/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
+++ b/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
@@ -33,7 +33,7 @@ module Gitlab
 
         def self.enabled?
           ::Feature::FlipperFeature.table_exists? &&
-            Feature.enabled?(:detect_cross_database_modification, default_enabled: :yaml)
+            Feature.enabled?(:detect_cross_database_modification)
         end
 
         def self.requires_tracking?(parsed)
diff --git a/lib/gitlab/database/reindexing.rb b/lib/gitlab/database/reindexing.rb
index 91c3fcc7d72..e13dd3b2058 100644
--- a/lib/gitlab/database/reindexing.rb
+++ b/lib/gitlab/database/reindexing.rb
@@ -16,7 +16,7 @@ module Gitlab
       REMOVE_INDEX_RETRY_CONFIG = [[1.minute, 9.minutes]] * 30
 
       def self.enabled?
-        Feature.enabled?(:database_reindexing, type: :ops, default_enabled: :yaml)
+        Feature.enabled?(:database_reindexing, type: :ops)
       end
 
       def self.invoke(database = nil)
diff --git a/lib/gitlab/database/shared_model.rb b/lib/gitlab/database/shared_model.rb
index 563fab692ef..f4c8fca8fa2 100644
--- a/lib/gitlab/database/shared_model.rb
+++ b/lib/gitlab/database/shared_model.rb
@@ -15,14 +15,16 @@ module Gitlab
           previous_connection = self.overriding_connection
 
           unless previous_connection.nil? || previous_connection.equal?(connection)
-            raise 'cannot nest connection overrides for shared models with different connections'
+            raise "Cannot change connection for Gitlab::Database::SharedModel "\
+              "from '#{Gitlab::Database.db_config_name(previous_connection)}' "\
+              "to '#{Gitlab::Database.db_config_name(connection)}'"
           end
 
           self.overriding_connection = connection
 
           yield
         ensure
-          self.overriding_connection = nil unless previous_connection.equal?(self.overriding_connection)
+          self.overriding_connection = previous_connection
         end
 
         def connection